From 4d817c48a8e207827046cbc7f4713d0721a705d9 Mon Sep 17 00:00:00 2001
From: Nacho Rivera <nachor1992@gmail.com>
Date: Mon, 21 Aug 2023 13:23:43 +0200
Subject: [PATCH] fix(get_checks_from_input_arn): fix function and add tests
 (#2749)

---
 prowler/providers/aws/aws_provider.py |  2 +
 tests/lib/check/check_test.py         | 79 ++++++++++++++++++++++++---
 2 files changed, 73 insertions(+), 8 deletions(-)

diff --git a/prowler/providers/aws/aws_provider.py b/prowler/providers/aws/aws_provider.py
index 9bb5168e..5fcc87c2 100644
--- a/prowler/providers/aws/aws_provider.py
+++ b/prowler/providers/aws/aws_provider.py
@@ -246,6 +246,8 @@ def get_checks_from_input_arn(audit_resources: list, provider: str) -> set:
             if any(sub_service in check for sub_service in sub_service_list):
                 if not (sub_service == "policy" and "password_policy" in check):
                     checks_from_arn.add(check)
+            else:
+                checks_from_arn.add(check)
 
     # Return final checks list
     return sorted(checks_from_arn)
diff --git a/tests/lib/check/check_test.py b/tests/lib/check/check_test.py
index 11ba4f40..9321d92c 100644
--- a/tests/lib/check/check_test.py
+++ b/tests/lib/check/check_test.py
@@ -137,6 +137,40 @@ def mock_recover_checks_from_aws_provider_lambda_service(*_):
     ]
 
 
+def mock_recover_checks_from_aws_provider_iam_service(*_):
+    return [
+        (
+            "iam_customer_attached_policy_no_administrative_privileges",
+            "/root_dir/fake_path/iam/iam_customer_attached_policy_no_administrative_privileges",
+        ),
+        (
+            "iam_check_saml_providers_sts",
+            "/root_dir/fake_path/iam/iam_check_saml_providers_sts",
+        ),
+        (
+            "iam_password_policy_minimum_length_14",
+            "/root_dir/fake_path/iam/iam_password_policy_minimum_length_14",
+        ),
+    ]
+
+
+def mock_recover_checks_from_aws_provider_s3_service(*_):
+    return [
+        (
+            "s3_account_level_public_access_blocks",
+            "/root_dir/fake_path/s3/s3_account_level_public_access_blocks",
+        ),
+        (
+            "s3_bucket_acl_prohibited",
+            "/root_dir/fake_path/s3/s3_bucket_acl_prohibited",
+        ),
+        (
+            "s3_bucket_policy_public_write_access",
+            "/root_dir/fake_path/s3/s3_bucket_policy_public_write_access",
+        ),
+    ]
+
+
 class Test_Check:
     def set_mocked_audit_info(self):
         audit_info = AWS_Audit_Info(
@@ -384,7 +418,7 @@ class Test_Check:
         "prowler.lib.check.check.recover_checks_from_provider",
         new=mock_recover_checks_from_aws_provider_lambda_service,
     )
-    def test_get_checks_from_input_arn(self):
+    def test_get_checks_from_input_arn_lambda(self):
         audit_resources = ["arn:aws:lambda:us-east-1:123456789:function:test-lambda"]
         provider = "aws"
         expected_checks = [
@@ -395,16 +429,45 @@ class Test_Check:
         recovered_checks = get_checks_from_input_arn(audit_resources, provider)
         assert recovered_checks == expected_checks
 
+    @patch(
+        "prowler.lib.check.check.recover_checks_from_provider",
+        new=mock_recover_checks_from_aws_provider_iam_service,
+    )
+    def test_get_checks_from_input_arn_iam(self):
+        audit_resources = [f"arn:aws:iam::{AWS_ACCOUNT_NUMBER}:user/user-name"]
+        provider = "aws"
+        expected_checks = [
+            "iam_check_saml_providers_sts",
+            "iam_customer_attached_policy_no_administrative_privileges",
+            "iam_password_policy_minimum_length_14",
+        ]
+        recovered_checks = get_checks_from_input_arn(audit_resources, provider)
+        assert recovered_checks == expected_checks
+
+    @patch(
+        "prowler.lib.check.check.recover_checks_from_provider",
+        new=mock_recover_checks_from_aws_provider_s3_service,
+    )
+    def test_get_checks_from_input_arn_s3(self):
+        audit_resources = ["arn:aws:s3:::bucket-name"]
+        provider = "aws"
+        expected_checks = [
+            "s3_account_level_public_access_blocks",
+            "s3_bucket_acl_prohibited",
+            "s3_bucket_policy_public_write_access",
+        ]
+        recovered_checks = get_checks_from_input_arn(audit_resources, provider)
+        assert recovered_checks == expected_checks
+
     def test_get_regions_from_audit_resources(self):
         audit_resources = [
-            "arn:aws:lambda:us-east-1:123456789:function:test-lambda",
-            "arn:aws:iam::106908755756:policy/test",
-            "arn:aws:ec2:eu-west-1:106908755756:security-group/sg-test",
-        ]
-        expected_regions = [
-            "us-east-1",
-            "eu-west-1",
+            f"arn:aws:lambda:us-east-1:{AWS_ACCOUNT_NUMBER}:function:test-lambda",
+            f"arn:aws:iam::{AWS_ACCOUNT_NUMBER}:policy/test",
+            f"arn:aws:ec2:eu-west-1:{AWS_ACCOUNT_NUMBER}:security-group/sg-test",
+            "arn:aws:s3:::bucket-name",
+            "arn:aws:apigateway:us-east-2::/restapis/api-id/stages/stage-name",
         ]
+        expected_regions = ["us-east-1", "eu-west-1", "us-east-2"]
         recovered_regions = get_regions_from_audit_resources(audit_resources)
         assert recovered_regions == expected_regions