cleanup more

local-app/python-tools/gfl-resource-actions/aws_resource_management/__init__.py

@@ -1,8 +1,19 @@
 """
 AWS Resource Management package.
-
-This package provides tools for managing AWS resources including stopping,
-starting, and listing resources across different AWS services.
 """
 
+# Re-export key modules for easier imports
+from aws_resource_management.logging_setup import (
+    LoggingContext,
+    configure_logging,
+    initialize_csv_log,
+    log_action_to_csv,
+    log_operation,
+    log_with_context,
+    setup_logging,
+)
+
+# Set up a default logger
+logger = setup_logging()
+
 __version__ = "0.1.0"

local-app/python-tools/gfl-resource-actions/aws_resource_management/aws_utils.py

@@ -8,11 +8,13 @@
 import os
 import threading
 import time
+from functools import lru_cache
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
 import boto3
 import botocore
 from aws_resource_management.logging_setup import setup_logging
+from botocore.exceptions import ClientError, NoCredentialsError, ProfileNotFound
 
 logger = setup_logging()
 
@@ -90,27 +92,22 @@ def safe_in(substring: Any, container: Any) -> bool:
 # ---------------------------------------------------------------------------
 
 
-def get_partition_for_region(region: Optional[str]) -> str:
-    """Get the AWS partition for a region (can be None)."""
-    # Default for GovCloud environment
-    DEFAULT_PARTITION = "aws-us-gov"
+@lru_cache(maxsize=128)
+def get_partition_for_region(region_name: str) -> str:
+    """
+    Determine AWS partition based on region name.
 
-    # Early return for None or non-string regions
-    if not is_string(region):
-        return DEFAULT_PARTITION
+    Args:
+        region_name: AWS region name
 
-    # Map region prefixes to partitions using simple lookups
-    if region.startswith("us-gov-"):
+    Returns:
+        AWS partition name (e.g., 'aws', 'aws-us-gov', 'aws-cn')
+    """
+    if region_name.startswith("us-gov-"):
         return "aws-us-gov"
-    elif region.startswith("cn-"):
+    elif region_name.startswith("cn-"):
         return "aws-cn"
-    elif any(
-        region.startswith(prefix)
-        for prefix in ("us-", "eu-", "ap-", "ca-", "sa-", "af-")
-    ):
-        return "aws"
-
-    return DEFAULT_PARTITION
+    return "aws"
 
 
 def get_regions_for_partition(partition: str) -> List[str]:
@@ -353,44 +350,120 @@ def detect_partition_from_credentials(credentials: Dict[str, str]) -> str:
 # ---------------------------------------------------------------------------
 
 
+@lru_cache(maxsize=128)
 def get_session_for_account(
     account_id: str, region: Optional[str] = None, profile_name: Optional[str] = None
 ) -> Optional[boto3.Session]:
-    """Get a boto3 session with improved caching to minimize API calls."""
+    """
+    Get a boto3 session for the specified account with proper role assumption.
+
+    Args:
+        account_id: AWS account ID
+        region: AWS region name (optional)
+        profile_name: AWS profile name to use (optional)
+
+    Returns:
+        Boto3 session with appropriate credentials
+    """
     cache_key = (
         f"session:{account_id}:{region or 'default'}:{profile_name or 'default'}"
     )
 
-    # Check cache first
+    # Return cached session if available
     with _session_cache_lock:
         if (
             cache_key in _session_cache
             and _session_cache[cache_key]["timestamp"] > time.time() - CACHE_EXPIRY
         ):
             return _session_cache[cache_key]["session"]
 
-    # Get credentials
-    credentials = get_credentials(account_id, profile_name)
-    if not credentials:
-        return None
-
-    # Create session
+    # Try getting a session using standard methods
     try:
-        session = boto3.Session(
-            aws_access_key_id=credentials["aws_access_key_id"],
-            aws_secret_access_key=credentials["aws_secret_access_key"],
-            aws_session_token=credentials.get("aws_session_token"),
-            region_name=region,
-        )
+        # Method 1: Try using a profile named after the account ID or the specified profile
+        try:
+            session = boto3.Session(
+                profile_name=profile_name or account_id, region_name=region
+            )
+            # Validate session by checking identity
+            sts = session.client("sts")
+            identity = sts.get_caller_identity()
+            if identity.get("Account") == account_id:
+                logger.debug(
+                    f"Using profile {profile_name or account_id} for authentication"
+                )
 
-        # Cache the session
-        with _session_cache_lock:
-            _session_cache[cache_key] = {"session": session, "timestamp": time.time()}
+                # Cache the session
+                with _session_cache_lock:
+                    _session_cache[cache_key] = {
+                        "session": session,
+                        "timestamp": time.time(),
+                    }
+
+                return session
+        except (ProfileNotFound, NoCredentialsError, ClientError):
+            # Profile not found or invalid, continue to next method
+            pass
+
+        # Method 2: Use credentials to create a session
+        credentials = get_credentials(account_id, profile_name)
+        if credentials:
+            session = boto3.Session(
+                aws_access_key_id=credentials["aws_access_key_id"],
+                aws_secret_access_key=credentials["aws_secret_access_key"],
+                aws_session_token=credentials.get("aws_session_token"),
+                region_name=region,
+            )
+
+            # Cache the session
+            with _session_cache_lock:
+                _session_cache[cache_key] = {
+                    "session": session,
+                    "timestamp": time.time(),
+                }
+
+            logger.debug(f"Created session for account {account_id} using credentials")
+            return session
+
+        # Method 3: Try to assume role in target account using default session
+        default_session = boto3.Session(region_name=region)
+        sts = default_session.client("sts")
+
+        # Try common cross-account roles
+        for role_name in ["OrganizationAccountAccessRole", "AWSControlTowerExecution"]:
+            try:
+                role_arn = f"arn:aws:iam::{account_id}:role/{role_name}"
+                response = sts.assume_role(
+                    RoleArn=role_arn, RoleSessionName="ResourceManagementSession"
+                )
+
+                # Create session with temporary credentials
+                credentials = response["Credentials"]
+                session = boto3.Session(
+                    aws_access_key_id=credentials["AccessKeyId"],
+                    aws_secret_access_key=credentials["SecretAccessKey"],
+                    aws_session_token=credentials["SessionToken"],
+                    region_name=region,
+                )
+
+                # Cache the session
+                with _session_cache_lock:
+                    _session_cache[cache_key] = {
+                        "session": session,
+                        "timestamp": time.time(),
+                    }
+
+                logger.debug(f"Assumed {role_name} in account {account_id}")
+                return session
+            except ClientError:
+                # Role assumption failed, try the next role
+                continue
+
+        # If all methods fail, raise exception
+        raise Exception(f"Could not authenticate to account {account_id}")
 
-        return session
     except Exception as e:
-        logger.error(f"Error creating session for account {account_id}: {str(e)}")
-        return None
+        logger.error(f"Failed to get session for account {account_id}: {str(e)}")
+        raise
 
 
 # ---------------------------------------------------------------------------
@@ -559,6 +632,47 @@ def get_enabled_regions(
         return get_regions_for_partition(partition)
 
 
+def list_enabled_regions(
+    session: boto3.Session, exclude_regions: Optional[List[str]] = None
+) -> List[str]:
+    """
+    List all enabled regions for the account, respecting partition.
+
+    Args:
+        session: Boto3 session
+        exclude_regions: List of regions to exclude
+
+    Returns:
+        List of enabled region names
+    """
+    if exclude_regions is None:
+        exclude_regions = []
+
+    try:
+        ec2 = session.client(
+            "ec2", region_name="us-east-1"
+        )  # Use US East 1 for global operations
+        regions_response = ec2.describe_regions(
+            AllRegions=False
+        )  # Only get enabled regions
+
+        regions = [
+            r["RegionName"]
+            for r in regions_response["Regions"]
+            if r["RegionName"] not in exclude_regions
+        ]
+        return regions
+    except Exception as e:
+        logger.warning(f"Could not retrieve regions, using defaults: {str(e)}")
+
+        # For GovCloud, provide sensible defaults
+        if session.region_name and session.region_name.startswith("us-gov-"):
+            return ["us-gov-east-1", "us-gov-west-1"]
+
+        # Default to common commercial regions
+        return ["us-east-1", "us-east-2", "us-west-1", "us-west-2"]
+
+
 def is_valid_region(region: str) -> bool:
     """Check if a region is valid without making an API call."""
     if not is_string(region):