fix(aws): lead the partition bootstrap regions with the configured region (#12764)

Co-authored-by: pedrooot <pedromarting3@gmail.com>

docs/user-guide/providers/aws/regions-and-partitions.mdx

@@ -21,10 +21,28 @@ When scanning the China (`aws-cn`), European Sovereign Cloud (`aws-eusc`) or Gov
 
 - Specify the regions to audit within that partition using the `-f/--region` flag.
 
+- Declare the partition with the `PROWLER_AWS_PARTITION` environment variable, set to `aws`, `aws-cn`, `aws-eusc` or `aws-us-gov`.
+
 <Note>
 Refer to: https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html#configuring-credentials for more information about the AWS credential configuration.
 
 </Note>
+### Declaring the Partition
+
+`PROWLER_AWS_PARTITION` tells Prowler which partition the scan runs against, without relying on a region being configured:
+
+```bash
+export PROWLER_AWS_PARTITION="aws-us-gov"
+```
+
+It matters most where nothing else says. Resolving an identity means calling STS before anything is known about the credentials, and with no region configured Prowler would otherwise start from the commercial endpoints. Declaring the partition makes that first call go to the right place, which is the difference between a scan that starts and one that fails on an endpoint the credentials cannot use.
+
+A region configured for the session still wins when it belongs to the declared partition, so a deployment in `us-gov-west-1` is not sent to `us-gov-east-1`. A region belonging to a different partition is ignored, since a partition that has been declared explicitly is the more deliberate statement of the two.
+
+<Note>
+Set it wherever the scan runs. For deployments that scan from containers, that means the environment of the containers doing the scanning, not only the one accepting the request.
+</Note>
+
 ### Scanning Specific Regions
 
 To scan a particular AWS region with Prowler, use:

prowler/changelog.d/aws-partition-bootstrap-region-honours-configured-region.fixed.md

@@ -0,0 +1 @@
+Bootstrap STS calls now use the session region when `PROWLER_AWS_PARTITION` is set and the region belongs to that partition, instead of always going to the partition's global STS region, which a deployment reached only through its own region's VPC endpoints cannot route to

prowler/providers/aws/aws_provider.py

@@ -576,8 +576,15 @@ class AwsProvider(Provider):
     ) -> str:
         excluded_regions = set(excluded_regions or ())
         session_region = session.region_name
+        env_partition_regions = get_env_partition_regions(session_region)
         if session_region and session_region not in excluded_regions:
-            return session_region
+            if not env_partition_regions or session_region in env_partition_regions:
+                return session_region
+        if env_partition_regions:
+            for region in env_partition_regions:
+                if region not in excluded_regions:
+                    return region
+            return env_partition_regions[0]
 
         for region in AwsProvider.get_bootstrap_region_candidates(session_region):
             if region not in excluded_regions:
@@ -673,7 +680,7 @@ class AwsProvider(Provider):
                 session = Session(**session_arguments)
                 session._session.set_default_client_config(session_config)
                 sts_region = (
-                    get_env_partition_bootstrap_region()
+                    get_env_partition_bootstrap_region(session.region_name)
                     or session.region_name
                     or AWS_STS_GLOBAL_ENDPOINT_REGION
                 )
@@ -1420,12 +1427,6 @@ class AwsProvider(Provider):
             Connection(is_connected=True, Error=None))
         """
         try:
-            if aws_region is None:
-                aws_region = (
-                    get_env_partition_bootstrap_region()
-                    or AWS_STS_GLOBAL_ENDPOINT_REGION
-                )
-
             session = AwsProvider.setup_session(
                 mfa=mfa_enabled,
                 profile=profile,
@@ -1434,6 +1435,12 @@ class AwsProvider(Provider):
                 aws_session_token=aws_session_token,
             )
 
+            if aws_region is None:
+                aws_region = (
+                    get_env_partition_bootstrap_region(session.region_name)
+                    or AWS_STS_GLOBAL_ENDPOINT_REGION
+                )
+
             if role_arn:
                 session_duration = validate_session_duration(session_duration)
                 role_session_name = validate_role_session_name(role_session_name)
@@ -1759,11 +1766,18 @@ def get_botocore_partition_regions() -> dict:
     return partition_regions
 
 
-def get_env_partition_regions() -> Optional[list]:
+def get_env_partition_regions(
+    session_region: Optional[str] = None,
+) -> Optional[list]:
     """
     Get the bootstrap region candidates for the partition set in the
     PROWLER_AWS_PARTITION environment variable.
 
+    Args:
+        session_region (Optional[str]): The region of the AWS session. It leads
+            the candidates when it belongs to the partition and is ignored
+            otherwise.
+
     Returns:
         Optional[list]: The regions of the configured partition, preferred
             bootstrap region first, or None when the environment variable is
@@ -1782,14 +1796,25 @@ def get_env_partition_regions() -> Optional[list]:
         raise AWSInvalidPartitionError(
             message=f"Invalid partition: {raw_partition} set in PROWLER_AWS_PARTITION. Valid partitions: {', '.join(sorted(partition_regions))}"
         )
+
+    # A deployment reached only through its own region's endpoints has no route
+    # to the partition's global STS region, so the session region goes first
+    if session_region in regions:
+        regions = [session_region] + [r for r in regions if r != session_region]
     return regions
 
 
-def get_env_partition_bootstrap_region() -> Optional[str]:
+def get_env_partition_bootstrap_region(
+    session_region: Optional[str] = None,
+) -> Optional[str]:
     """
     Get the STS bootstrap region for the partition set in the
     PROWLER_AWS_PARTITION environment variable.
 
+    Args:
+        session_region (Optional[str]): The region of the AWS session, preferred
+            when it belongs to the partition.
+
     Returns:
         Optional[str]: The preferred bootstrap region of the configured
             partition, or None when the environment variable is not set.
@@ -1797,7 +1822,7 @@ def get_env_partition_bootstrap_region() -> Optional[str]:
     Raises:
         AWSInvalidPartitionError: If the value is not a partition known to botocore.
     """
-    regions = get_env_partition_regions()
+    regions = get_env_partition_regions(session_region)
     return regions[0] if regions else None
 
 
@@ -1833,7 +1858,7 @@ def get_aws_region_for_sts(
             if region not in excluded_regions:
                 return region
 
-    env_partition_regions = get_env_partition_regions()
+    env_partition_regions = get_env_partition_regions(session_region)
     if env_partition_regions:
         # The configured partition constrains the whole fallback chain: prefer
         # a non-excluded region, but never leave the partition

tests/providers/aws/aws_provider_test.py

@@ -16,7 +16,12 @@ from moto import mock_aws
 from pytest import raises
 from tzlocal import get_localzone
 
-from prowler.providers.aws.aws_provider import AwsProvider, get_aws_region_for_sts
+from prowler.providers.aws.aws_provider import (
+    AwsProvider,
+    get_aws_region_for_sts,
+    get_env_partition_bootstrap_region,
+    get_env_partition_regions,
+)
 from prowler.providers.aws.config import (
     AWS_STS_GLOBAL_ENDPOINT_REGION,
     BOTO3_USER_AGENT_EXTRA,
@@ -56,6 +61,7 @@ from tests.providers.aws.utils import (
     AWS_REGION_EU_WEST_1,
     AWS_REGION_EUSC_DE_EAST_1,
     AWS_REGION_GOV_CLOUD_US_EAST_1,
+    AWS_REGION_GOV_CLOUD_US_WEST_1,
     AWS_REGION_ISO_GLOBAL,
     AWS_REGION_US_EAST_1,
     AWS_REGION_US_EAST_2,
@@ -2447,6 +2453,245 @@ aws:
                 == AWS_REGION_GOV_CLOUD_US_EAST_1
             )
 
+    def test_get_env_partition_regions_leads_with_session_region(self):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            regions = get_env_partition_regions(AWS_REGION_GOV_CLOUD_US_WEST_1)
+
+            assert regions[0] == AWS_REGION_GOV_CLOUD_US_WEST_1
+            assert set(regions) == set(get_env_partition_regions())
+
+    def test_get_env_partition_regions_ignores_session_region_outside_partition(
+        self,
+    ):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            regions = get_env_partition_regions(AWS_REGION_EU_WEST_1)
+
+            assert regions[0] == AWS_REGION_GOV_CLOUD_US_EAST_1
+            assert AWS_REGION_EU_WEST_1 not in regions
+
+    def test_get_env_partition_bootstrap_region_prefers_session_region(self):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            assert (
+                get_env_partition_bootstrap_region(AWS_REGION_GOV_CLOUD_US_WEST_1)
+                == AWS_REGION_GOV_CLOUD_US_WEST_1
+            )
+
+    def test_get_env_partition_bootstrap_region_without_session_region(self):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            assert (
+                get_env_partition_bootstrap_region() == AWS_REGION_GOV_CLOUD_US_EAST_1
+            )
+
+    def test_get_env_partition_bootstrap_region_without_partition(self):
+        with mock.patch.dict(os.environ, {"PROWLER_AWS_PARTITION": ""}, clear=False):
+            assert (
+                get_env_partition_bootstrap_region(AWS_REGION_GOV_CLOUD_US_WEST_1)
+                is None
+            )
+
+    def test_get_aws_region_for_sts_env_partition_prefers_session_region(self):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            assert (
+                get_aws_region_for_sts(AWS_REGION_GOV_CLOUD_US_WEST_1, None)
+                == AWS_REGION_GOV_CLOUD_US_WEST_1
+            )
+
+    def test_get_profile_region_env_partition_keeps_session_region_inside_partition(
+        self,
+    ):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            aws_session = session.Session(region_name=AWS_REGION_GOV_CLOUD_US_WEST_1)
+
+            assert (
+                AwsProvider.get_profile_region(aws_session)
+                == AWS_REGION_GOV_CLOUD_US_WEST_1
+            )
+
+    def test_get_profile_region_env_partition_ignores_session_region_outside_partition(
+        self,
+    ):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            aws_session = session.Session(region_name=AWS_REGION_US_EAST_1)
+
+            assert (
+                AwsProvider.get_profile_region(aws_session)
+                == AWS_REGION_GOV_CLOUD_US_EAST_1
+            )
+
+    def test_get_profile_region_env_partition_excluded_session_region_stays_in_partition(
+        self,
+    ):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            aws_session = session.Session(region_name=AWS_REGION_GOV_CLOUD_US_WEST_1)
+
+            assert (
+                AwsProvider.get_profile_region(
+                    aws_session, {AWS_REGION_GOV_CLOUD_US_WEST_1}
+                )
+                == AWS_REGION_GOV_CLOUD_US_EAST_1
+            )
+
+    def test_get_profile_region_env_partition_all_regions_excluded_stays_in_partition(
+        self,
+    ):
+        with mock.patch.dict(
+            os.environ,
+            {"PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION},
+            clear=False,
+        ):
+            aws_session = session.Session(region_name=AWS_REGION_GOV_CLOUD_US_WEST_1)
+            gov_cloud_regions = set(get_env_partition_regions())
+
+            assert (
+                AwsProvider.get_profile_region(aws_session, gov_cloud_regions)
+                == AWS_REGION_GOV_CLOUD_US_WEST_1
+            )
+
+    @mock_aws
+    def test_test_connection_env_partition_prefers_session_region(self):
+        with (
+            mock.patch.dict(
+                os.environ,
+                {
+                    "PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION,
+                    "AWS_DEFAULT_REGION": AWS_REGION_GOV_CLOUD_US_WEST_1,
+                },
+                clear=False,
+            ),
+            mock.patch.object(
+                AwsProvider,
+                "validate_credentials",
+                return_value=AWSCallerIdentity(
+                    user_id="test-user-id",
+                    account=AWS_ACCOUNT_NUMBER,
+                    arn=ARN(AWS_GOV_CLOUD_ACCOUNT_ARN),
+                    region=AWS_REGION_GOV_CLOUD_US_WEST_1,
+                ),
+            ) as mock_validate_credentials,
+        ):
+            connection = AwsProvider.test_connection(
+                aws_access_key_id="test-access-key",
+                aws_secret_access_key="test-secret-key",
+                raise_on_exception=False,
+            )
+
+            assert connection.is_connected
+            assert (
+                mock_validate_credentials.call_args.args[1]
+                == AWS_REGION_GOV_CLOUD_US_WEST_1
+            )
+
+    @mock_aws
+    def test_test_connection_role_env_partition_prefers_session_region(self):
+        with (
+            mock.patch.dict(
+                os.environ,
+                {
+                    "PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION,
+                    "AWS_DEFAULT_REGION": AWS_REGION_GOV_CLOUD_US_WEST_1,
+                },
+                clear=False,
+            ),
+            mock.patch.object(
+                AwsProvider,
+                "assume_role",
+                return_value=AWSCredentials(
+                    aws_access_key_id="assumed-access-key",
+                    aws_secret_access_key="assumed-secret-key",
+                    aws_session_token="assumed-session-token",
+                    expiration=datetime.now(),
+                ),
+            ) as mock_assume_role,
+            mock.patch.object(
+                AwsProvider,
+                "validate_credentials",
+                return_value=AWSCallerIdentity(
+                    user_id="test-user-id",
+                    account=AWS_ACCOUNT_NUMBER,
+                    arn=ARN(AWS_GOV_CLOUD_ACCOUNT_ARN),
+                    region=AWS_REGION_GOV_CLOUD_US_WEST_1,
+                ),
+            ),
+        ):
+            connection = AwsProvider.test_connection(
+                role_arn=f"arn:{AWS_GOV_CLOUD_PARTITION}:iam::{AWS_ACCOUNT_NUMBER}:role/test-role",
+                aws_access_key_id="test-access-key",
+                aws_secret_access_key="test-secret-key",
+                raise_on_exception=False,
+            )
+
+            assert connection.is_connected
+            assumed_role_info = mock_assume_role.call_args.args[1]
+            assert assumed_role_info.sts_region == AWS_REGION_GOV_CLOUD_US_WEST_1
+
+    @mock_aws
+    def test_setup_session_mfa_env_partition_prefers_session_region(self):
+        with (
+            mock.patch.dict(
+                os.environ,
+                {
+                    "PROWLER_AWS_PARTITION": AWS_GOV_CLOUD_PARTITION,
+                    "AWS_DEFAULT_REGION": AWS_REGION_GOV_CLOUD_US_WEST_1,
+                },
+                clear=False,
+            ),
+            mock.patch.object(
+                AwsProvider,
+                "input_role_mfa_token_and_code",
+                return_value=AWSMFAInfo(
+                    arn=f"arn:{AWS_GOV_CLOUD_PARTITION}:iam::{AWS_ACCOUNT_NUMBER}:mfa/test",
+                    totp="123456",
+                ),
+            ),
+            mock.patch.object(
+                AwsProvider,
+                "create_sts_session",
+                side_effect=AwsProvider.create_sts_session,
+            ) as mock_create_sts_session,
+        ):
+            AwsProvider.setup_session(
+                mfa=True,
+                aws_access_key_id="test-access-key",
+                aws_secret_access_key="test-secret-key",
+            )
+
+            assert (
+                mock_create_sts_session.call_args.args[1]
+                == AWS_REGION_GOV_CLOUD_US_WEST_1
+            )
+
     @mock_aws
     def test_test_connection_env_partition_mismatch(self):
         with (

tests/providers/aws/utils.py

@@ -51,6 +51,7 @@ AWS_REGION_CN_NORTH_1 = "cn-north-1"
 
 # Gov Cloud Regions
 AWS_REGION_GOV_CLOUD_US_EAST_1 = "us-gov-east-1"
+AWS_REGION_GOV_CLOUD_US_WEST_1 = "us-gov-west-1"
 
 # Iso Regions
 AWS_REGION_ISO_GLOBAL = "aws-iso-global"