diff --git a/local-app/python-tools/gfl-resource-actions/README.md b/local-app/python-tools/gfl-resource-actions/README.md index 41167bed..083afbbf 100644 --- a/local-app/python-tools/gfl-resource-actions/README.md +++ b/local-app/python-tools/gfl-resource-actions/README.md @@ -29,9 +29,9 @@ pip install boto3 ### Basic Usage ```python -import aws_utils +import aws_resource_management.aws_utils as aws_utils -# Create a client using SSO profiles or role assumption +# Get an EC2 client for a specific account ec2_client = aws_utils.create_boto3_client_for_account( '123456789012', # AWS account ID 'ec2', # AWS service diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/aws_utils.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/aws_utils.py index fc0c15e8..e39baf71 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/aws_utils.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/aws_utils.py @@ -1,207 +1,241 @@ """ AWS utility functions for credential management and account discovery. +Simplified version with reduced API calls and streamlined functionality. """ -import concurrent.futures import configparser import logging import os import threading -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Any, Dict, List, Optional, Tuple +import time +from typing import Any, Dict, List, Optional, Set, Tuple, Union import boto3 import botocore from aws_resource_management.logging_setup import setup_logging logger = setup_logging() -# Thread-local storage for session caching -_thread_local = threading.local() +# --------------------------------------------------------------------------- +# Global caches to reduce API calls +# --------------------------------------------------------------------------- +# Thread-safe session cache +_session_cache = {} +_session_cache_lock = threading.Lock() + +# Region and partition information cache +_region_cache = { + "timestamp": 0, + "all_regions": {}, + "enabled_regions": {}, + "partition_regions": { + "aws-us-gov": ["us-gov-east-1", "us-gov-west-1"], + "aws-cn": ["cn-north-1", "cn-northwest-1"], + "aws": [ + "us-east-1", + "us-east-2", + "us-west-1", + "us-west-2", + "eu-west-1", + "eu-central-1", + "eu-west-2", + "eu-west-3", + "ap-northeast-1", + "ap-northeast-2", + "ap-southeast-1", + "ap-southeast-2", + ], + }, +} + +# Cache expiration time in seconds (1 hour) +CACHE_EXPIRY = 3600 + +# --------------------------------------------------------------------------- +# String utility functions for safer string handling +# --------------------------------------------------------------------------- + + +def is_string(value: Any) -> bool: + """Check if a value is a string.""" + return isinstance(value, str) + + +def safe_startswith(value: Any, prefix: Union[str, Tuple[str, ...]]) -> bool: + """Safely check if a value starts with a prefix, handling None values.""" + if not is_string(value): + return False -def get_account_list() -> List[Dict[str, str]]: - """ - Get list of accounts from AWS Organizations. + if isinstance(prefix, str): + return value.startswith(prefix) + elif isinstance(prefix, tuple) and all(is_string(p) for p in prefix): + return value.startswith(prefix) - Returns: - List of dictionaries containing account_id and account_name - """ - try: - # First try to use organizations API - session = boto3.Session() - org_client = session.client("organizations") + return False - accounts = [] - paginator = org_client.get_paginator("list_accounts") - for page in paginator.paginate(): - for account in page["Accounts"]: - if account["Status"] == "ACTIVE": - accounts.append( - {"account_id": account["Id"], "account_name": account["Name"]} - ) +def safe_in(substring: Any, container: Any) -> bool: + """Safely check if a substring is in a container, handling None values.""" + if substring is None or container is None: + return False - return accounts + try: + return substring in container + except (TypeError, ValueError): + return False - except (botocore.exceptions.ClientError, botocore.exceptions.BotoCoreError) as e: - logger.warning(f"Failed to get accounts from Organizations API: {str(e)}") - logger.info("Falling back to configured AWS profiles") - # Fall back to configured profiles if Organizations API access fails - return get_accounts_from_profiles() +# --------------------------------------------------------------------------- +# Simplified region and partition handling +# --------------------------------------------------------------------------- -def get_accounts_from_profiles() -> List[Dict[str, str]]: - """ - Get list of accounts from configured AWS profiles in ~/.aws/config. +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" - Returns: - List of dictionaries containing account_id and account_name - """ - accounts = [] - # Dictionary to store the first profile seen for each account ID - account_profiles = {} + # Early return for None or non-string regions + if not is_string(region): + return DEFAULT_PARTITION - try: - # Read AWS config file - config_path = os.path.expanduser("~/.aws/config") - if not os.path.exists(config_path): - logger.warning(f"AWS config file not found at {config_path}") - return [] + # Map region prefixes to partitions using simple lookups + if region.startswith("us-gov-"): + return "aws-us-gov" + elif region.startswith("cn-"): + return "aws-cn" + elif any( + region.startswith(prefix) + for prefix in ("us-", "eu-", "ap-", "ca-", "sa-", "af-") + ): + return "aws" - config = configparser.ConfigParser() - config.read(config_path) + return DEFAULT_PARTITION - # Extract account information from profiles - for section in config.sections(): - # Skip sections that don't have account ID - if not config.has_option(section, "sso_account_id"): - continue - account_id = config.get(section, "sso_account_id") - section_name = section - if section.startswith("profile "): - section_name = section[8:] # Remove "profile " prefix - - # For each account ID, remember only the first profile we encounter - # Unless it's a profile with a preferred role like AdministratorAccess or inf-admin-t2 - if account_id not in account_profiles or ( - config.has_option(section, "sso_role_name") - and config.get(section, "sso_role_name") - in ["AdministratorAccess", "inf-admin-t2"] - ): - # Get account name if available, otherwise use account ID - account_name = account_id - if config.has_option(section, "sso_account_name"): - account_name = config.get(section, "sso_account_name") +def get_regions_for_partition(partition: str) -> List[str]: + """Get regions for a partition using cached data instead of API calls.""" + if partition in _region_cache["partition_regions"]: + return _region_cache["partition_regions"][partition] - account_profiles[account_id] = { - "account_id": account_id, - "account_name": account_name, - "profile": section_name, - } + # Default to GovCloud + return _region_cache["partition_regions"]["aws-us-gov"] - logger.debug( - f"Using profile {section_name} for account {account_id} ({account_name})" - ) - # Convert dictionary values to list - accounts = list(account_profiles.values()) +def get_all_regions(partition: Optional[str] = None) -> List[str]: + """Get all regions, using cache to minimize API calls.""" + # Use predefined regions for special partitions + if partition in ("aws-us-gov", "aws-cn"): + return get_regions_for_partition(partition) + + # Check cache first + current_time = time.time() + cache_key = partition or "all" + + if ( + cache_key in _region_cache["all_regions"] + and current_time - _region_cache["timestamp"] < CACHE_EXPIRY + ): + return _region_cache["all_regions"][cache_key] + + # Cache miss: fetch regions + try: + ec2 = boto3.client("ec2", region_name="us-east-1") + all_regions = [ + region["RegionName"] for region in ec2.describe_regions()["Regions"] + ] - logger.info(f"Found {len(accounts)} accounts from AWS profiles") - return accounts + # Cache the full list + _region_cache["all_regions"]["all"] = all_regions + _region_cache["timestamp"] = current_time + # If partition specified, filter and cache that too + if partition: + filtered_regions = [ + r + for r in all_regions + if is_string(r) and get_partition_for_region(r) == partition + ] + _region_cache["all_regions"][partition] = filtered_regions + return filtered_regions + + return all_regions except Exception as e: - logger.error(f"Error reading AWS profiles: {str(e)}") - return [] + logger.warning(f"Error getting regions: {str(e)}") + return get_regions_for_partition(partition or "aws-us-gov") + + +# --------------------------------------------------------------------------- +# Simplified credential management +# --------------------------------------------------------------------------- def get_credentials( - account_id: str, profile_name: Optional[str] = None + account_id: str, profile_name: Optional[str] = None, region: Optional[str] = None ) -> Optional[Dict[str, Any]]: - """ - Get AWS credentials for the specified account. + """Get AWS credentials with simplified logic and better caching.""" + cache_key = f"creds:{account_id}:{profile_name or 'default'}" - Args: - account_id: AWS account ID - profile_name: Optional AWS profile name to use + 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]["credentials"] - Returns: - Dict with AWS credentials or None if failed - """ - try: - # If profile name is provided, use it directly - if profile_name: - logger.debug(f"Using specified profile: {profile_name}") - try: - session = boto3.Session(profile_name=profile_name) + # Try profile approach first + credentials = _get_credentials_from_profile(account_id, profile_name) + if credentials: + _cache_credentials(cache_key, credentials) + return credentials - # Verify we can get credentials from the session - if session.get_credentials() is None: - logger.warning( - f"No credentials available for profile {profile_name}" - ) - return None + # Fall back to role assumption + credentials = _assume_role_for_account(account_id) + if credentials: + _cache_credentials(cache_key, credentials) + return credentials - return { - "aws_access_key_id": session.get_credentials().access_key, - "aws_secret_access_key": session.get_credentials().secret_key, - "aws_session_token": session.get_credentials().token, - } - except ( - botocore.exceptions.ProfileNotFound, - botocore.exceptions.ClientError, - ) as e: - logger.warning(f"Error using profile {profile_name}: {str(e)}") - # Fall through to try finding another profile - - # Find profile for the specified account - matching_profile = find_profile_by_account_id(account_id) - if matching_profile: - logger.debug(f"Using profile {matching_profile} for account {account_id}") - try: - session = boto3.Session(profile_name=matching_profile) + return None - # Verify we can get credentials from the session - if session.get_credentials() is None: - logger.warning( - f"No credentials available for profile {matching_profile}" - ) - return None - return { - "aws_access_key_id": session.get_credentials().access_key, - "aws_secret_access_key": session.get_credentials().secret_key, - "aws_session_token": session.get_credentials().token, - } - except ( - botocore.exceptions.ProfileNotFound, - botocore.exceptions.ClientError, - ) as e: - logger.warning(f"Error using profile {matching_profile}: {str(e)}") - # Fall through to role assumption - - # If no profile is found or profile didn't work, fall back to role assumption - logger.debug( - f"No working profile found for account {account_id}, falling back to role assumption" - ) - return assume_role_credentials(account_id) +def _cache_credentials(cache_key: str, credentials: Dict[str, Any]) -> None: + """Cache credentials with timestamp.""" + with _session_cache_lock: + _session_cache[cache_key] = { + "credentials": credentials, + "timestamp": time.time(), + } - except Exception as e: - logger.error(f"Error getting credentials for account {account_id}: {str(e)}") - return None +def _get_credentials_from_profile( + account_id: str, profile_name: Optional[str] = None +) -> Optional[Dict[str, Any]]: + """Get credentials from a profile.""" + try: + # Use specified profile or find one + profile_to_use = profile_name or _find_profile_for_account(account_id) + if not profile_to_use: + return None -def find_profile_by_account_id(account_id: str) -> Optional[str]: - """ - Find an AWS profile name that matches the specified account ID. + # Get credentials from the profile + session = boto3.Session(profile_name=profile_to_use) + if session.get_credentials() is None: + return None - Args: - account_id: AWS account ID + return { + "aws_access_key_id": session.get_credentials().access_key, + "aws_secret_access_key": session.get_credentials().secret_key, + "aws_session_token": session.get_credentials().token, + } + except Exception as e: + logger.debug( + f"Error getting credentials from profile for {account_id}: {str(e)}" + ) + return None - Returns: - Profile name or None if not found - """ + +def _find_profile_for_account(account_id: str) -> Optional[str]: + """Find an AWS profile for an account with simplified logic.""" try: config_path = os.path.expanduser("~/.aws/config") if not os.path.exists(config_path): @@ -210,69 +244,44 @@ def find_profile_by_account_id(account_id: str) -> Optional[str]: config = configparser.ConfigParser() config.read(config_path) - # Try to find profiles in order of preference - preferred_profiles = [ + # Prioritized profile patterns + profile_patterns = [ f"{account_id}.AdministratorAccess", f"{account_id}.inf-admin-t2", f"{account_id}-gov.administratoraccess", f"{account_id}-gov.inf-admin-t2", ] - # First check for preferred profiles - for profile_name in preferred_profiles: - for section in config.sections(): - section_name = section - if section.startswith("profile "): - section_name = section[8:] # Remove "profile " prefix - - if section_name.lower() == profile_name.lower(): - logger.debug( - f"Found preferred profile {section_name} for account {account_id}" - ) - return section_name - - # If no preferred profile found, look for any profile with this account ID + # Check for exact matches of preferred profiles first for section in config.sections(): - if not config.has_option(section, "sso_account_id"): - continue - - profile_account_id = config.get(section, "sso_account_id") - if profile_account_id == account_id: - section_name = section - if section.startswith("profile "): - section_name = section[8:] # Remove "profile " prefix + section_name = ( + section[8:] if safe_startswith(section, "profile ") else section + ) - logger.debug(f"Found profile {section_name} for account {account_id}") + if section_name.lower() in [p.lower() for p in profile_patterns]: return section_name - logger.debug(f"No profile found for account {account_id}") - return None + # Then check for any profile with matching account ID + for section in config.sections(): + if ( + config.has_option(section, "sso_account_id") + and config.get(section, "sso_account_id") == account_id + ): + return section[8:] if safe_startswith(section, "profile ") else section - except Exception as e: - logger.error(f"Error finding profile for account {account_id}: {str(e)}") + return None + except Exception: return None -def assume_role_credentials(account_id: str) -> Optional[Dict[str, Any]]: - """ - Get credentials by assuming a role in the target account. - - Args: - account_id: AWS account ID - - Returns: - Dict with AWS credentials or None if failed - """ +def _assume_role_for_account(account_id: str) -> Optional[Dict[str, Any]]: + """Assume a role to get credentials.""" try: - sts_client = boto3.client("sts") - - # Try common role names - role_names = ["OrganizationAccountAccessRole", "AWSControlTowerExecution"] - - for role_name in role_names: + # Try common role names in order of preference + for role_name in ["OrganizationAccountAccessRole", "AWSControlTowerExecution"]: try: role_arn = f"arn:aws:iam::{account_id}:role/{role_name}" - response = sts_client.assume_role( + response = boto3.client("sts").assume_role( RoleArn=role_arn, RoleSessionName="ResourceManagementSession" ) @@ -285,329 +294,313 @@ def assume_role_credentials(account_id: str) -> Optional[Dict[str, Any]]: except botocore.exceptions.ClientError: continue - logger.warning(f"Could not assume any role in account {account_id}") return None - - except Exception as e: - logger.error(f"Error assuming role for account {account_id}: {str(e)}") + except Exception: return None -def get_partition_for_region(region: str) -> str: - """ - Get the AWS partition for a region. - - Args: - region: AWS region name - - Returns: - AWS partition (aws, aws-us-gov, etc.) - """ - # Default to commercial AWS - partition = "aws" - - # Check for GovCloud regions - if region.startswith("us-gov"): - partition = "aws-us-gov" - elif region.startswith("cn-"): - partition = "aws-cn" - - return partition - - -def get_all_regions(partition: Optional[str] = None) -> List[str]: - """ - Get all available AWS regions, filtered by partition if specified. +def detect_partition_from_credentials(credentials: Dict[str, str]) -> str: + """Detect partition from credentials with minimal API calls.""" + if not credentials: + return "aws-us-gov" # Default for environment + + # Check credential format first (fast, no API calls) + access_key = str(credentials.get("aws_access_key_id", "")).lower() + session_token = str(credentials.get("aws_session_token", "")).lower() + + # Check for GovCloud indicators + if any( + safe_in(indicator, access_key) or safe_in(indicator, session_token) + for indicator in ["gov", "usgovcloud"] + ): + return "aws-us-gov" - Args: - partition: Optional AWS partition to filter by + # Check for China indicators + if any( + safe_in(indicator, access_key) or safe_in(indicator, session_token) + for indicator in ["cn-", "china"] + ): + return "aws-cn" - Returns: - List of region names - """ + # Try one API call to most likely partition based on environment try: - if partition == "aws-us-gov": - return ["us-gov-east-1", "us-gov-west-1"] - elif partition == "aws-cn": - return ["cn-north-1", "cn-northwest-1"] - - # For standard AWS or no partition specified, try to get all regions - ec2 = boto3.client("ec2", region_name="us-east-1") - regions = [region["RegionName"] for region in ec2.describe_regions()["Regions"]] + boto3.client( + "sts", + region_name="us-gov-west-1", # Try GovCloud first based on environment + aws_access_key_id=credentials.get("aws_access_key_id"), + aws_secret_access_key=credentials.get("aws_secret_access_key"), + aws_session_token=credentials.get("aws_session_token"), + ).get_caller_identity() + return "aws-us-gov" + except Exception: + # If GovCloud fails, try standard AWS + try: + boto3.client( + "sts", + region_name="us-east-1", + aws_access_key_id=credentials.get("aws_access_key_id"), + aws_secret_access_key=credentials.get("aws_secret_access_key"), + aws_session_token=credentials.get("aws_session_token"), + ).get_caller_identity() + return "aws" + except Exception: + # Default to GovCloud for environment + return "aws-us-gov" - # If a partition was specified, filter the results - if partition == "aws": - regions = [ - r - for r in regions - if not (r.startswith("us-gov-") or r.startswith("cn-")) - ] - return regions - except Exception as e: - logger.warning(f"Error getting all regions: {e}") - # Default to common regions based on partition - if partition == "aws-us-gov": - return ["us-gov-east-1", "us-gov-west-1"] - elif partition == "aws-cn": - return ["cn-north-1", "cn-northwest-1"] - else: - return ["us-east-1", "us-east-2", "us-west-1", "us-west-2"] +# --------------------------------------------------------------------------- +# Simplified session management +# --------------------------------------------------------------------------- -def get_enabled_regions( - credentials: Dict[str, str], partition: Optional[str] = None -) -> List[str]: - """ - Get list of AWS regions enabled for the account. +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.""" + cache_key = ( + f"session:{account_id}:{region or 'default'}:{profile_name or 'default'}" + ) - Args: - credentials: AWS credentials dictionary - partition: AWS partition (aws, aws-us-gov, aws-cn) + # Check cache first + 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 - Returns: - List of enabled region names - """ + # Create session try: - # If no partition was specified, detect it from credentials - if not partition: - partition = detect_partition_from_credentials(credentials) - - # For GovCloud or China partitions, return their standard regions - # as the describe_regions call might not work with these credentials - if partition == "aws-us-gov": - return ["us-gov-east-1", "us-gov-west-1"] - elif partition == "aws-cn": - return ["cn-north-1", "cn-northwest-1"] - - # For standard partition, try to discover all enabled regions session = boto3.Session( - aws_access_key_id=credentials.get("aws_access_key_id"), - aws_secret_access_key=credentials.get("aws_secret_access_key"), + 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, ) - ec2_client = session.client("ec2", region_name="us-east-1") - regions = [ - region["RegionName"] for region in ec2_client.describe_regions()["Regions"] - ] - - # Filter to only enabled regions - enabled_regions = [] - for region in regions: - if _is_region_enabled(region, credentials): - enabled_regions.append(region) + # Cache the session + with _session_cache_lock: + _session_cache[cache_key] = {"session": session, "timestamp": time.time()} - return enabled_regions + return session except Exception as e: - logger.warning(f"Error getting all regions: {str(e)}") - # Fallback to default regions based on partition - if partition == "aws-us-gov": - return ["us-gov-east-1", "us-gov-west-1"] - elif partition == "aws-cn": - return ["cn-north-1", "cn-northwest-1"] - return ["us-east-1", "us-east-2", "us-west-1", "us-west-2"] + logger.error(f"Error creating session for account {account_id}: {str(e)}") + return None -def _is_region_enabled(region: str, credentials: Dict[str, str]) -> bool: - """ - Check if a region is enabled for the account. +# --------------------------------------------------------------------------- +# Account discovery with caching +# --------------------------------------------------------------------------- - Args: - region: AWS region name - credentials: AWS credentials dictionary - Returns: - True if the region is enabled, False otherwise - """ +def get_account_list() -> List[Dict[str, str]]: + """Get AWS accounts with caching to minimize API calls.""" + cache_key = "accounts" + + # Check cache first + 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]["accounts"] + + # Try Organizations API first try: - session = boto3.Session( - aws_access_key_id=credentials.get("aws_access_key_id"), - aws_secret_access_key=credentials.get("aws_secret_access_key"), - aws_session_token=credentials.get("aws_session_token"), - ) + session = boto3.Session() + accounts = [] + for account in ( + session.client("organizations").get_paginator("list_accounts").paginate() + ): + accounts.extend( + [ + {"account_id": acc["Id"], "account_name": acc["Name"]} + for acc in account["Accounts"] + if acc["Status"] == "ACTIVE" + ] + ) + + # Cache the accounts + with _session_cache_lock: + _session_cache[cache_key] = {"accounts": accounts, "timestamp": time.time()} - # Try to create a client and make a lightweight API call - ec2_client = session.client("ec2", region_name=region) - try: - # Try a lightweight call first - ec2_client.describe_regions(RegionNames=[region]) - return True - except Exception: - # If that fails, try another lightweight call - try: - ec2_client.describe_availability_zones(ZoneNames=[f"{region}a"]) - return True - except Exception: - # Both calls failed, region might not be enabled - return False + return accounts except Exception: - # Could not create client or all calls failed - return False + # Fall back to profiles + accounts = _get_accounts_from_profiles() + # Cache these accounts too + with _session_cache_lock: + _session_cache[cache_key] = {"accounts": accounts, "timestamp": time.time()} -def is_valid_region(region: str) -> bool: - """ - Check if the given region is valid. - - Args: - region: AWS region name - - Returns: - True if valid, False otherwise - """ - try: - boto3.session.Session().client("ec2", region_name=region) - return True - except botocore.exceptions.ClientError: - return False + return accounts -def get_session_for_account( - account_id: str, region: str = None, profile_name: Optional[str] = None -) -> Optional[boto3.Session]: +def get_organization_accounts() -> List[Dict[str, str]]: """ - Get or create a boto3 session for the specified account and region. - Uses thread-local storage to cache sessions for better performance. - - Args: - account_id: AWS account ID - region: Region name (optional) - profile_name: AWS profile name (optional) + Get list of accounts from AWS Organization. Returns: - boto3.Session object or None if failed + List of dictionaries with account information """ - # Get thread-local storage for sessions - if not hasattr(_thread_local, "sessions"): - _thread_local.sessions = {} + try: + accounts = _get_accounts_from_profiles() + return accounts + except Exception as e: + logger.error(f"Error getting accounts from profiles: {str(e)}") + return [] - # Create a key from account ID, region, and profile name - key = f"{account_id}:{region or 'default'}:{profile_name or 'default'}" - # Return cached session if it exists - if key in _thread_local.sessions: - return _thread_local.sessions[key] +def _get_accounts_from_profiles() -> List[Dict[str, str]]: + """Get accounts from local AWS config with simplified logic.""" + accounts_by_id = {} # Use dict to avoid duplicates try: - # Get credentials for the account - credentials = get_credentials(account_id, profile_name=profile_name) - if not credentials: - return None + config_path = os.path.expanduser("~/.aws/config") + if not os.path.exists(config_path): + return [] - # Create a new session - 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, - ) + config = configparser.ConfigParser() + config.read(config_path) - # Cache the session - _thread_local.sessions[key] = session - return session + for section in config.sections(): + if not config.has_option(section, "sso_account_id"): + continue + account_id = config.get(section, "sso_account_id") + account_name = ( + config.get(section, "sso_account_name") + if config.has_option(section, "sso_account_name") + else account_id + ) + + # Get profile name + profile = section[8:] if safe_startswith(section, "profile ") else section + + # Keep track if we should replace an existing entry + replace_existing = account_id in accounts_by_id + if replace_existing and config.has_option(section, "sso_role_name"): + role_name = config.get(section, "sso_role_name") + replace_existing = role_name in ["AdministratorAccess", "inf-admin-t2"] + + # Store or replace account info + if replace_existing or account_id not in accounts_by_id: + accounts_by_id[account_id] = { + "account_id": account_id, + "account_name": account_name, + "profile": profile, + } + + return list(accounts_by_id.values()) except Exception as e: - logger.error(f"Error creating session for account {account_id}: {str(e)}") - return None + logger.error(f"Error reading AWS profiles: {str(e)}") + return [] -def detect_partition_from_credentials(credentials: Dict[str, str]) -> str: - """ - Detect which AWS partition these credentials are valid for. - Tests credentials against key regions from different partitions to determine where they're valid. +# --------------------------------------------------------------------------- +# Simplified region utilities +# --------------------------------------------------------------------------- - Args: - credentials: AWS credentials dictionary - Returns: - AWS partition name (aws, aws-us-gov, aws-cn) - """ - logger.info("Detecting AWS partition from credentials...") +def get_enabled_regions( + credentials: Dict[str, str], partition: Optional[str] = None +) -> List[str]: + """Get enabled regions with minimized API calls using cache.""" + # Determine partition if not specified + if not partition: + partition = detect_partition_from_credentials(credentials) + + # For GovCloud/China, return predefined regions to avoid API calls + if partition in ("aws-us-gov", "aws-cn"): + return get_regions_for_partition(partition) + + # Generate cache key from credentials + creds_hash = hash( + f"{credentials.get('aws_access_key_id', '')}:{credentials.get('aws_secret_access_key', '')}" + ) + cache_key = f"enabled_regions:{creds_hash}:{partition}" - # Try GovCloud first since that's a common use case in this system - try: - client = boto3.client( - "sts", - region_name="us-gov-west-1", - aws_access_key_id=credentials.get("aws_access_key_id"), - aws_secret_access_key=credentials.get("aws_secret_access_key"), - aws_session_token=credentials.get("aws_session_token"), - ) - client.get_caller_identity() - logger.info("Credentials valid for GovCloud partition (aws-us-gov)") - return "aws-us-gov" - except Exception: - logger.debug("Credentials not valid for GovCloud partition") + # Check cache + with _session_cache_lock: + if ( + cache_key in _region_cache["enabled_regions"] + and time.time() - _region_cache["timestamp"] < CACHE_EXPIRY + ): + return _region_cache["enabled_regions"][cache_key] - # Try commercial AWS + # Make a single API call to describe_regions rather than checking each region try: - client = boto3.client( - "sts", - region_name="us-east-1", + session = boto3.Session( aws_access_key_id=credentials.get("aws_access_key_id"), aws_secret_access_key=credentials.get("aws_secret_access_key"), aws_session_token=credentials.get("aws_session_token"), ) - client.get_caller_identity() - logger.info("Credentials valid for standard AWS partition (aws)") - return "aws" + + regions = [ + region["RegionName"] + for region in session.client( + "ec2", region_name="us-east-1" + ).describe_regions()["Regions"] + ] + + # Cache the result + with _session_cache_lock: + _region_cache["enabled_regions"][cache_key] = regions + _region_cache["timestamp"] = time.time() + + return regions except Exception: - logger.debug("Credentials not valid for standard AWS partition") + # Fall back to default regions + return get_regions_for_partition(partition) + - # Try China partition +def is_valid_region(region: str) -> bool: + """Check if a region is valid without making an API call.""" + if not is_string(region): + return False + + # Check against our cached region lists + for partition, regions in _region_cache["partition_regions"].items(): + if region in regions: + return True + + # If not in predefined lists, make an API call as last resort try: - client = boto3.client( - "sts", - region_name="cn-north-1", - aws_access_key_id=credentials.get("aws_access_key_id"), - aws_secret_access_key=credentials.get("aws_secret_access_key"), - aws_session_token=credentials.get("aws_session_token"), - ) - client.get_caller_identity() - logger.info("Credentials valid for China AWS partition (aws-cn)") - return "aws-cn" - except Exception: - logger.debug("Credentials not valid for China AWS partition") + boto3.session.Session().client("ec2", region_name=region) + return True + except: + return False - # Default to standard AWS if we couldn't determine - logger.warning( - "Could not determine valid partition for credentials - defaulting to aws" - ) - return "aws" +# --------------------------------------------------------------------------- +# Function aliases for backward compatibility +# --------------------------------------------------------------------------- + +# Alias get_enabled_regions as get_available_regions for backward compatibility +get_available_regions = get_enabled_regions -def get_regions_for_partition(partition: str) -> List[str]: - """ - Get list of regions available in the specified AWS partition. + +# Ensure create_boto3_client function exists +def create_boto3_client(service, credentials, region=None): + """Create a boto3 client with provided credentials. Args: - partition: AWS partition (aws, aws-us-gov, aws-cn) + service: AWS service name + credentials: Dict containing AWS credentials + region: AWS region name Returns: - List of region names for the partition + Boto3 client """ - if partition == "aws-us-gov": - return ["us-gov-east-1", "us-gov-west-1"] - elif partition == "aws-cn": - return ["cn-north-1", "cn-northwest-1"] - else: # Default to commercial AWS regions - try: - ec2 = boto3.client("ec2", region_name="us-east-1") - regions = [ - region["RegionName"] for region in ec2.describe_regions()["Regions"] - ] - return regions - except Exception as e: - logger.warning(f"Could not fetch regions for partition {partition}: {e}") - # Return common commercial regions as fallback - return [ - "us-east-1", - "us-east-2", - "us-west-1", - "us-west-2", - "eu-west-1", - "eu-central-1", - "ap-southeast-1", - "ap-southeast-2", - ] + return boto3.client( + service, + region_name=region, + aws_access_key_id=credentials.get("aws_access_key_id"), + aws_secret_access_key=credentials.get("aws_secret_access_key"), + aws_session_token=credentials.get("aws_session_token"), + ) diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/cli_optimized.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/cli_optimized.py index 7fb7d46a..f981bc8a 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/cli_optimized.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/cli_optimized.py @@ -30,7 +30,7 @@ def parse_args(): # Resource selection parser.add_argument( "--resource-type", - choices=["ec2", "rds", "eks", "emr", "all"], + choices=["ecr", "ec2", "rds", "eks", "emr", "all"], default="all", help="Resource type to manage (default: all)", ) @@ -65,7 +65,8 @@ def parse_args(): parser.add_argument( "--partition", choices=["aws", "aws-us-gov", "aws-cn"], - help="AWS partition to operate in", + default="aws-us-gov", # Added missing comma + help="AWS partition to operate in (default: aws-us-gov)", ) # Authentication @@ -98,7 +99,7 @@ def main(): exclude_resources = [] if args.resource_type != "all": # If specific resource type is selected, exclude all others - all_resource_types = ["ec2", "rds", "eks", "emr"] + all_resource_types = ["ecr", "ec2", "rds", "eks", "emr"] exclude_resources = [ rt for rt in all_resource_types if rt != args.resource_type ] diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/core.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/core.py index 5ab3a97d..133d0ff5 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/core.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/core.py @@ -2,17 +2,19 @@ Core business logic for AWS Resource Management tool. """ +import concurrent.futures import logging import sys -from typing import Any, Dict, List, Optional, Union +from collections import defaultdict +from typing import Any, Dict, List, Optional, Tuple, Union from aws_resource_management import utils from aws_resource_management.aws_utils import ( detect_partition_from_credentials, get_account_list, - get_accounts_from_profiles, + get_available_regions, get_credentials, - get_enabled_regions, + get_organization_accounts, ) from aws_resource_management.config_manager import get_config from aws_resource_management.discovery import get_account_resources @@ -23,13 +25,19 @@ EMRManager, RDSManager, ) -from aws_resource_management.reporting import print_resource_summary, initialize_stats +from aws_resource_management.reporting import initialize_stats, print_resource_summary +from botocore.exceptions import ClientError logger = setup_logging() config = get_config() + class ResourceManager: """Main resource manager that orchestrates operations across accounts and resources.""" + + # Resource type constants for cleaner checking + RESOURCE_TYPES = ["ec2", "rds", "eks", "emr", "ecr"] + def __init__( self, profile_name: Optional[str] = None, @@ -37,26 +45,15 @@ def __init__( discover_regions: bool = False, partition: Optional[str] = None, ) -> None: - """ - Initialize the resource manager. - - Args: - profile_name: Optional AWS profile name to use for all operations - use_profiles: Whether to use AWS profiles from ~/.aws/config - discover_regions: Whether to automatically discover enabled regions - partition: AWS partition to operate in (aws, aws-us-gov, aws-cn) - """ + """Initialize the resource manager.""" self.config = get_config() - # Initialize account list self.account_list = [] self.profile_name = profile_name self.use_profiles = use_profiles self.discover_regions = discover_regions self.partition = partition - # Cache for discovered regions by account self.region_cache = {} - # If no partition was specified, we'll auto-detect based on credentials - # when we process the first account + self.MAX_CONCURRENT_ACCOUNTS = 3 # Limit concurrent account processing def process_accounts( self, @@ -68,489 +65,491 @@ def process_accounts( dry_run: bool = False, stats: Dict[str, Any] = None, ) -> Dict[str, Any]: - """ - Process accounts and perform the specified action. - - Args: - action: Action to perform (stop or start) - regions: List of regions to scan - exclude_accounts: List of account IDs to exclude - exclude_resources: List of resource types to exclude - exclude_regions: List of regions to exclude - dry_run: Whether to perform a dry run - stats: Statistics dictionary to update - - Returns: - Updated statistics dictionary - """ + """Process accounts and perform the specified action.""" try: - if exclude_accounts is None: - exclude_accounts = [] - if exclude_resources is None: - exclude_resources = [] - if exclude_regions is None: - exclude_regions = [] - # Initialize stats if not provided - if stats is None: - stats = initialize_stats() + # Initialize parameters + exclude_accounts = exclude_accounts or [] + exclude_resources = exclude_resources or [] + exclude_regions = exclude_regions or [] + stats = stats or initialize_stats() # Get account list - either from profiles or Organizations API - if self.use_profiles: - logger.info("Using AWS SSO profiles for account discovery") - accounts = get_accounts_from_profiles() - else: - logger.info("Using AWS Organizations API for account discovery") - accounts = get_account_list() + accounts = self._get_accounts() # Log summary of accounts logger.info(f"Found {len(accounts)} accounts to process") if exclude_accounts: logger.info(f"Excluding {len(exclude_accounts)} accounts") - # Process each account - for account in accounts: - account_name = account.get("account_name", "Unknown") - account_id = account.get("account_id") - # Skip excluded accounts - if account_id in exclude_accounts: - logger.info( - f"Skipping excluded account {account_id} ({account_name})" - ) - stats["accounts_skipped"] += 1 - continue - - try: - # Get credentials for the account - credentials = get_credentials(account_id, self.profile_name) - if not credentials: - logger.warning( - f"Could not obtain credentials for account {account_id}" - ) + # Process accounts with limited concurrency + with concurrent.futures.ThreadPoolExecutor( + max_workers=self.MAX_CONCURRENT_ACCOUNTS + ) as executor: + # Create futures for each account + futures = [] + for account in accounts: + account_id = account.get("account_id") + if account_id in exclude_accounts: stats["accounts_skipped"] += 1 continue - # Determine regions to scan for this account - account_regions = regions - if self.discover_regions: - try: - logger.info( - f"Discovering enabled regions for account {account_id}" - ) - account_regions = get_enabled_regions( - credentials, self.partition - ) - if exclude_regions: - account_regions = [ - r - for r in account_regions - if r not in exclude_regions - ] - logger.info( - f"Found {len(account_regions)} enabled regions in account {account_id}" - ) - # Cache discovered regions - self.region_cache[account_id] = account_regions - except Exception as e: - logger.error( - f"Error discovering regions for account {account_id}: {str(e)}" - ) - if regions: - logger.info( - f"Falling back to provided regions: {', '.join(regions)}" - ) - account_regions = regions - else: - logger.info( - "No regions provided and region discovery failed. Skipping account." - ) - stats["accounts_skipped"] += 1 - continue - elif exclude_regions: - account_regions = [ - r for r in account_regions if r not in exclude_regions - ] - - # Track regions processed for this account - stats["regions_processed"] += len(account_regions) - stats["regions_by_account"][account_id] = account_regions - - # Process this account with its regions - logger.info( - f"Processing account: {account_name} ({account_id}) in {len(account_regions)} regions" - ) - - self._process_single_account( + future = executor.submit( + self._process_account_safely, account=account, - credentials=credentials, - regions=account_regions, + regions=regions, action=action, exclude_resources=exclude_resources, + exclude_regions=exclude_regions, dry_run=dry_run, - stats=stats, - ) - stats["accounts_processed"] += 1 - - except Exception as e: - logger.error( - f"Error processing account {account_id}: {str(e)}", - exc_info=True, ) - stats["accounts_skipped"] += 1 + futures.append((future, account_id)) + + # Process results as they complete + for future, account_id in futures: + try: + account_stats = future.result() + # Merge the account's stats into the main stats + if account_stats: + self._merge_stats(stats, account_stats) + stats["accounts_processed"] += 1 + else: + stats["accounts_skipped"] += 1 + except Exception as e: + logger.error( + f"Error processing account {account_id}: {str(e)}", + exc_info=True, + ) + stats["accounts_skipped"] += 1 + stats["errors"].append( + f"Error processing account {account_id}: {str(e)}" + ) - # Log summary of regions processed - if self.discover_regions: - logger.info( - f"Total accounts processed: {stats.get('accounts_processed', 0)}" - ) - logger.info( - f"Total regions processed: {stats.get('regions_processed', 0)}" - ) - for account_id, account_regions in stats.get( - "regions_by_account", {} - ).items(): - logger.debug( - f"Account {account_id} regions: {', '.join(account_regions)}" - ) - self._print_resource_summary(stats, action) + # Log summary + self._log_summary(stats, action) return stats except KeyboardInterrupt: logger.warning("Account processing interrupted by user") raise - def _print_resource_summary(self, stats: Dict[str, Any], action: str) -> None: - """Print a summary of the resources processed.""" - print_resource_summary(stats, action, stats.get("dry_run", False)) - - def _process_single_account( + def _get_accounts(self) -> List[Dict[str, str]]: + """Get list of accounts from profiles or Organizations API.""" + if self.use_profiles: + logger.info("Using AWS SSO profiles for account discovery") + return ( + get_organization_accounts() + ) # Changed from get_accounts_from_profiles + else: + logger.info("Using AWS Organizations API for account discovery") + return get_account_list() + + def _process_account_safely( self, account: Dict[str, str], - credentials: Dict[str, str], regions: List[str], action: str, - dry_run: bool, exclude_resources: List[str], - stats: Dict[str, Any], - ) -> None: - """Process a single account for the specified action.""" - # Fix keys - use consistent naming with process_accounts + exclude_regions: List[str], + dry_run: bool, + ) -> Optional[Dict[str, Any]]: + """Process a single account with error handling.""" account_id = account.get("account_id") - account_name = account.get("account_name", account_id) + account_name = account.get("account_name", "Unknown") try: - # Auto-detect partition if not specified - if not self.partition: - self.partition = detect_partition_from_credentials(credentials) - logger.info(f"Auto-detected AWS partition: {self.partition}") - - # Ensure we're using regions from the right partition - if regions and regions[0] != "auto": - # Filter regions to match the detected partition - filtered_regions = [] - for region in regions: - if self.partition == "aws-us-gov" and not region.startswith( - "us-gov-" - ): - logger.debug( - f"Skipping region {region} - not in {self.partition} partition" - ) - continue - elif self.partition == "aws-cn" and not region.startswith("cn-"): - logger.debug( - f"Skipping region {region} - not in {self.partition} partition" - ) - continue - elif self.partition == "aws" and ( - region.startswith("us-gov-") or region.startswith("cn-") - ): - logger.debug( - f"Skipping region {region} - not in {self.partition} partition" - ) - continue - filtered_regions.append(region) - - # If filtering removed all regions, use defaults for the partition - if not filtered_regions: - if self.partition == "aws-us-gov": - filtered_regions = ["us-gov-east-1", "us-gov-west-1"] - elif self.partition == "aws-cn": - filtered_regions = ["cn-north-1", "cn-northwest-1"] - else: - filtered_regions = [ - "us-east-1", - "us-east-2", - "us-west-1", - "us-west-2", - ] - logger.info( - f"Using default regions for partition {self.partition}: {', '.join(filtered_regions)}" - ) + # Get credentials + credentials = get_credentials(account_id, self.profile_name) + if not credentials: + logger.warning(f"Could not obtain credentials for account {account_id}") + return None + + # Get regions for this account + account_regions = self._get_account_regions( + account_id, credentials, regions, exclude_regions + ) + if not account_regions: + logger.info(f"No valid regions for account {account_id}, skipping") + return None - regions = filtered_regions - elif not regions: - # If no regions provided, use defaults for the detected partition - if self.partition == "aws-us-gov": - regions = ["us-gov-east-1", "us-gov-west-1"] - elif self.partition == "aws-cn": - regions = ["cn-north-1", "cn-northwest-1"] - else: - regions = ["us-east-1", "us-east-2", "us-west-1", "us-west-2"] - logger.info( - f"No regions specified, using defaults for partition {self.partition}: {', '.join(regions)}" - ) + # Initialize account-specific stats + account_stats = initialize_stats() + account_stats["regions_processed"] = len(account_regions) + account_stats["regions_by_account"][account_id] = account_regions + # Process the account with its determined regions logger.info( - f"Processing account: {account_name} ({account_id}) in {len(regions)} regions" + f"Processing account: {account_name} ({account_id}) in {len(account_regions)} regions" ) - # Define special handling for EMR clusters - remove 'STOPPED' from valid states as it's not accepted by the API - emr_states = None - if "emr" not in exclude_resources: - emr_states = [ - "STARTING", - "BOOTSTRAPPING", - "RUNNING", - "TERMINATING", - "WAITING", - ] - logger.debug( - f"Using valid EMR states for discovery: {', '.join(emr_states)}" - ) + # Auto-detect partition if not specified + partition = ( + detect_partition_from_credentials(credentials) + if not self.partition + else self.partition + ) # Get resources with special handling for EMR - resources = get_account_resources( - credentials, - regions, - exclude_resources, - account_id, - account_name, - emr_cluster_states=emr_states, - ) - if not resources: - logger.error(f"Failed to get resources for account {account_id}") - return - - # Initialize stat counters if they don't exist - for stat_type in [ - "ec2_found", - "ec2_skipped", - "ec2_stopped", - "ec2_started", - "ec2_errors", - "rds_found", - "rds_skipped", - "rds_stopped", - "rds_started", - "rds_errors", - "eks_found", - "eks_skipped", - "eks_stopped", - "eks_started", - "eks_errors", - "emr_found", - "emr_skipped", - "emr_stopped", - "emr_started", - "emr_errors", - "ecr_images_found", - "ecr_old_images_found", - ]: - if stat_type not in stats: - stats[stat_type] = 0 - - # Initialize RDS engines stats if it doesn't exist - if "rds_engines" not in stats: - stats["rds_engines"] = {} - - # Initialize errors array if it doesn't exist - if "errors" not in stats: - stats["errors"] = [] - - # Normalize state and status fields for all resource types - ec2_instances = resources.get("ec2_instances", []) - for instance in ec2_instances: - if "state" not in instance and "status" in instance: - instance["state"] = instance["status"] - elif "status" not in instance and "state" in instance: - instance["status"] = instance["state"] - elif "state" not in instance and "status" not in instance: - instance["status"] = "unknown" - instance["state"] = "unknown" - - rds_instances = resources.get("rds_instances", []) - for instance in rds_instances: - if "state" not in instance and "status" in instance: - instance["state"] = instance["status"] - elif "status" not in instance and "state" in instance: - instance["status"] = instance["state"] - elif "state" not in instance and "status" not in instance: - instance["status"] = "unknown" - instance["state"] = "unknown" - - eks_clusters = resources.get("eks_clusters", []) - for cluster in eks_clusters: - if "state" not in cluster and "status" in cluster: - cluster["state"] = cluster["status"] - elif "status" not in cluster and "state" in cluster: - cluster["status"] = cluster["state"] - elif "state" not in cluster and "status" not in cluster: - cluster["status"] = "unknown" - cluster["state"] = "unknown" - - emr_clusters = resources.get("emr_clusters", []) - for cluster in emr_clusters: - if "state" not in cluster: - cluster["state"] = cluster.get("status", "UNKNOWN") - if "status" not in cluster: - cluster["status"] = cluster.get("state", "UNKNOWN") - - # Count service-managed EC2 instances - eks_tag = self.config.get("eks_tag", "kubernetes.io/cluster/") - exclusion_tag = self.config.get("exclusion_tag", "DoNotStop") - emr_tag = self.config.get("emr_tag", "aws:elasticmapreduce:job-flow-id") - - eks_managed = sum( - 1 - for i in resources.get("ec2_instances", []) - if any(tag.startswith(eks_tag) for tag in i.get("tags", {})) - ) - emr_managed = sum( - 1 - for i in resources.get("ec2_instances", []) - if emr_tag in i.get("tags", {}) - ) + emr_states = self._get_emr_states(exclude_resources) + + try: + # Get all resources - this might raise AuthFailure + resources = get_account_resources( + credentials, + account_regions, + exclude_resources, + account_id, + account_name, + emr_cluster_states=emr_states, + ) - # Update resource counts - total_ec2 = len(resources.get("ec2_instances", [])) - excluded_ec2 = sum( - 1 - for i in resources.get("ec2_instances", []) - if exclusion_tag in i.get("tags", {}) - ) - stats["ec2_found"] += total_ec2 - stats["ec2_skipped"] += excluded_ec2 - - stats["rds_found"] += len(resources.get("rds_instances", [])) - stats["emr_found"] += len(resources.get("emr_clusters", [])) - stats["eks_found"] += len(resources.get("eks_clusters", [])) - stats["ecr_images_found"] += len(resources.get("ecr_images", [])) - stats["ecr_old_images_found"] += len(resources.get("ecr_old_images", [])) - - # Track RDS engines - for instance in resources.get("rds_instances", []): - if instance and "engine" in instance: - engine = instance["engine"] - if engine in stats["rds_engines"]: - stats["rds_engines"][engine] += 1 - else: - stats["rds_engines"][engine] = 1 - - # Log discovered resources with service-managed counts - logger.info( - f"Account {account_id} - Found {total_ec2} EC2 instances ({excluded_ec2} explicitly excluded, {eks_managed} EKS-managed, {emr_managed} EMR-managed)" - ) + if not resources: + logger.error(f"Failed to get resources for account {account_id}") + return None + + # Normalize states and collect statistics + self._process_resources(resources, account_stats) + + # Log resource counts + self._log_resource_counts(account_id, resources, account_stats) + + # Create resource managers only once per account + managers = self._create_resource_managers( + credentials, account_id, account_name + ) + + # Process actions on resources based on the selected mode + self._execute_actions( + action, + managers, + resources, + exclude_resources, + dry_run, + account_stats, + ) + + return account_stats + except ClientError as e: + # Special handling for authentication failures + error_code = e.response.get("Error", {}).get("Code", "") + if error_code in [ + "AuthFailure", + "InvalidClientTokenId", + "UnauthorizedOperation", + "AccessDenied", + ]: + logger.error( + f"Authentication failure for account {account_id}: {e}" + ) + account_stats["errors"].append( + f"Authentication failure for account {account_id}: {str(e)}" + ) + return account_stats + raise # Re-raise other client errors + + except Exception as e: + logger.error(f"Error processing account {account_id}: {str(e)}") + if account_stats := locals().get("account_stats"): + account_stats["errors"].append( + f"Error processing account {account_id}: {str(e)}" + ) + return account_stats + return None + + def _get_account_regions( + self, + account_id: str, + credentials: Dict[str, str], + provided_regions: List[str], + exclude_regions: List[str], + ) -> List[str]: + """Determine regions to use for an account.""" + if self.discover_regions: + try: + logger.info(f"Discovering enabled regions for account {account_id}") + discovered_regions = get_available_regions( + credentials, self.partition + ) # Changed from get_enabled_regions + account_regions = [ + r for r in discovered_regions if r not in exclude_regions + ] + logger.info( + f"Found {len(account_regions)} enabled regions in account {account_id}" + ) + self.region_cache[account_id] = account_regions + return account_regions + except Exception as e: + logger.error( + f"Error discovering regions for account {account_id}: {str(e)}" + ) + if provided_regions: + logger.info( + f"Falling back to provided regions: {', '.join(provided_regions)}" + ) + return provided_regions + return [] + elif provided_regions: + filtered_regions = [r for r in provided_regions if r not in exclude_regions] + if not filtered_regions: + filtered_regions = self._get_default_regions(credentials) + logger.info( + f"All provided regions were excluded, using defaults: {', '.join(filtered_regions)}" + ) + return filtered_regions + else: + default_regions = self._get_default_regions(credentials) logger.info( - f"Account {account_id} - Found {len(resources.get('rds_instances', []))} RDS instances" + f"No regions specified, using defaults: {', '.join(default_regions)}" ) - logger.info( - f"Account {account_id} - Found {len(resources.get('eks_clusters', []))} EKS clusters" + return default_regions + + def _get_default_regions( + self, credentials: Optional[Dict[str, str]] = None + ) -> List[str]: + """ + Get default regions for the current partition. + + Args: + credentials: Optional credentials to use for partition detection + + Returns: + List of default regions for the partition + """ + # If partition is not explicitly set, try to detect it from credentials + partition = self.partition + if not partition and credentials: + try: + partition = detect_partition_from_credentials(credentials) + logger.info(f"Auto-detected partition for default regions: {partition}") + except Exception as e: + logger.warning(f"Failed to auto-detect partition from credentials: {e}") + + # If still no partition, use commercial AWS as fallback + partition = partition or "aws" + + if partition == "aws-us-gov": + return ["us-gov-east-1", "us-gov-west-1"] + elif partition == "aws-cn": + return ["cn-north-1", "cn-northwest-1"] + else: + return ["us-east-1", "us-east-2", "us-west-1", "us-west-2"] + + def _get_emr_states(self, exclude_resources: List[str]) -> Optional[List[str]]: + """Get valid EMR states for discovery.""" + if "emr" not in exclude_resources: + emr_states = [ + "STARTING", + "BOOTSTRAPPING", + "RUNNING", + "WAITING", + "TERMINATING", + ] + logger.debug( + f"Using valid EMR states for discovery: {', '.join(emr_states)}" ) - logger.info( - f"Account {account_id} - Found {len(resources.get('emr_clusters', []))} EMR clusters" + return emr_states + return None + + def _process_resources( + self, resources: Dict[str, List[Dict[str, Any]]], stats: Dict[str, Any] + ) -> None: + """Normalize resource states and collect statistics.""" + # Normalize state and status fields for consistency + for resource_key, state_mappings in [ + ("ec2_instances", "state"), + ("rds_instances", "state"), + ("eks_clusters", "state"), + ("emr_clusters", "state"), + ]: + for resource in resources.get(resource_key, []): + self._normalize_resource_state(resource) + + # Count service-managed EC2 instances + eks_tag = self.config.get("eks_tag", "kubernetes.io/cluster/") + exclusion_tag = self.config.get("exclusion_tag", "DoNotStop") + emr_tag = self.config.get("emr_tag", "aws:elasticmapreduce:job-flow-id") + ec2_instances = resources.get("ec2_instances", []) + eks_managed = sum( + 1 + for i in ec2_instances + if any(tag.startswith(eks_tag) for tag in i.get("tags", {})) + ) + emr_managed = sum(1 for i in ec2_instances if emr_tag in i.get("tags", {})) + excluded_ec2 = sum( + 1 for i in ec2_instances if exclusion_tag in i.get("tags", {}) + ) + # Update resource counts + stats["ec2_found"] += len(ec2_instances) + stats["ec2_skipped"] += excluded_ec2 + stats["rds_found"] += len(resources.get("rds_instances", [])) + stats["emr_found"] += len(resources.get("emr_clusters", [])) + stats["eks_found"] += len(resources.get("eks_clusters", [])) + stats["ecr_images_found"] += len(resources.get("ecr_images", [])) + stats["ecr_old_images_found"] += len(resources.get("ecr_old_images", [])) + + # Track RDS engines + for instance in resources.get("rds_instances", []): + if instance and "engine" in instance: + engine = instance["engine"] + stats["rds_engines"][engine] = stats["rds_engines"].get(engine, 0) + 1 + + def _normalize_resource_state(self, resource: Dict[str, Any]) -> None: + """Ensure resources have both state and status fields.""" + if "state" not in resource and "status" in resource: + resource["state"] = resource["status"] + elif "status" not in resource and "state" in resource: + resource["status"] = resource["state"] + elif "state" not in resource and "status" not in resource: + resource["status"] = "unknown" + resource["state"] = "unknown" + + def _log_resource_counts( + self, + account_id: str, + resources: Dict[str, List[Dict[str, Any]]], + stats: Dict[str, Any], + ) -> None: + """Log discovered resource counts.""" + total_ec2 = len(resources.get("ec2_instances", [])) + excluded_ec2 = sum( + 1 + for i in resources.get("ec2_instances", []) + if self.config.get("exclusion_tag") in i.get("tags", {}) + ) + eks_tag = self.config.get("eks_tag", "kubernetes.io/cluster/") + emr_tag = self.config.get("emr_tag", "aws:elasticmapreduce:job-flow-id") + eks_managed = sum( + 1 + for i in resources.get("ec2_instances", []) + if any(tag.startswith(eks_tag) for tag in i.get("tags", {})) + ) + emr_managed = sum( + 1 + for i in resources.get("ec2_instances", []) + if emr_tag in i.get("tags", {}) + ) + logger.info( + f"Account {account_id} - Found {total_ec2} EC2 instances ({excluded_ec2} explicitly excluded, {eks_managed} EKS-managed, {emr_managed} EMR-managed)" + ) + logger.info( + f"Account {account_id} - Found {len(resources.get('rds_instances', []))} RDS instances" + ) + logger.info( + f"Account {account_id} - Found {len(resources.get('eks_clusters', []))} EKS clusters" + ) + logger.info( + f"Account {account_id} - Found {len(resources.get('emr_clusters', []))} EMR clusters" + ) + logger.info( + f"Account {account_id} - Found {len(resources.get('ecr_images', []))} ECR images " + f"({len(resources.get('ecr_old_images', []))} older than 1 year)" + ) + + def _create_resource_managers( + self, credentials: Dict[str, str], account_id: str, account_name: str + ) -> Dict[str, Any]: + """Create all resource managers for an account.""" + return { + "ec2": EC2Manager(credentials, account_id, account_name), + "rds": RDSManager(credentials, account_id, account_name), + "eks": EKSManager(credentials, account_id, account_name), + "emr": EMRManager(credentials, account_id, account_name), + } + + def _execute_actions( + self, + action: str, + managers: Dict[str, Any], + resources: Dict[str, List[Dict[str, Any]]], + exclude_resources: List[str], + dry_run: bool, + stats: Dict[str, Any], + ) -> None: + """Execute start/stop actions on resources.""" + + # Helper function to safely process manager results + def safe_get_result(result: Optional[Dict[str, int]], key: str) -> int: + if result is None: + return 0 + return result.get(key, 0) + + # Execute actions on each resource type + for resource_type in self.RESOURCE_TYPES: + if resource_type == "ecr": # ECR doesn't have start/stop actions + continue + + if resource_type in exclude_resources: + # Skip this resource type + resource_list = resources.get( + ( + f"{resource_type}_instances" + if resource_type != "eks" + else "eks_clusters" + ), + [], + ) + stats[f"{resource_type}_skipped"] += len(resource_list) + continue + + # Get the appropriate manager and resource list + manager = managers.get(resource_type) + if not manager: + logger.warning(f"No manager found for {resource_type}, skipping") + continue + + # Get the right resource list name + resource_key = ( + f"{resource_type}_instances" + if resource_type != "eks" + else "eks_clusters" ) - logger.info( - f"Account {account_id} - Found {len(resources.get('ecr_images', []))} ECR images " - f"({len(resources.get('ecr_old_images', []))} older than 1 year)" + resource_list = resources.get(resource_key, []) + + # Execute the action + if action == "stop": + result = manager.stop(resource_list, dry_run) + stats[f"{resource_type}_stopped"] += safe_get_result(result, "success") + stats[f"{resource_type}_errors"] += safe_get_result(result, "errors") + stats[f"{resource_type}_skipped"] += safe_get_result(result, "skipped") + else: # action == "start" + result = manager.start(resource_list, dry_run) + stats[f"{resource_type}_started"] += safe_get_result(result, "success") + stats[f"{resource_type}_errors"] += safe_get_result(result, "errors") + stats[f"{resource_type}_skipped"] += safe_get_result(result, "skipped") + + def _merge_stats( + self, main_stats: Dict[str, Any], account_stats: Dict[str, Any] + ) -> None: + """Merge account-specific stats into main stats.""" + # Merge numeric values + for key, value in account_stats.items(): + if isinstance(value, (int, float)) and key in main_stats: + main_stats[key] += value + + # Merge errors list + main_stats["errors"].extend(account_stats.get("errors", [])) + + # Merge RDS engines + for engine, count in account_stats.get("rds_engines", {}).items(): + main_stats["rds_engines"][engine] = ( + main_stats["rds_engines"].get(engine, 0) + count ) - # Initialize resource managers with account name - ec2_manager = EC2Manager(credentials, account_id, account_name) - rds_manager = RDSManager(credentials, account_id, account_name) - eks_manager = EKSManager(credentials, account_id, account_name) - emr_manager = EMRManager(credentials, account_id, account_name) + # Merge regions_by_account + for account_id, regions in account_stats.get("regions_by_account", {}).items(): + main_stats["regions_by_account"][account_id] = regions - # Helper function to safely process manager results - def safe_get_result(result: Optional[Dict[str, int]], key: str) -> int: - if result is None: - return 0 - return result.get(key, 0) + def _log_summary(self, stats: Dict[str, Any], action: str) -> None: + """Log summary information after processing.""" + if self.discover_regions: + logger.info( + f"Total accounts processed: {stats.get('accounts_processed', 0)}" + ) + logger.info(f"Total regions processed: {stats.get('regions_processed', 0)}") + for account_id, account_regions in stats.get( + "regions_by_account", {} + ).items(): + logger.debug( + f"Account {account_id} regions: {', '.join(account_regions)}" + ) - # Perform actions based on selected mode - if action == "stop": - if "ec2" not in exclude_resources: - result = ec2_manager.stop( - resources.get("ec2_instances", []), dry_run - ) - stats["ec2_stopped"] += safe_get_result(result, "success") - stats["ec2_errors"] += safe_get_result(result, "errors") - else: - stats["ec2_skipped"] += total_ec2 - excluded_ec2 - - if "rds" not in exclude_resources: - result = rds_manager.stop( - resources.get("rds_instances", []), dry_run - ) - stats["rds_stopped"] += safe_get_result(result, "success") - stats["rds_errors"] += safe_get_result(result, "errors") - else: - stats["rds_skipped"] += len(resources.get("rds_instances", [])) - - if "eks" not in exclude_resources: - result = eks_manager.stop( - resources.get("eks_clusters", []), dry_run - ) - stats["eks_stopped"] += safe_get_result(result, "success") - stats["eks_errors"] += safe_get_result(result, "errors") - else: - stats["eks_skipped"] += len(resources.get("eks_clusters", [])) - - if "emr" not in exclude_resources: - result = emr_manager.stop( - resources.get("emr_clusters", []), dry_run - ) - stats["emr_stopped"] += safe_get_result(result, "success") - stats["emr_errors"] += safe_get_result(result, "errors") - else: - stats["emr_skipped"] += len(resources.get("emr_clusters", [])) - else: # action == "start" - if "ec2" not in exclude_resources: - result = ec2_manager.start( - resources.get("ec2_instances", []), dry_run - ) - stats["ec2_started"] += safe_get_result(result, "success") - stats["ec2_errors"] += safe_get_result(result, "errors") - else: - stats["ec2_skipped"] += total_ec2 - excluded_ec2 - - if "rds" not in exclude_resources: - result = rds_manager.start( - resources.get("rds_instances", []), dry_run - ) - stats["rds_started"] += safe_get_result(result, "success") - stats["rds_errors"] += safe_get_result(result, "errors") - else: - stats["rds_skipped"] += len(resources.get("rds_instances", [])) - - if "eks" not in exclude_resources: - result = eks_manager.start( - resources.get("eks_clusters", []), dry_run - ) - stats["eks_started"] += safe_get_result(result, "success") - stats["eks_errors"] += safe_get_result(result, "errors") - else: - stats["eks_skipped"] += len(resources.get("eks_clusters", [])) - - if "emr" not in exclude_resources: - result = emr_manager.start( - resources.get("emr_clusters", []), dry_run - ) - stats["emr_started"] += safe_get_result(result, "success") - stats["emr_errors"] += safe_get_result(result, "errors") - else: - stats["emr_skipped"] += len(resources.get("emr_clusters", [])) - except Exception as e: - logger.error(f"Error processing account {account_id}: {str(e)}") - if "errors" not in stats: - stats["errors"] = [] - stats["errors"].append(f"Error processing account {account_id}: {str(e)}") - return + print_resource_summary(stats, action, stats.get("dry_run", False)) diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/discovery.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/discovery.py index e7c1c183..175f37f5 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/discovery.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/discovery.py @@ -17,8 +17,8 @@ from aws_resource_management.config_manager import get_config from aws_resource_management.logging_setup import setup_logging from aws_resource_management.managers import ( - ECRManager, EC2Manager, + ECRManager, EKSManager, EMRManager, RDSManager, @@ -518,31 +518,43 @@ def get_account_resources( # Discover EC2 instances if "ec2" not in exclude_resources: - logger.info(f"Discovering EC2 instances for account {account_id} in {len(regions)} regions") + logger.info( + f"Discovering EC2 instances for account {account_id} in {len(regions)} regions" + ) ec2_manager = EC2Manager(credentials, account_id, account_name) resources["ec2_instances"] = ec2_manager.discover_instances(regions) # Discover RDS instances if "rds" not in exclude_resources: - logger.info(f"Discovering RDS instances for account {account_id} in {len(regions)} regions") + logger.info( + f"Discovering RDS instances for account {account_id} in {len(regions)} regions" + ) rds_manager = RDSManager(credentials, account_id, account_name) resources["rds_instances"] = rds_manager.discover_instances(regions) # Discover EKS clusters if "eks" not in exclude_resources: - logger.info(f"Discovering EKS clusters for account {account_id} in {len(regions)} regions") + logger.info( + f"Discovering EKS clusters for account {account_id} in {len(regions)} regions" + ) eks_manager = EKSManager(credentials, account_id, account_name) resources["eks_clusters"] = eks_manager.discover_clusters(regions) # Discover EMR clusters if "emr" not in exclude_resources: - logger.info(f"Discovering EMR clusters for account {account_id} in {len(regions)} regions") + logger.info( + f"Discovering EMR clusters for account {account_id} in {len(regions)} regions" + ) emr_manager = EMRManager(credentials, account_id, account_name) - resources["emr_clusters"] = emr_manager.discover_clusters(regions, emr_cluster_states) + resources["emr_clusters"] = emr_manager.discover_clusters( + regions, emr_cluster_states + ) # Discover ECR images if "ecr" not in exclude_resources: - logger.info(f"Discovering ECR images for account {account_id} in {len(regions)} regions") + logger.info( + f"Discovering ECR images for account {account_id} in {len(regions)} regions" + ) ecr_manager = ECRManager(credentials, account_id, account_name) result = ecr_manager.discover_images_by_age(regions) resources["ecr_images"] = result["all_images"] diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/__init__.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/__init__.py index 0cbfac73..0583be56 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/__init__.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/__init__.py @@ -3,8 +3,8 @@ """ from aws_resource_management.managers.base import ResourceManager -from aws_resource_management.managers.ecr import ECRManager from aws_resource_management.managers.ec2 import EC2Manager +from aws_resource_management.managers.ecr import ECRManager from aws_resource_management.managers.eks import EKSManager from aws_resource_management.managers.emr import EMRManager from aws_resource_management.managers.rds import RDSManager diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/base.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/base.py index 9a85ec99..9e626b2e 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/base.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/base.py @@ -3,50 +3,48 @@ """ import logging +import sys import time from datetime import datetime from typing import Any, Dict, List, Optional, Union import boto3 from aws_resource_management.config_manager import get_config +from botocore.exceptions import ClientError logger = logging.getLogger(__name__) config = get_config() +# Define authentication failure error codes +AUTH_FAILURE_ERRORS = [ + "AuthFailure", + "InvalidClientTokenId", + "UnauthorizedOperation", + "AccessDenied", +] + + class ResourceManager: """Base class for AWS resource managers.""" def __init__( - self, credentials: Dict[str, str], account_id: str, account_name: Optional[str] = None + self, + credentials: Dict[str, str], + account_id: str, + account_name: Optional[str] = None, ): - """ - Initialize the resource manager. - - Args: - credentials: AWS credentials dictionary - account_id: AWS account ID - account_name: AWS account name (optional) - """ + """Initialize the resource manager.""" self.credentials = credentials self.account_id = account_id self.account_name = account_name self.clients = {} # Cache for boto3 clients - - def get_boto3_client(self, service_name: str, region: str) -> Any: - """ - Get a boto3 client for the specified service and region. - - Args: - service_name: AWS service name - region: AWS region - Returns: - Boto3 client object or None if creation fails - """ + def get_boto3_client(self, service_name: str, region: str) -> Any: + """Get a boto3 client for the specified service and region.""" client_key = f"{service_name}-{region}" if client_key in self.clients: return self.clients[client_key] - + try: client = boto3.client( service_name, @@ -55,15 +53,74 @@ def get_boto3_client(self, service_name: str, region: str) -> Any: aws_secret_access_key=self.credentials.get("aws_secret_access_key"), aws_session_token=self.credentials.get("aws_session_token"), ) + + # Immediately validate credentials by making a simple API call + self._validate_credentials(client, service_name, region) + + # Cache client if validation passes self.clients[client_key] = client return client + except ClientError as e: + # Check if it's an authentication error + error_code = e.response.get("Error", {}).get("Code", "") + error_message = e.response.get("Error", {}).get("Message", "") + + if error_code in AUTH_FAILURE_ERRORS: + logger.error( + f"Authentication failed for {service_name} in {region}: {error_message}" + ) + logger.error( + f"This account likely doesn't have valid credentials. Skipping account {self.account_id}." + ) + # Return None to signal authentication failure + return None + else: + logger.error( + f"Error creating {service_name} client in {region}: {error_code} - {error_message}" + ) + return None except Exception as e: logger.error(f"Error creating {service_name} client in {region}: {e}") return None - + + def _validate_credentials( + self, client: Any, service_name: str, region: str + ) -> None: + """ + Validate credentials by making a minimal API call. + Raises ClientError if authentication fails. + """ + # Make service-specific minimal API call to verify credentials + try: + if service_name == "ec2": + client.describe_regions(MaxResults=5) + elif service_name == "rds": + client.describe_db_engine_versions(MaxRecords=1) + elif service_name == "eks": + client.list_clusters(maxResults=1) + elif service_name == "emr": + client.list_clusters(MaxResults=1) + elif service_name == "ecr": + client.describe_repositories(maxResults=1) + else: + # For other services, don't validate (better than failing) + pass + except ClientError as e: + # Check if it's an authentication error + error_code = e.response.get("Error", {}).get("Code", "") + if error_code in AUTH_FAILURE_ERRORS: + logger.error( + f"Authentication validation failed for {service_name} in {region}" + ) + raise # Re-raise to be caught by get_boto3_client + # Otherwise ignore other errors during validation + except Exception: + # Ignore any other exceptions during validation + pass + # Alias for backwards compatibility create_client = get_boto3_client - + def should_exclude(self, resource: Dict[str, Any]) -> bool: """ Check if a resource should be excluded based on tags. @@ -76,9 +133,9 @@ def should_exclude(self, resource: Dict[str, Any]) -> bool: """ tags = resource.get("tags", {}) exclusion_tag = config.get("exclusion_tag") - + return exclusion_tag in tags - + def get_timestamp(self) -> str: """ Get current timestamp in ISO 8601 format. @@ -87,7 +144,7 @@ def get_timestamp(self) -> str: ISO 8601 timestamp string """ return datetime.utcnow().isoformat() - + def get_partition_for_region(self, region: str) -> str: """ Get AWS partition for a region. @@ -104,16 +161,16 @@ def get_partition_for_region(self, region: str) -> str: return "aws-cn" else: return "aws" - + def log_action( - self, - resource_id: str, - region: str, - action: str, + self, + resource_id: str, + region: str, + action: str, resource_name: Optional[str] = None, - details: Optional[str] = None, - dry_run: bool = False, - existing_schedule: Optional[str] = None + details: Optional[str] = None, + dry_run: bool = False, + existing_schedule: Optional[str] = None, ) -> None: """ Log an action taken on a resource. @@ -132,28 +189,24 @@ def log_action( details_str = f": {details}" if details else "" schedule_str = f" [Schedule: {existing_schedule}]" if existing_schedule else "" dry_run_str = "[DRY RUN] " if dry_run else "" - + logger.info( f"{dry_run_str}{action_str} {self.resource_type} {resource_id}{name_str} " f"in {region}{details_str}{schedule_str}" ) def paginate_boto3(self, client, method_name, result_key, **kwargs): - """ - Helper method to handle pagination for boto3 API calls. - - Args: - client: Boto3 client - method_name: API method name to call - result_key: Key in the response that contains the results - **kwargs: Additional arguments to pass to the method - - Returns: - Combined list of all items across pages - """ + """Helper method to handle pagination for boto3 API calls.""" + # If client is None (authentication failed earlier), return empty list + if client is None: + logger.warning( + f"Skipping {method_name} due to previous authentication failure" + ) + return [] + method = getattr(client, method_name) items = [] - + try: # Try to use paginator if available try: @@ -162,43 +215,66 @@ def paginate_boto3(self, client, method_name, result_key, **kwargs): if result_key in page: items.extend(page[result_key]) return items - except (AttributeError, client.exceptions.ClientError): - # Fall back to manual pagination if paginator not available + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "") + if error_code in AUTH_FAILURE_ERRORS: + logger.error( + f"Authentication failure in paginator for {method_name}: {error_code}" + ) + return [] # Return empty result on auth failure + # Fall back to manual pagination + pass + except AttributeError: + # No paginator available, fall back to manual pagination pass - + # Manual pagination - response = method(**kwargs) - if result_key in response: - items.extend(response[result_key]) - - # Check for different pagination keys - pagination_keys = [ - "NextToken", "nextToken", "Marker", "marker", - "next_token", "next_marker" - ] - - pagination_key = next( - (k for k in pagination_keys if k in response), None - ) - - while pagination_key and response.get(pagination_key): - kwargs[pagination_key] = response[pagination_key] + try: response = method(**kwargs) if result_key in response: items.extend(response[result_key]) - - return items + + pagination_keys = [ + "NextToken", + "nextToken", + "Marker", + "marker", + "next_token", + "next_marker", + ] + pagination_key = next( + (k for k in pagination_keys if k in response), None + ) + + while pagination_key and response.get(pagination_key): + kwargs[pagination_key] = response[pagination_key] + response = method(**kwargs) + if result_key in response: + items.extend(response[result_key]) + + except ClientError as e: + error_code = e.response.get("Error", {}).get("Code", "") + if error_code in AUTH_FAILURE_ERRORS: + logger.error( + f"Authentication failure in manual pagination for {method_name}: {error_code}" + ) + return [] # Return empty result on auth failure + logger.error(f"Error during manual pagination for {method_name}: {e}") + except Exception as e: - logger.error( - f"Error paginating through {method_name} for {client._service_model.service_name}: {str(e)}" - ) - return items + logger.error(f"Error paginating through {method_name}: {str(e)}") + + return items # Abstract methods that should be implemented by subclasses - def start(self, resources: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: + def start( + self, resources: List[Dict[str, Any]], dry_run: bool = False + ) -> Dict[str, int]: """Start resources - abstract method to be implemented by subclasses.""" raise NotImplementedError("Subclasses must implement start()") - - def stop(self, resources: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: + + def stop( + self, resources: List[Dict[str, Any]], dry_run: bool = False + ) -> Dict[str, int]: """Stop resources - abstract method to be implemented by subclasses.""" raise NotImplementedError("Subclasses must implement stop()") diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ec2.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ec2.py index 4368975c..8fe39c37 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ec2.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ec2.py @@ -2,8 +2,8 @@ EC2 resource manager class. """ -from typing import Any, Dict, List, Optional from collections import defaultdict +from typing import Any, Dict, List, Optional from aws_resource_management.config_manager import get_config from aws_resource_management.logging_setup import setup_logging @@ -33,27 +33,31 @@ def should_exclude(self, instance: Dict[str, Any]) -> bool: exclusion_tag = config.get("exclusion_tag") eks_tag = config.get("eks_tag") emr_tag = config.get("emr_tag") - + if exclusion_tag in tags: - logger.debug(f"Skipping EC2 instance {instance['id']} with exclusion tag {exclusion_tag}") + logger.debug( + f"Skipping EC2 instance {instance['id']} with exclusion tag {exclusion_tag}" + ) return True - + if any(tag.startswith(eks_tag) for tag in tags): logger.debug(f"Skipping EC2 instance {instance['id']} - EKS managed node") return True - + if emr_tag in tags: logger.debug(f"Skipping EC2 instance {instance['id']} - EMR managed node") return True return False - def stop(self, instances: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: + def stop( + self, instances: List[Dict[str, Any]], dry_run: bool = False + ) -> Dict[str, int]: """Stop running EC2 instances in batches for efficiency.""" # Group instances by region for batch processing running_instances_by_region = defaultdict(list) stats = {"success": 0, "errors": 0, "skipped": 0} - + # Filter instances - keep only running instances that should not be excluded for instance in instances: if self.should_exclude(instance): @@ -63,11 +67,13 @@ def stop(self, instances: List[Dict[str, Any]], dry_run: bool = False) -> Dict[s if instance["state"] == "running": running_instances_by_region[instance["region"]].append(instance) else: - logger.debug(f"Skipping EC2 instance {instance['id']} in state {instance['state']} (not running)") + logger.debug( + f"Skipping EC2 instance {instance['id']} in state {instance['state']} (not running)" + ) stats["skipped"] += 1 - + timestamp = self.get_timestamp() - + # Process each region's instances in batches for region, region_instances in running_instances_by_region.items(): ec2_client = self.get_boto3_client("ec2", region) @@ -75,52 +81,62 @@ def stop(self, instances: List[Dict[str, Any]], dry_run: bool = False) -> Dict[s logger.error(f"Failed to create EC2 client for region {region}") stats["errors"] += len(region_instances) continue - + # Process instances in batches for i in range(0, len(region_instances), self.MAX_INSTANCES_PER_API_CALL): - batch = region_instances[i:i + self.MAX_INSTANCES_PER_API_CALL] + batch = region_instances[i : i + self.MAX_INSTANCES_PER_API_CALL] instance_ids = [instance["id"] for instance in batch] - + try: # Tag all instances in batch if not dry_run: - logger.info(f"Tagging {len(instance_ids)} EC2 instances in {region}") + logger.info( + f"Tagging {len(instance_ids)} EC2 instances in {region}" + ) ec2_client.create_tags( Resources=instance_ids, - Tags=[{"Key": config.get("stop_tag"), "Value": timestamp}] + Tags=[{"Key": config.get("stop_tag"), "Value": timestamp}], ) - + # Stop all instances in batch if not dry_run: - logger.info(f"Stopping {len(instance_ids)} EC2 instances in {region}") + logger.info( + f"Stopping {len(instance_ids)} EC2 instances in {region}" + ) ec2_client.stop_instances(InstanceIds=instance_ids) else: - logger.info(f"[DRY RUN] Would stop {len(instance_ids)} EC2 instances in {region}") - + logger.info( + f"[DRY RUN] Would stop {len(instance_ids)} EC2 instances in {region}" + ) + # Log individual instances for tracking purposes for instance in batch: self.log_action( - instance["id"], - region, - "stop", + instance["id"], + region, + "stop", resource_name=instance.get("name", "Unnamed"), - dry_run=dry_run + dry_run=dry_run, ) - + stats["success"] += len(batch) except Exception as e: logger.error(f"Batch EC2 stop operation failed in {region}: {e}") stats["errors"] += len(batch) - logger.info(f"EC2 stop summary: {stats['success']} stopped, {stats['skipped']} skipped, {stats['errors']} errors") + logger.info( + f"EC2 stop summary: {stats['success']} stopped, {stats['skipped']} skipped, {stats['errors']} errors" + ) return stats - def start(self, instances: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: + def start( + self, instances: List[Dict[str, Any]], dry_run: bool = False + ) -> Dict[str, int]: """Start stopped EC2 instances in batches for efficiency.""" # Group instances by region for batch processing stopped_instances_by_region = defaultdict(list) stats = {"success": 0, "errors": 0, "skipped": 0} - + # Filter instances - keep only stopped instances that should not be excluded for instance in instances: if self.should_exclude(instance): @@ -130,9 +146,11 @@ def start(self, instances: List[Dict[str, Any]], dry_run: bool = False) -> Dict[ if instance["state"] == "stopped": stopped_instances_by_region[instance["region"]].append(instance) else: - logger.debug(f"Skipping EC2 instance {instance['id']} in state {instance['state']} (not stopped)") + logger.debug( + f"Skipping EC2 instance {instance['id']} in state {instance['state']} (not stopped)" + ) stats["skipped"] += 1 - + # Process each region's instances in batches for region, region_instances in stopped_instances_by_region.items(): ec2_client = self.get_boto3_client("ec2", region) @@ -140,63 +158,69 @@ def start(self, instances: List[Dict[str, Any]], dry_run: bool = False) -> Dict[ logger.error(f"Failed to create EC2 client for region {region}") stats["errors"] += len(region_instances) continue - + # Process instances in batches for i in range(0, len(region_instances), self.MAX_INSTANCES_PER_API_CALL): - batch = region_instances[i:i + self.MAX_INSTANCES_PER_API_CALL] + batch = region_instances[i : i + self.MAX_INSTANCES_PER_API_CALL] instance_ids = [instance["id"] for instance in batch] - + try: # Start all instances in batch if not dry_run: - logger.info(f"Starting {len(instance_ids)} EC2 instances in {region}") + logger.info( + f"Starting {len(instance_ids)} EC2 instances in {region}" + ) ec2_client.start_instances(InstanceIds=instance_ids) else: - logger.info(f"[DRY RUN] Would start {len(instance_ids)} EC2 instances in {region}") - + logger.info( + f"[DRY RUN] Would start {len(instance_ids)} EC2 instances in {region}" + ) + # Remove tags in batch if not dry_run: - logger.info(f"Removing stop tag from {len(instance_ids)} EC2 instances in {region}") + logger.info( + f"Removing stop tag from {len(instance_ids)} EC2 instances in {region}" + ) ec2_client.delete_tags( Resources=instance_ids, - Tags=[{"Key": config.get("stop_tag")}] + Tags=[{"Key": config.get("stop_tag")}], ) - + # Log individual instances for tracking purposes for instance in batch: self.log_action( - instance["id"], - region, - "start", + instance["id"], + region, + "start", resource_name=instance.get("name", "Unnamed"), - dry_run=dry_run + dry_run=dry_run, ) - + stats["success"] += len(batch) except Exception as e: logger.error(f"Batch EC2 start operation failed in {region}: {e}") stats["errors"] += len(batch) - logger.info(f"EC2 start summary: {stats['success']} started, {stats['skipped']} skipped, {stats['errors']} errors") + logger.info( + f"EC2 start summary: {stats['success']} started, {stats['skipped']} skipped, {stats['errors']} errors" + ) return stats def discover_instances(self, regions: List[str]) -> List[Dict[str, Any]]: """Discover EC2 instances across multiple regions.""" all_instances = [] - + for region in regions: try: - ec2_client = self.get_boto3_client('ec2', region) + ec2_client = self.get_boto3_client("ec2", region) if not ec2_client: continue - + # Use pagination helper from base class reservations = self.paginate_boto3( - ec2_client, - 'describe_instances', - 'Reservations' + ec2_client, "describe_instances", "Reservations" ) - + # Extract and process instances instances = [] for reservation in reservations: @@ -207,25 +231,31 @@ def discover_instances(self, regions: List[str]) -> List[Dict[str, Any]]: tags[tag["Key"]] = tag["Value"] # Create a standardized instance dictionary - instances.append({ - "id": instance["InstanceId"], - "name": tags.get("Name", "Unnamed"), - "type": instance["InstanceType"], - "status": instance["State"]["Name"], - "state": instance["State"]["Name"], - "region": region, - "tags": tags, - "launch_time": instance["LaunchTime"].isoformat() if "LaunchTime" in instance else None, - "public_ip": instance.get("PublicIpAddress"), - "private_ip": instance.get("PrivateIpAddress"), - "accountId": self.account_id, - "accountName": self.account_name - }) - + instances.append( + { + "id": instance["InstanceId"], + "name": tags.get("Name", "Unnamed"), + "type": instance["InstanceType"], + "status": instance["State"]["Name"], + "state": instance["State"]["Name"], + "region": region, + "tags": tags, + "launch_time": ( + instance["LaunchTime"].isoformat() + if "LaunchTime" in instance + else None + ), + "public_ip": instance.get("PublicIpAddress"), + "private_ip": instance.get("PrivateIpAddress"), + "accountId": self.account_id, + "accountName": self.account_name, + } + ) + logger.info(f"Found {len(instances)} EC2 instances in {region}") all_instances.extend(instances) - + except Exception as e: logger.error(f"Error discovering EC2 instances in {region}: {str(e)}") - + return all_instances diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ecr.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ecr.py index 9989b76f..b3b959be 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ecr.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/ecr.py @@ -1,6 +1,6 @@ -from typing import Dict, List, Any, Optional import logging -from datetime import datetime, timedelta +from datetime import datetime, timedelta, timezone +from typing import Any, Dict, List, Optional from aws_resource_management.managers.base import ResourceManager @@ -9,110 +9,122 @@ class ECRManager(ResourceManager): """Manager for Amazon ECR (Elastic Container Registry) resources.""" - - def __init__(self, credentials: Dict[str, str], account_id: str, account_name: Optional[str] = None): + + def __init__( + self, + credentials: Dict[str, str], + account_id: str, + account_name: Optional[str] = None, + ): """Initialize the ECR resource manager.""" super().__init__(credentials, account_id, account_name) self.resource_type = "ecr_image" - - def start(self, resources: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: + + def start( + self, resources: List[Dict[str, Any]], dry_run: bool = False + ) -> Dict[str, int]: """No-op implementation as ECR images don't have a start operation.""" return {"success": 0, "errors": 0, "skipped": len(resources)} - - def stop(self, resources: List[Dict[str, Any]], dry_run: bool = False) -> Dict[str, int]: + + def stop( + self, resources: List[Dict[str, Any]], dry_run: bool = False + ) -> Dict[str, int]: """No-op implementation as ECR images don't have a stop operation.""" return {"success": 0, "errors": 0, "skipped": len(resources)} - - def discover_images_by_age(self, regions: List[str], age_threshold_days: int = 365) -> Dict[str, List[Dict[str, Any]]]: + + def discover_images_by_age( + self, regions: List[str], age_threshold_days: int = 365 + ) -> Dict[str, List[Dict[str, Any]]]: """Discover ECR images across regions, categorized by age.""" all_images = [] old_images = [] - + for region in regions: images_result = self._discover_images_in_region(region, age_threshold_days) - all_images.extend(images_result['all_images']) - old_images.extend(images_result['old_images']) - - logger.info(f"Found {len(all_images)} ECR images across {len(regions)} regions " - f"({len(old_images)} older than {age_threshold_days} days)") - - return { - 'all_images': all_images, - 'old_images': old_images - } - - def _discover_images_in_region(self, region: str, age_threshold_days: int = 365) -> Dict[str, List[Dict[str, Any]]]: + all_images.extend(images_result["all_images"]) + old_images.extend(images_result["old_images"]) + + logger.info( + f"Found {len(all_images)} ECR images across {len(regions)} regions " + f"({len(old_images)} older than {age_threshold_days} days)" + ) + + return {"all_images": all_images, "old_images": old_images} + + def _discover_images_in_region( + self, region: str, age_threshold_days: int = 365 + ) -> Dict[str, List[Dict[str, Any]]]: """Discover ECR images in a specific region, categorized by age.""" ecr_client = self.get_boto3_client("ecr", region) if not ecr_client: - return {'all_images': [], 'old_images': []} - + return {"all_images": [], "old_images": []} + all_images = [] old_images = [] - threshold_date = datetime.now() - timedelta(days=age_threshold_days) - + # Create timezone-aware datetime for threshold to avoid comparison issues + threshold_date = datetime.now(timezone.utc) - timedelta(days=age_threshold_days) + try: # Get all repositories efficiently using pagination helper repositories = self.paginate_boto3( - ecr_client, - 'describe_repositories', - 'repositories' + ecr_client, "describe_repositories", "repositories" + ) + + repository_names = [repo["repositoryName"] for repo in repositories] + logger.info( + f"Found {len(repository_names)} ECR repositories in region {region}" ) - - repository_names = [repo['repositoryName'] for repo in repositories] - logger.info(f"Found {len(repository_names)} ECR repositories in {region}") - + # For each repository, efficiently get all image IDs for repo_name in repository_names: try: # Get image IDs using pagination helper image_ids = self.paginate_boto3( - ecr_client, - 'list_images', - 'imageIds', - repositoryName=repo_name + ecr_client, "list_images", "imageIds", repositoryName=repo_name ) - + # Process images in chunks due to API limitations chunk_size = 100 for i in range(0, len(image_ids), chunk_size): - chunk = image_ids[i:i + chunk_size] + chunk = image_ids[i : i + chunk_size] if not chunk: continue - + try: # Get details for all images in this chunk response = ecr_client.describe_images( - repositoryName=repo_name, - imageIds=chunk + repositoryName=repo_name, imageIds=chunk ) - + # Process each image - for image_detail in response.get('imageDetails', []): + for image_detail in response.get("imageDetails", []): # Add metadata - image_detail['repositoryName'] = repo_name - image_detail['region'] = region - image_detail['accountId'] = self.account_id + image_detail["repositoryName"] = repo_name + image_detail["region"] = region + image_detail["accountId"] = self.account_id if self.account_name: - image_detail['accountName'] = self.account_name - + image_detail["accountName"] = self.account_name + all_images.append(image_detail) - + # Check if image is older than threshold - if ('imagePushedAt' in image_detail and - image_detail['imagePushedAt'] < threshold_date): - old_images.append(image_detail) + # Make sure we're comparing compatible datetime objects + if "imagePushedAt" in image_detail: + # Both are now timezone-aware so comparison is safe + if image_detail["imagePushedAt"] < threshold_date: + old_images.append(image_detail) except Exception as e: - logger.error(f"Error describing images for repository {repo_name} chunk {i//chunk_size + 1}: {e}") + logger.error( + f"Error describing images for repository {repo_name} chunk {i//chunk_size + 1}: {e}" + ) except Exception as e: logger.error(f"Error processing repository {repo_name}: {e}") - - logger.info(f"Found {len(all_images)} ECR images in {region} ({len(old_images)} older than {age_threshold_days} days)") - + + logger.info( + f"Found {len(all_images)} ECR images in {region} ({len(old_images)} older than {age_threshold_days} days)" + ) + except Exception as e: logger.error(f"Error discovering ECR images in {region}: {e}") - - return { - 'all_images': all_images, - 'old_images': old_images - } + + return {"all_images": all_images, "old_images": old_images} diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/eks.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/eks.py index 3a9913b7..78e7c27c 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/eks.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/eks.py @@ -507,48 +507,53 @@ def _scale_nodegroups( def discover_clusters(self, regions: List[str]) -> List[Dict[str, Any]]: """ Discover EKS clusters across multiple regions. - + Args: regions: List of AWS regions to scan - + Returns: List of EKS cluster dictionaries """ all_clusters = [] - + for region in regions: try: - eks_client = self.get_boto3_client('eks', region) + eks_client = self.get_boto3_client("eks", region) if not eks_client: logger.warning(f"Could not create EKS client in {region}") continue - + # Get clusters response = eks_client.list_clusters() cluster_names = response.get("clusters", []) - + # Handle pagination while "nextToken" in response: response = eks_client.list_clusters(nextToken=response["nextToken"]) cluster_names.extend(response.get("clusters", [])) - + # Get detailed information for each cluster clusters = [] for name in cluster_names: try: # Get cluster details cluster = eks_client.describe_cluster(name=name)["cluster"] - + # Get tags tags = cluster.get("tags", {}) - + # Create a simplified cluster dictionary cluster_dict = { "id": name, "name": name, - "arn": cluster.get("arn", f"arn:aws:eks:{region}:{self.account_id}:cluster/{name}"), + "arn": cluster.get( + "arn", + f"arn:aws:eks:{region}:{self.account_id}:cluster/{name}", + ), "status": cluster.get("status", "UNKNOWN"), - "state": cluster.get("status", "UNKNOWN"), # Add both for consistency + "state": cluster.get( + "status", "UNKNOWN" + ), # Add both for consistency "region": region, "tags": tags, "version": cluster.get("version"), @@ -558,16 +563,20 @@ def discover_clusters(self, regions: List[str]) -> List[Dict[str, Any]]: } if self.account_name: cluster_dict["accountName"] = self.account_name - + clusters.append(cluster_dict) except Exception as e: - logger.warning(f"Error getting details for EKS cluster {name}: {e}") - + logger.warning( + f"Error getting details for EKS cluster {name}: {e}" + ) + logger.info(f"Found {len(clusters)} EKS clusters in {region}") all_clusters.extend(clusters) - + except Exception as e: logger.error(f"Error discovering EKS clusters in {region}: {str(e)}") - - logger.info(f"Found a total of {len(all_clusters)} EKS clusters across all regions") + + logger.info( + f"Found a total of {len(all_clusters)} EKS clusters across all regions" + ) return all_clusters diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py index 20c091d9..f4ffdfed 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/emr.py @@ -33,55 +33,64 @@ def __init__( super().__init__(credentials, account_id, account_name) self.resource_type = "emr_cluster" - def discover_clusters(self, regions: List[str], cluster_states: Optional[List[str]] = None) -> List[Dict[str, Any]]: + def discover_clusters( + self, regions: List[str], cluster_states: Optional[List[str]] = None + ) -> List[Dict[str, Any]]: """ Discover EMR clusters across multiple regions. - + Args: regions: List of AWS regions to scan cluster_states: List of cluster states to include (optional) - + Returns: List of EMR cluster dictionaries """ all_clusters = [] - + # Default cluster states if not provided if cluster_states is None: - cluster_states = ["STARTING", "BOOTSTRAPPING", "RUNNING", "WAITING", "TERMINATING"] - + cluster_states = [ + "STARTING", + "BOOTSTRAPPING", + "RUNNING", + "WAITING", + "TERMINATING", + ] + for region in regions: try: - emr_client = self.get_boto3_client('emr', region) + emr_client = self.get_boto3_client("emr", region) if not emr_client: logger.warning(f"Could not create EMR client in {region}") continue - + # Get clusters response = emr_client.list_clusters(ClusterStates=cluster_states) clusters = response.get("Clusters", []) - + # Handle pagination while "Marker" in response: response = emr_client.list_clusters( - Marker=response["Marker"], - ClusterStates=cluster_states + Marker=response["Marker"], ClusterStates=cluster_states ) clusters.extend(response.get("Clusters", [])) - + # Process each cluster processed_clusters = [] for cluster in clusters: try: # Get cluster details - cluster_details = emr_client.describe_cluster(ClusterId=cluster["Id"]) + cluster_details = emr_client.describe_cluster( + ClusterId=cluster["Id"] + ) cluster_info = cluster_details.get("Cluster", {}) - + # Convert tags tags = {} for tag in cluster_info.get("Tags", []): tags[tag.get("Key", "")] = tag.get("Value", "") - + # Create a simplified cluster dictionary cluster_dict = { "id": cluster["Id"], @@ -90,25 +99,35 @@ def discover_clusters(self, regions: List[str], cluster_states: Optional[List[st "state": cluster.get("Status", {}).get("State", "UNKNOWN"), "region": region, "tags": tags, - "creation_time": cluster.get("Status", {}).get("Timeline", {}).get("CreationDateTime"), - "termination_time": cluster.get("Status", {}).get("Timeline", {}).get("EndDateTime"), - "cluster_type": cluster_info.get("InstanceCollectionType", "Unknown"), + "creation_time": cluster.get("Status", {}) + .get("Timeline", {}) + .get("CreationDateTime"), + "termination_time": cluster.get("Status", {}) + .get("Timeline", {}) + .get("EndDateTime"), + "cluster_type": cluster_info.get( + "InstanceCollectionType", "Unknown" + ), "accountId": self.account_id, } if self.account_name: cluster_dict["accountName"] = self.account_name - + processed_clusters.append(cluster_dict) except Exception as e: - logger.warning(f"Error getting details for EMR cluster {cluster['Id']}: {e}") - + logger.warning( + f"Error getting details for EMR cluster {cluster['Id']}: {e}" + ) + logger.info(f"Found {len(processed_clusters)} EMR clusters in {region}") all_clusters.extend(processed_clusters) - + except Exception as e: logger.error(f"Error discovering EMR clusters in {region}: {str(e)}") - - logger.info(f"Found a total of {len(all_clusters)} EMR clusters across all regions") + + logger.info( + f"Found a total of {len(all_clusters)} EMR clusters across all regions" + ) return all_clusters def discover_all_clusters(self, regions: List[str]) -> List[Dict[str, Any]]: diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/rds.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/rds.py index 6127b1a3..d5add442 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/rds.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/managers/rds.py @@ -2,8 +2,8 @@ RDS resource manager class. """ -from typing import Any, Dict, List, Optional from collections import defaultdict +from typing import Any, Dict, List, Optional from aws_resource_management.config_manager import get_config from aws_resource_management.logging_setup import setup_logging @@ -34,21 +34,23 @@ def stop( # Group instances by region available_instances_by_region = defaultdict(list) stats = {"success": 0, "errors": 0, "skipped": 0} - + # Filter instances - keep only those in 'available' state and not excluded for instance in instances: if self.should_exclude(instance): stats["skipped"] += 1 continue - + if instance["status"] == "available": available_instances_by_region[instance["region"]].append(instance) else: - logger.debug(f"Skipping RDS instance {instance['id']} in state {instance['status']} (not available)") + logger.debug( + f"Skipping RDS instance {instance['id']} in state {instance['status']} (not available)" + ) stats["skipped"] += 1 - + timestamp = self.get_timestamp() - + # Process each region's instances for region, region_instances in available_instances_by_region.items(): rds_client = self.get_boto3_client("rds", region) @@ -56,44 +58,52 @@ def stop( logger.error(f"Failed to create RDS client for region {region}") stats["errors"] += len(region_instances) continue - + # Process instances in batches for i in range(0, len(region_instances), self.MAX_DB_OPERATIONS): - batch = region_instances[i:i + self.MAX_DB_OPERATIONS] - + batch = region_instances[i : i + self.MAX_DB_OPERATIONS] + # Unfortunately, RDS doesn't support bulk tagging or stopping # Process each instance individually but within a batch context for instance in batch: try: db_id = instance["id"] arn = f"arn:{self.get_partition_for_region(region)}:rds:{region}:{self.account_id}:db:{db_id}" - + # Tag the instance before stopping if not dry_run: logger.debug(f"Tagging RDS instance {db_id}") rds_client.add_tags_to_resource( ResourceName=arn, - Tags=[{"Key": config.get("stop_tag"), "Value": timestamp}], + Tags=[ + {"Key": config.get("stop_tag"), "Value": timestamp} + ], ) logger.info(f"Stopping RDS instance {db_id} in {region}") rds_client.stop_db_instance(DBInstanceIdentifier=db_id) else: - logger.info(f"[DRY RUN] Would stop RDS instance {db_id} in {region}") - + logger.info( + f"[DRY RUN] Would stop RDS instance {db_id} in {region}" + ) + self.log_action( db_id, region, "stop", resource_name=instance.get("name"), - dry_run=dry_run + dry_run=dry_run, ) stats["success"] += 1 except Exception as e: - logger.error(f"Failed to stop RDS instance {instance['id']}: {e}") + logger.error( + f"Failed to stop RDS instance {instance['id']}: {e}" + ) stats["errors"] += 1 - logger.info(f"RDS stop summary: {stats['success']} stopped, {stats['skipped']} skipped, {stats['errors']} errors") + logger.info( + f"RDS stop summary: {stats['success']} stopped, {stats['skipped']} skipped, {stats['errors']} errors" + ) return stats def start( @@ -103,19 +113,21 @@ def start( # Group instances by region stopped_instances_by_region = defaultdict(list) stats = {"success": 0, "errors": 0, "skipped": 0} - + # Filter instances - keep only those in 'stopped' state and not excluded for instance in instances: if self.should_exclude(instance): stats["skipped"] += 1 continue - + if instance["status"] == "stopped": stopped_instances_by_region[instance["region"]].append(instance) else: - logger.debug(f"Skipping RDS instance {instance['id']} in state {instance['status']} (not stopped)") + logger.debug( + f"Skipping RDS instance {instance['id']} in state {instance['status']} (not stopped)" + ) stats["skipped"] += 1 - + # Process each region's instances for region, region_instances in stopped_instances_by_region.items(): rds_client = self.get_boto3_client("rds", region) @@ -123,22 +135,22 @@ def start( logger.error(f"Failed to create RDS client for region {region}") stats["errors"] += len(region_instances) continue - + # Process instances in batches for i in range(0, len(region_instances), self.MAX_DB_OPERATIONS): - batch = region_instances[i:i + self.MAX_DB_OPERATIONS] - + batch = region_instances[i : i + self.MAX_DB_OPERATIONS] + # Unfortunately, RDS doesn't support bulk operations for instance in batch: try: db_id = instance["id"] arn = f"arn:{self.get_partition_for_region(region)}:rds:{region}:{self.account_id}:db:{db_id}" - + # Start the instance if not dry_run: logger.info(f"Starting RDS instance {db_id} in {region}") rds_client.start_db_instance(DBInstanceIdentifier=db_id) - + # Remove the stop tag logger.debug(f"Removing stop tag from RDS instance {db_id}") rds_client.remove_tags_from_resource( @@ -146,40 +158,44 @@ def start( TagKeys=[config.get("stop_tag")], ) else: - logger.info(f"[DRY RUN] Would start RDS instance {db_id} in {region}") - + logger.info( + f"[DRY RUN] Would start RDS instance {db_id} in {region}" + ) + self.log_action( db_id, region, "start", resource_name=instance.get("name"), - dry_run=dry_run + dry_run=dry_run, ) stats["success"] += 1 except Exception as e: - logger.error(f"Failed to start RDS instance {instance['id']}: {e}") + logger.error( + f"Failed to start RDS instance {instance['id']}: {e}" + ) stats["errors"] += 1 - logger.info(f"RDS start summary: {stats['success']} started, {stats['skipped']} skipped, {stats['errors']} errors") + logger.info( + f"RDS start summary: {stats['success']} started, {stats['skipped']} skipped, {stats['errors']} errors" + ) return stats def discover_instances(self, regions: List[str]) -> List[Dict[str, Any]]: """Discover RDS instances across regions using efficient pagination.""" all_instances = [] - + for region in regions: try: - rds_client = self.get_boto3_client('rds', region) + rds_client = self.get_boto3_client("rds", region) if not rds_client: continue - + # Use pagination helper from base class db_instances = self.paginate_boto3( - rds_client, - 'describe_db_instances', - 'DBInstances' + rds_client, "describe_db_instances", "DBInstances" ) - + # Process each instance instances = [] for db in db_instances: @@ -193,33 +209,37 @@ def discover_instances(self, regions: List[str]) -> List[Dict[str, Any]]: for item in tags_response.get("TagList", []) } except Exception as e: - logger.warning(f"Error getting tags for RDS instance {db['DBInstanceIdentifier']}: {e}") + logger.warning( + f"Error getting tags for RDS instance {db['DBInstanceIdentifier']}: {e}" + ) tags = {} # Create standardized instance dictionary - instances.append({ - "id": db["DBInstanceIdentifier"], - "name": db["DBInstanceIdentifier"], - "engine": db["Engine"], - "engine_version": db.get("EngineVersion"), - "status": db["DBInstanceStatus"], - "state": db["DBInstanceStatus"], - "region": region, - "tags": tags, - "instance_class": db.get("DBInstanceClass"), - "storage": db.get("AllocatedStorage"), - "multi_az": db.get("MultiAZ", False), - "public": db.get("PubliclyAccessible", False), - "endpoint": db.get("Endpoint", {}).get("Address"), - "port": db.get("Endpoint", {}).get("Port"), - "accountId": self.account_id, - "accountName": self.account_name - }) - + instances.append( + { + "id": db["DBInstanceIdentifier"], + "name": db["DBInstanceIdentifier"], + "engine": db["Engine"], + "engine_version": db.get("EngineVersion"), + "status": db["DBInstanceStatus"], + "state": db["DBInstanceStatus"], + "region": region, + "tags": tags, + "instance_class": db.get("DBInstanceClass"), + "storage": db.get("AllocatedStorage"), + "multi_az": db.get("MultiAZ", False), + "public": db.get("PubliclyAccessible", False), + "endpoint": db.get("Endpoint", {}).get("Address"), + "port": db.get("Endpoint", {}).get("Port"), + "accountId": self.account_id, + "accountName": self.account_name, + } + ) + logger.info(f"Found {len(instances)} RDS instances in {region}") all_instances.extend(instances) - + except Exception as e: logger.error(f"Error discovering RDS instances in {region}: {str(e)}") - + return all_instances diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/reporting.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/reporting.py index dce2fc3e..25c611f3 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/reporting.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/reporting.py @@ -3,11 +3,14 @@ """ import logging -from typing import Dict, Any, List +from typing import Any, Dict, List logger = logging.getLogger(__name__) -def print_resource_summary(stats: Dict[str, Any], action: str, dry_run: bool = False) -> None: + +def print_resource_summary( + stats: Dict[str, Any], action: str, dry_run: bool = False +) -> None: """ Print a detailed summary of resources processed, broken down by resource type. @@ -17,9 +20,7 @@ def print_resource_summary(stats: Dict[str, Any], action: str, dry_run: bool = F dry_run: Whether this was a dry run """ logger.info(f"\n{'=' * 30} SUMMARY {'=' * 30}") - logger.info( - f"ACTION: {action.upper()} {'(DRY RUN)' if dry_run else ''}" - ) + logger.info(f"ACTION: {action.upper()} {'(DRY RUN)' if dry_run else ''}") # General stats logger.info(f"\nGENERAL STATISTICS:") @@ -81,13 +82,14 @@ def print_resource_summary(stats: Dict[str, Any], action: str, dry_run: bool = F # Error summary if there were any errors print_error_summary(stats.get("errors", [])) - + logger.info(f"{'=' * 68}") + def print_error_summary(errors: List[str]) -> None: """ Print a summary of errors that occurred during processing. - + Args: errors: List of error messages """ @@ -101,10 +103,11 @@ def print_error_summary(errors: List[str]) -> None: f" ... and {len(errors) - 5} more errors (see log for details)" ) + def initialize_stats() -> Dict[str, Any]: """ Initialize a statistics dictionary with default values. - + Returns: A dictionary with initialized statistics counters """ @@ -120,7 +123,7 @@ def initialize_stats() -> Dict[str, Any]: "ecr_images_found": 0, "ecr_old_images_found": 0, } - + # Initialize resource-specific counters for resource_type in ["ecr", "ec2", "rds", "eks", "emr"]: for stat_type in [ @@ -131,5 +134,5 @@ def initialize_stats() -> Dict[str, Any]: "errors", ]: stats[f"{resource_type}_{stat_type}"] = 0 - + return stats diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/resource_controller.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/resource_controller.py index 82f6a459..f40023b5 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/resource_controller.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/resource_controller.py @@ -16,10 +16,12 @@ import sys from typing import Any, Dict, List, Optional, Union +from aws_resource_management.aws_utils import ( + get_available_regions, # Changed from get_enabled_regions +) from aws_resource_management.aws_utils import ( get_account_list, get_credentials, - get_enabled_regions, ) from aws_resource_management.config_manager import get_config from aws_resource_management.core import ResourceManager @@ -260,7 +262,9 @@ def get_regions_for_account(self, account_id: str) -> List[str]: return [] # Get enabled regions - regions = get_enabled_regions(credentials, partition=self.partition) + regions = get_available_regions( + credentials, partition=self.partition + ) # Changed from get_enabled_regions return regions except Exception as e: diff --git a/local-app/python-tools/gfl-resource-actions/discovery.py b/local-app/python-tools/gfl-resource-actions/discovery.py index 6e01aeb5..78839ace 100644 --- a/local-app/python-tools/gfl-resource-actions/discovery.py +++ b/local-app/python-tools/gfl-resource-actions/discovery.py @@ -3,7 +3,7 @@ """ import config -from aws_utils import create_boto3_client +from aws_resource_management.aws_utils import create_boto3_client from logging_utils import log_action_to_csv, setup_logging logger = setup_logging() diff --git a/local-app/python-tools/gfl-resource-actions/main.py b/local-app/python-tools/gfl-resource-actions/main.py index 3ed280d7..a60ead0a 100644 --- a/local-app/python-tools/gfl-resource-actions/main.py +++ b/local-app/python-tools/gfl-resource-actions/main.py @@ -4,10 +4,16 @@ """ import argparse +import concurrent.futures # Add missing import +import sys # Add missing import from concurrent.futures import ThreadPoolExecutor import config -from aws_utils import assume_role, get_available_regions, get_organization_accounts +from aws_resource_management.aws_utils import ( + assume_role, + get_available_regions, + get_organization_accounts, +) from discovery import get_account_resources from logging_utils import setup_logging from managers import EC2Manager, EKSManager, EMRManager, RDSManager diff --git a/local-app/python-tools/gfl-resource-actions/managers/base.py b/local-app/python-tools/gfl-resource-actions/managers/base.py index 961d1b10..4712357d 100644 --- a/local-app/python-tools/gfl-resource-actions/managers/base.py +++ b/local-app/python-tools/gfl-resource-actions/managers/base.py @@ -10,7 +10,7 @@ import boto3 import config -from aws_utils import get_partition_for_region +from aws_resource_management.aws_utils import get_partition_for_region from botocore.exceptions import ClientError from logging_utils import log_action_to_csv, setup_logging