fix(aws): try the rest of the partition when the bootstrap region is unreachable (#12822)

Co-authored-by: César Arroba <19954079+cesararroba@users.noreply.github.com>
Co-authored-by: pedrooot <pedromarting3@gmail.com>

Author: 3 people

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

@@ -39,6 +39,8 @@ It matters most where nothing else says. Resolving an identity means calling STS
 
 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.
 
+When no configured region says which one to prefer, the first region of the partition is tried, and up to two more follow if it cannot be reached. A network that routes to only one region of its partition therefore works without having to declare which one that is. Only a connection failure moves on to the next region: a credential error is reported from the first, since it would be the same everywhere. A region excluded from the scan is tried last, so it is avoided whenever another region of the partition answers.
+
 <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>

prowler/changelog.d/aws-partition-bootstrap-falls-back-to-the-next-region.fixed.md

@@ -0,0 +1 @@
+Bootstrap STS calls now try up to two more regions of the partition declared in `PROWLER_AWS_PARTITION` when the first one cannot be reached, so a deployment that routes to only one region of its partition no longer fails on an endpoint it has no path to. This covers validating credentials, assuming a role and getting an MFA session token

prowler/providers/aws/aws_provider.py

@@ -3,12 +3,19 @@
 from datetime import datetime
 from functools import lru_cache
 from re import fullmatch
-from typing import Optional
+from typing import Any, Callable, Optional
 
 from boto3.session import Session
 from botocore.config import Config
 from botocore.credentials import RefreshableCredentials
-from botocore.exceptions import ClientError, NoCredentialsError, ProfileNotFound
+from botocore.exceptions import (
+    ClientError,
+    ConnectTimeoutError,
+    EndpointConnectionError,
+    NoCredentialsError,
+    ProfileNotFound,
+    ReadTimeoutError,
+)
 from botocore.session import Session as BotocoreSession
 from colorama import Fore, Style
 from pytz import utc
@@ -263,7 +270,10 @@ def __init__(
         caller_identity = self.validate_credentials(
             session=self.session.current_session,
             aws_region=sts_region,
+            excluded_regions=excluded_regions,
         )
+        # Later STS calls go where validation got an answer, not where it timed out
+        sts_region = caller_identity.region
 
         logger.info("Credentials validated")
         ########
@@ -690,17 +700,19 @@ def setup_session(
                     or session.region_name
                     or AWS_STS_GLOBAL_ENDPOINT_REGION
                 )
-                sts_client = AwsProvider.create_sts_session(session, sts_region)
-
                 # TODO: pass values from the input
                 mfa_info = AwsProvider.input_role_mfa_token_and_code()
                 # TODO: validate MFA ARN here
                 get_session_token_arguments = {
                     "SerialNumber": mfa_info.arn,
                     "TokenCode": mfa_info.totp,
                 }
-                session_credentials = sts_client.get_session_token(
-                    **get_session_token_arguments
+                _, session_credentials = AwsProvider.sts_call_with_partition_failover(
+                    session,
+                    sts_region,
+                    lambda sts_client: sts_client.get_session_token(
+                        **get_session_token_arguments
+                    ),
                 )
                 mfa_session = Session(
                     aws_access_key_id=session_credentials["Credentials"]["AccessKeyId"],
@@ -1244,10 +1256,11 @@ def assume_role(
                 mfa_info = AwsProvider.input_role_mfa_token_and_code()
                 assume_role_arguments["SerialNumber"] = mfa_info.arn
                 assume_role_arguments["TokenCode"] = mfa_info.totp
-            sts_client = AwsProvider.create_sts_session(
-                session, assumed_role_info.sts_region
+            _, assumed_credentials = AwsProvider.sts_call_with_partition_failover(
+                session,
+                assumed_role_info.sts_region,
+                lambda sts_client: sts_client.assume_role(**assume_role_arguments),
             )
-            assumed_credentials = sts_client.assume_role(**assume_role_arguments)
             # Convert the UTC datetime object to your local timezone
             credentials_expiration_local_time = (
                 assumed_credentials["Credentials"]["Expiration"]
@@ -1326,30 +1339,98 @@ def get_checks_to_execute_by_audit_resources(self) -> set[str]:
             )
             raise error
 
+    @staticmethod
+    def sts_call_with_partition_failover(
+        session: Session,
+        aws_region: str,
+        operation: Callable[[Any], Any],
+        excluded_regions: set[str] | None = None,
+    ) -> tuple[str, Any]:
+        """
+        Run a bootstrap STS call, moving on when a region cannot be reached.
+
+        Bootstrap calls happen before anything is known about the credentials, so
+        the region they go to is a guess whenever none was configured. On a network
+        that routes to only one region of its partition that guess is fatal, and the
+        remaining regions of the partition declared in PROWLER_AWS_PARTITION are the
+        ones worth trying.
+
+        Args:
+            session (Session): The AWS session object.
+            aws_region (str): The region to try first.
+            operation (Callable[[Any], Any]): Receives an STS client and performs
+                the call.
+            excluded_regions (set[str] | None): Regions excluded from the scan,
+                tried after the rest of the partition.
+
+        Returns:
+            tuple[str, Any]: The region that answered and whatever the operation
+                returned.
+
+        Raises:
+            Exception: Whatever the operation raises, or the last connection error
+                when no region could be reached.
+        """
+        *fallback_regions, last_region = get_partition_bootstrap_candidates(
+            aws_region, session.region_name, excluded_regions
+        )
+
+        for candidate_region in fallback_regions:
+            try:
+                sts_client = AwsProvider.create_sts_session(session, candidate_region)
+                return candidate_region, operation(sts_client)
+            # The credentials are not at fault, so the next region is worth trying
+            except (
+                EndpointConnectionError,
+                ConnectTimeoutError,
+                ReadTimeoutError,
+            ) as unreachable:
+                logger.warning(
+                    f"{unreachable.__class__.__name__}[{unreachable.__traceback__.tb_lineno}]: {unreachable}"
+                )
+
+        # Nothing is left to try after the last region, so its error is the answer
+        sts_client = AwsProvider.create_sts_session(session, last_region)
+        return last_region, operation(sts_client)
+
     @staticmethod
     def validate_credentials(
         session: Session,
         aws_region: str,
+        excluded_regions: set[str] | None = None,
     ) -> AWSCallerIdentity:
         """
         Validates the AWS credentials using the provided session and AWS region.
+
+        When the region cannot be reached, the remaining regions of the partition
+        declared in PROWLER_AWS_PARTITION are tried before giving up. A credential
+        error is returned from the first region instead, since it would be the same
+        everywhere.
+
         Args:
             session (Session): The AWS session object.
             aws_region (str): The AWS region to validate the credentials.
+            excluded_regions (set[str] | None): Regions excluded from the scan,
+                tried after the rest of the partition.
         Returns:
-            AWSCallerIdentity: An object containing the caller identity information.
+            AWSCallerIdentity: An object containing the caller identity information,
+                including the region that answered.
         Raises:
             Exception: If an error occurs during the validation process.
         """
         try:
-            sts_client = AwsProvider.create_sts_session(session, aws_region)
-            caller_identity = sts_client.get_caller_identity()
+            sts_region, caller_identity = AwsProvider.sts_call_with_partition_failover(
+                session,
+                aws_region,
+                lambda sts_client: sts_client.get_caller_identity(),
+                excluded_regions,
+            )
             # Include the region where the caller_identity has validated the credentials
             return AWSCallerIdentity(
                 user_id=caller_identity.get("UserId"),
                 account=caller_identity.get("Account"),
                 arn=ARN(caller_identity.get("Arn")),
-                region=aws_region,
+                region=sts_region,
             )
         except ClientError as client_error:
             logger.error(
@@ -1846,6 +1927,49 @@ def get_env_partition_bootstrap_region(
     return regions[0] if regions else None
 
 
+# An unreachable endpoint costs a connection timeout, so a partition with many
+# regions is not walked in full
+MAX_STS_BOOTSTRAP_ATTEMPTS = 3
+
+
+def get_partition_bootstrap_candidates(
+    aws_region: str,
+    session_region: Optional[str] = None,
+    excluded_regions: set[str] | None = None,
+) -> list:
+    """
+    Get the STS bootstrap regions to try, in order, starting with the chosen one.
+
+    A deployment reached only through its own region's endpoints has no route to
+    the rest of its partition, and which region that is cannot be known from the
+    environment alone: a container may carry a region belonging to no partition
+    it scans. Offering the remaining regions of the declared partition lets the
+    bootstrap succeed without anything having to declare the right one.
+
+    Args:
+        aws_region (str): The region already chosen for the bootstrap call.
+        session_region (Optional[str]): The region of the AWS session.
+        excluded_regions (set[str] | None): Regions excluded from the scan. They
+            go after the rest of the partition, so the bootstrap avoids them
+            whenever another region answers and still has them as a last resort.
+
+    Returns:
+        list: The regions to try, preferred first, capped at
+            MAX_STS_BOOTSTRAP_ATTEMPTS.
+    """
+    excluded_regions = set(excluded_regions or ())
+    partition_regions = get_env_partition_regions(session_region) or []
+    # sorted() is stable, so the partition order survives on each side of the split
+    ordered_regions = sorted(
+        partition_regions, key=lambda region: region in excluded_regions
+    )
+    candidates = [aws_region]
+    for region in ordered_regions:
+        if region not in candidates:
+            candidates.append(region)
+    return candidates[:MAX_STS_BOOTSTRAP_ATTEMPTS]
+
+
 # TODO: This can be moved to another class since it doesn't need self
 def get_aws_region_for_sts(
     session_region: str,