diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/account_manager.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/account_manager.py new file mode 100644 index 00000000..0b234380 --- /dev/null +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/account_manager.py @@ -0,0 +1,111 @@ +"""Account management for AWS Resource Management.""" + +from typing import Dict, List, Optional + +from aws_resource_management.utils import ( + ensure_valid_account_name, + get_account_list, + get_credentials, + get_organization_accounts, + logger, +) + + +class AccountManager: + """Manages AWS account discovery and filtering.""" + + def __init__( + self, + profile_name: Optional[str] = None, + use_profiles: bool = False, + max_concurrent_accounts: int = 3 + ) -> None: + """Initialize the account manager. + + Args: + profile_name: AWS profile name + use_profiles: Whether to use AWS profiles + max_concurrent_accounts: Maximum number of concurrent account operations + """ + self.profile_name = profile_name + self.use_profiles = use_profiles + self.MAX_CONCURRENT_ACCOUNTS = max_concurrent_accounts + + def get_accounts(self) -> List[Dict[str, str]]: + """Get list of accounts from the current organization or profiles.""" + if self.use_profiles: + logger.info("Using AWS SSO profiles for account discovery") + return get_organization_accounts() + else: + logger.info("Using AWS Organizations API for account discovery") + return get_account_list() + + def pre_validate_accounts( + self, accounts: List[Dict[str, str]] + ) -> List[Dict[str, str]]: + """Pre-validate which accounts we can authenticate with.""" + valid_accounts = [] + for account in accounts: + account_id = account.get("account_id") + try: + credentials = get_credentials(account_id, self.profile_name) + if credentials: + valid_accounts.append(account) + else: + logger.debug( + f"No credentials available for account {account_id}, skipping" + ) + except Exception as e: + logger.debug(f"Cannot authenticate with account {account_id}: {str(e)}") + return valid_accounts + + def filter_accounts( + self, + accounts: List[Dict[str, str]], + exclude_accounts: List[str], + specific_account: Optional[str] = None + ) -> List[Dict[str, str]]: + """Filter accounts based on exclusions and specific account.""" + filtered_accounts = accounts + + # Filter for specific account if provided + if specific_account: + filtered_accounts = [ + acc for acc in filtered_accounts + if acc.get("account_id") == specific_account + ] + if not filtered_accounts: + logger.error(f"Specified account {specific_account} not found or not accessible") + return [] + + # Apply exclusions + return [ + acc for acc in filtered_accounts + if acc.get("account_id") not in exclude_accounts + ] + + def normalize_account_names( + self, accounts: List[Dict[str, str]] + ) -> List[Dict[str, str]]: + """Ensure all accounts have valid names.""" + for account in accounts: + account_id = account.get("account_id") + account_name = ensure_valid_account_name( + account_id, account.get("account_name") + ) + account["account_name"] = account_name + return accounts + + def log_account_summary( + self, all_accounts: List[Dict[str, str]], + valid_accounts: List[Dict[str, str]], + exclude_accounts: List[str] + ) -> None: + """Log summary of accounts being processed.""" + logger.info(f"Found {len(valid_accounts)} accessible accounts to process") + if len(all_accounts) > len(valid_accounts): + logger.info( + f"Skipped {len(all_accounts) - len(valid_accounts)} inaccessible accounts" + ) + if exclude_accounts: + logger.info(f"Excluding {len(exclude_accounts)} accounts") 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 afacdf7a..f933f469 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 @@ -27,7 +27,7 @@ class ResourceManager: - """Main resource manager that orchestrates + """Main resource manager that orchestrates operations across accounts and resources.""" def __init__( @@ -39,13 +39,12 @@ def __init__( ) -> None: """Initialize the resource manager.""" self.config = get_config() - self.account_list = [] self.profile_name = profile_name self.use_profiles = use_profiles self.discover_regions = discover_regions self.partition = partition self.region_cache = {} - self.MAX_CONCURRENT_ACCOUNTS = 20 # Limit concurrent account processing + self.MAX_CONCURRENT_ACCOUNTS = 3 # Limit concurrent account processing # Get resource registry self.resource_registry = get_registry() @@ -65,102 +64,58 @@ def process_accounts( stats: Dict[str, Any] = None, ) -> Dict[str, Any]: """Process accounts and perform the specified action.""" - try: - # Initialize parameters - exclude_accounts = exclude_accounts or [] - exclude_resources = exclude_resources or [] - exclude_regions = exclude_regions or [] - stats = stats or initialize_stats() - - # Log which resource types are being excluded for debugging - if exclude_resources: - logger.info(f"Excluding resource types: {', '.join(exclude_resources)}") - - # Get account list from the organization you're currently authenticated with - accounts = self._get_accounts() - - # Pre-filter accounts that we can authenticate with - valid_accounts = self._pre_validate_accounts(accounts) - - # Log summary of accounts - logger.info(f"Found {len(valid_accounts)} accessible accounts to process") - if len(accounts) > len(valid_accounts): - logger.info( - f"Skipped {len(accounts) - len(valid_accounts)} \ - inaccessible accounts" - ) - if exclude_accounts: - logger.info(f"Excluding {len(exclude_accounts)} accounts") - - # 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 valid_accounts: - account_id = account.get("account_id") - if account_id in exclude_accounts: - stats["accounts_skipped"] += 1 - continue - - future = executor.submit( - self._process_account_safely, - account=account, - regions=regions, - action=action, - exclude_resources=exclude_resources, - exclude_regions=exclude_regions, - dry_run=dry_run, - ) - 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 - self._log_summary(stats, action) - return stats - - except KeyboardInterrupt: - logger.warning("Account processing interrupted by user") - raise + # Initialize parameters with defaults + exclude_accounts = exclude_accounts or [] + exclude_resources = exclude_resources or [] + exclude_regions = exclude_regions or [] + stats = stats or initialize_stats() + + # Log excluded resource types + if exclude_resources: + logger.info(f"Excluding resource types: {', '.join(exclude_resources)}") + + # Get and validate accounts + accounts = self._get_accounts() + valid_accounts = self._pre_validate_accounts(accounts) + + # Log account summary + self._log_account_summary(accounts, valid_accounts, exclude_accounts) + + # Process accounts using thread pool + with concurrent.futures.ThreadPoolExecutor( + max_workers=self.MAX_CONCURRENT_ACCOUNTS + ) as executor: + futures = self._create_account_futures( + executor, valid_accounts, exclude_accounts, regions, + action, exclude_resources, exclude_regions, dry_run + ) + + # Process results + for future, account_id in futures: + self._process_account_result(future, account_id, stats) + + # Log summary and return stats + self._log_summary(stats, action) + return stats + + def _get_accounts(self) -> List[Dict[str, str]]: + """Get list of accounts from the current organization or profiles.""" + if self.use_profiles: + logger.info("Using AWS SSO profiles for account discovery") + return get_organization_accounts() + else: + logger.info("Using AWS Organizations API for account discovery") + return get_account_list() def _pre_validate_accounts( self, accounts: List[Dict[str, str]] ) -> List[Dict[str, str]]: - """ - Pre-validate which accounts we can actually authenticate with. - - Args: - accounts: List of account dictionaries with account_id and account_name - - Returns: - List of accounts we can authenticate with - """ + """Pre-validate which accounts we can authenticate with.""" valid_accounts = [] for account in accounts: account_id = account.get("account_id") try: - # Try to get credentials without actually making any API calls - credentials = get_credentials(account_id, self.profile_name) + credentials = get_credentials(account_id, self) if credentials: valid_accounts.append(account) else: @@ -169,19 +124,72 @@ def _pre_validate_accounts( ) except Exception as e: logger.debug(f"Cannot authenticate with account {account_id}: {str(e)}") - return valid_accounts - def _get_accounts(self) -> List[Dict[str, str]]: - """Get list of accounts from the current organization only.""" - # Always use the Organizations API by default to get accounts from current org - # Only use profiles if explicitly requested (which should be rare) - if self.use_profiles: - logger.info("Using AWS SSO profiles for account discovery") - return get_organization_accounts() - else: - logger.info("Using AWS Organizations API for account discovery") - return get_account_list() + def _log_account_summary( + self, all_accounts: List[Dict[str, str]], + valid_accounts: List[Dict[str, str]], + exclude_accounts: List[str] + ) -> None: + """Log summary of accounts being processed.""" + logger.info(f"Found {len(valid_accounts)} accessible accounts to process") + if len(all_accounts) > len(valid_accounts): + logger.info( + f"Skipped {len(all_accounts) - len(valid_accounts)} inaccessible accounts" + ) + if exclude_accounts: + logger.info(f"Excluding {len(exclude_accounts)} accounts") + + def _create_account_futures( + self, + executor: concurrent.futures.ThreadPoolExecutor, + valid_accounts: List[Dict[str, str]], + exclude_accounts: List[str], + regions: List[str], + action: str, + exclude_resources: List[str], + exclude_regions: List[str], + dry_run: bool + ) -> List[tuple]: + """Create futures for each account to be processed.""" + futures = [] + for account in valid_accounts: + account_id = account.get("account_id") + if account_id in exclude_accounts: + continue + + future = executor.submit( + self._process_account_safely, + account=account, + regions=regions, + action=action, + exclude_resources=exclude_resources, + exclude_regions=exclude_regions, + dry_run=dry_run, + ) + futures.append((future, account_id)) + return futures + + def _process_account_result( + self, future: concurrent.futures.Future, account_id: str, stats: Dict[str, Any] + ) -> None: + """Process a completed account future and update stats.""" + try: + account_stats = future.result() + 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)}" + ) def _process_account_safely( self, @@ -194,16 +202,11 @@ def _process_account_safely( ) -> Optional[Dict[str, Any]]: """Process a single account with error handling.""" account_id = account.get("account_id") - # Ensure account_name is valid and update the original account dictionary account_name = ensure_valid_account_name( account_id, account.get("account_name") ) account["account_name"] = account_name # Update in original dict - # Add debug logging to trace the account name through the process - logger.debug(f"Processing account safely with name: {account_name}") - - # Use the validated account_name when logging logger.info(f"Processing account {account_name} ({account_id})") try: @@ -213,15 +216,6 @@ def _process_account_safely( logger.warning(f"Could not obtain credentials for account {account_id}") return None - # Get account alias - try: - from aws_resource_management.utils import get_account_alias - account_alias = get_account_alias(account_id, credentials) - logger.debug(f"Retrieved account alias for {account_id}: {account_alias}") - except Exception as e: - logger.debug(f"Error getting account alias for {account_id}: {str(e)}") - account_alias = None - # Get regions for this account account_regions = self._get_account_regions( account_id, credentials, regions, exclude_regions @@ -234,50 +228,32 @@ def _process_account_safely( account_stats = initialize_stats() account_stats["regions_processed"] = len(account_regions) account_stats["regions_by_account"][account_id] = account_regions - # Store account name in stats for better reporting account_stats["account_name"] = account_name - # Process the account with its determined regions - logger.info( - f"Processing account: {account_name} ({account_id}) in \ - {len(account_regions)} regions" - ) - # Auto-detect partition if not specified partition = ( detect_partition_from_credentials(credentials) if not self.partition else self.partition ) - - # Log partition info logger.info( f"Using partition {partition} for account {account_name} ({account_id})" ) - # Get EMR states from registry instead of hardcoding - emr_states = self._get_emr_states(exclude_resources) - try: - # Get all resources - this might raise AuthFailure - resources = get_account_resources( - credentials=credentials, - regions=account_regions, - exclude_resources=exclude_resources, - account_id=account_id, - account_name=account_name, # Explicitly pass as kwarg - account_alias=account_alias, # Pass account alias - emr_cluster_states=emr_states, + # Get account resources + resources = self._get_account_resources( + credentials, account_regions, exclude_resources, + account_id, account_name, partition ) - + if not resources: logger.error( - f"Failed to get resources for account \ - {account_id} ({account_name})" + f"Failed to get resources for account {account_id} ({account_name})" ) return None - # Update statistics using the registry instead of hardcoded logic + # Update statistics using the registry self._process_resources_with_registry(resources, account_stats) # Log resource counts @@ -285,45 +261,29 @@ def _process_account_safely( account_id, account_name, resources, account_stats ) - # Execute actions using registry pattern - # instead of manual manager creation + # Execute actions self._execute_actions_with_registry( - action, - credentials, - account_id, - account_name, # Make sure we're passing the account_name here - resources, - exclude_resources, - dry_run, - account_stats, - ) - - # Log discovery initialization with account name - logger.info( - f"Resource discovery initialized for account \ - {account_name} ({account_id}) in partition {partition}" + action, credentials, account_id, account_name, + 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", + "AuthFailure", "InvalidClientTokenId", + "UnauthorizedOperation", "AccessDenied", ]: logger.error( - f"Authentication failure for account \ - {account_id} ({account_name}): {e}" + f"Authentication failure for account {account_id} ({account_name}): {e}" ) account_stats["errors"].append( - f"Authentication failure for account \ - {account_id} ({account_name}): {str(e)}" + f"Authentication failure for account {account_id} ({account_name}): {str(e)}" ) return account_stats - raise # Re-raise other client errors + raise except Exception as e: logger.error( @@ -336,6 +296,27 @@ def _process_account_safely( return account_stats return None + def _get_account_resources( + self, + credentials: Dict[str, str], + account_regions: List[str], + exclude_resources: List[str], + account_id: str, + account_name: str, + partition: str + ) -> Dict[str, List[Dict[str, Any]]]: + """Get all resources for an account.""" + emr_states = self._get_emr_states(exclude_resources) + + return get_account_resources( + credentials=credentials, + regions=account_regions, + exclude_resources=exclude_resources, + account_id=account_id, + account_name=account_name, + emr_cluster_states=emr_states, + ) + def _get_account_regions( self, account_id: str, @@ -344,60 +325,19 @@ def _get_account_regions( exclude_regions: List[str], ) -> List[str]: """Determine regions to use for an account.""" + # If region discovery is enabled, try to discover regions if self.discover_regions: - try: - logger.info(f"Discovering enabled regions for account {account_id}") - discovered_regions = get_enabled_regions(credentials, self.partition) - # Filter discovered regions for the partition - if self.partition: - discovered_regions = filter_regions_for_partition( - discovered_regions, self.partition - ) - - # Apply exclusions - 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 [] - # If we have provided regions, filter them for the partition + return self._discover_regions( + account_id, credentials, exclude_regions + ) + + # If provided regions, filter them elif provided_regions: - if self.partition: - # Use centralized filter - valid_regions = filter_regions_for_partition( - provided_regions, self.partition - ) - filtered_regions = [ - r for r in valid_regions if r not in exclude_regions - ] - else: - 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 or invalid, using \ - defaults: {', '.join(filtered_regions)}" - ) - return filtered_regions + return self._filter_provided_regions( + provided_regions, exclude_regions, credentials + ) + + # Otherwise use default regions else: default_regions = self._get_default_regions(credentials) logger.info( @@ -405,19 +345,68 @@ def _get_account_regions( ) return default_regions - def _get_default_regions( - self, credentials: Optional[Dict[str, str]] = None + def _discover_regions( + self, account_id: str, credentials: Dict[str, str], exclude_regions: List[str] ) -> List[str]: - """ - Get default regions for the current partition. + """Discover enabled regions for an account.""" + try: + logger.info(f"Discovering enabled regions for account {account_id}") + discovered_regions = get_enabled_regions(credentials, self.partition) + + # Filter discovered regions for the partition + if self.partition: + discovered_regions = filter_regions_for_partition( + discovered_regions, self.partition + ) - Args: - credentials: Optional credentials to use for partition detection + # Apply exclusions + account_regions = [ + r for r in discovered_regions if r not in exclude_regions + ] - Returns: - List of default regions for the partition - """ - # If partition is not explicitly set, try to detect it from credentials + 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)}" + ) + return [] + + def _filter_provided_regions( + self, provided_regions: List[str], exclude_regions: List[str], + credentials: Dict[str, str] + ) -> List[str]: + """Filter provided regions based on partition and exclusions.""" + if self.partition: + # Use centralized filter + valid_regions = filter_regions_for_partition( + provided_regions, self.partition + ) + filtered_regions = [ + r for r in valid_regions if r not in exclude_regions + ] + else: + 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 or invalid, using defaults: " + f"{', '.join(filtered_regions)}" + ) + + return filtered_regions + + def _get_default_regions( + self, credentials: Optional[Dict[str, str]] = None + ) -> List[str]: + """Get default regions for the current partition.""" + # Try to detect partition from credentials if not explicitly set partition = self.partition if not partition and credentials: try: @@ -426,107 +415,86 @@ def _get_default_regions( except Exception as e: logger.warning(f"Failed to auto-detect partition from credentials: {e}") - # If still no partition, use GovCloud AWS (default for this environment) + # Use GovCloud AWS as the default partition partition = partition or DEFAULT_PARTITION - - # Use centralized utility return get_default_regions_for_partition(partition) def _get_emr_states(self, exclude_resources: List[str]) -> Optional[List[str]]: """Get valid EMR states for discovery.""" if "emr" not in exclude_resources: - # Get EMR states from registry or config instead of hardcoding + # Get EMR states from registry or config emr_manager = self.resource_registry.get_manager_class("emr") if hasattr(emr_manager, "VALID_CLUSTER_STATES"): return emr_manager.VALID_CLUSTER_STATES else: - # Fallback to hardcoded states if not defined in manager + # Fallback to default states return [ - "STARTING", - "BOOTSTRAPPING", - "RUNNING", - "WAITING", - "TERMINATING", + "STARTING", "BOOTSTRAPPING", "RUNNING", + "WAITING", "TERMINATING", ] return None def _process_resources_with_registry( self, resources: Dict[str, List[Dict[str, Any]]], stats: Dict[str, Any] ) -> None: - """Process resources using resource registry instead of hardcoded logic.""" - # Use registry to process resources and update stats + """Process resources using resource registry.""" for resource_type in self.RESOURCE_TYPES: manager_class = self.resource_registry.get_manager_class(resource_type) if not manager_class: continue - # Get the appropriate resource key based on resource type + + # Get resource list for this type resource_key = self.resource_registry.get_resource_key(resource_type) resource_list = resources.get(resource_key, []) - # Update stats - use the proper key format (resource_type_found) + # Update stats stats_key = f"{resource_type}_found" if stats_key not in stats: stats[stats_key] = 0 stats[stats_key] += len(resource_list) - # Log the resource count for debugging - logger.debug( - f"Added {len(resource_list)} {resource_type} resources to stats" - ) - - # Let the manager normalize resources if needed + # Normalize resources if needed if hasattr(manager_class, "normalize_resources"): manager_class.normalize_resources(resource_list) - # Let the manager update specific stats + + # Let manager update specific stats if hasattr(manager_class, "update_stats"): manager_class.update_stats(resource_list, stats) def _log_resource_counts( self, account_id: str, - account_name: str, # Added account_name parameter + account_name: str, resources: Dict[str, List[Dict[str, Any]]], stats: Dict[str, Any], ) -> None: """Log discovered resource counts.""" - # Use registry to get resource keys and log counts for resource_type in self.RESOURCE_TYPES: resource_key = self.resource_registry.get_resource_key(resource_type) - # Special handling for ECR which may have both regular and old images + + # Special handling for ECR if resource_type == "ecr": ecr_images = len(resources.get(resource_key, [])) if ecr_images > 0: logger.info(f"Total ecr resources discovered: {ecr_images}") - # Update the account-specific stats stats["ecr_found"] = stats.get("ecr_found", 0) + ecr_images - # Also track old images if available + + # Track old images old_images = resources.get("ecr_old_images", []) if old_images: - stats["ecr_old_images"] = stats.get("ecr_old_images", 0) + len( - old_images - ) + stats["ecr_old_images"] = stats.get("ecr_old_images", 0) + len(old_images) else: resource_count = len(resources.get(resource_key, [])) if resource_count > 0: - logger.info( - f"Total {resource_type} resources discovered: {resource_count}" - ) - # Update the account-specific stats with the explicit key name - stats[f"{resource_type}_found"] = ( - stats.get(f"{resource_type}_found", 0) + resource_count - ) - # Debug: log the updated stat - logger.debug( - f"Updated stats[{resource_type}_found] to \ - {stats.get(f'{resource_type}_found')}" - ) + logger.info(f"Total {resource_type} resources discovered: {resource_count}") + stats[f"{resource_type}_found"] = stats.get(f"{resource_type}_found", 0) + resource_count def _execute_actions_with_registry( self, action: str, credentials: Dict[str, str], account_id: str, - account_name: str, # Ensure we're using account_name + account_name: str, resources: Dict[str, List[Dict[str, Any]]], exclude_resources: List[str], dry_run: bool, @@ -534,77 +502,79 @@ def _execute_actions_with_registry( ) -> None: """Execute resource actions using the resource registry.""" for resource_type in self.RESOURCE_TYPES: + # Skip excluded resource types if resource_type in exclude_resources: - # Skip this resource type resource_key = self.resource_registry.get_resource_key(resource_type) resource_list = resources.get(resource_key, []) - stats[f"{resource_type}_skipped"] = stats.get( - f"{resource_type}_skipped", 0 - ) + len(resource_list) + stats[f"{resource_type}_skipped"] = stats.get(f"{resource_type}_skipped", 0) + len(resource_list) continue - # Get the resource manager from registry + # Get manager class and resources manager_class = self.resource_registry.get_manager_class(resource_type) if not manager_class: continue - # Get appropriate resource list resource_key = self.resource_registry.get_resource_key(resource_type) resource_list = resources.get(resource_key, []) - - # Special handling for ECR images - - # ensure they're counted in stats even if no action is performed + + # Special handling for ECR if resource_type == "ecr" and resource_list: - # Count total ECR images stats["ecr_found"] = stats.get("ecr_found", 0) + len(resource_list) - - # Count old images if available old_images = resources.get("ecr_old_images", []) - stats["ecr_old_images"] = stats.get("ecr_old_images", 0) + len( - old_images - ) + stats["ecr_old_images"] = stats.get("ecr_old_images", 0) + len(old_images) if not resource_list: continue - for region in self.region_cache.get(account_id, []): - # Filter resources for this region - region_resources = [ - r for r in resource_list if r.get("region") == region - ] - if not region_resources: - continue - - # Create manager instance for this region - manager = manager_class(account_id, region, credentials) - - # Set account name in manager if supported - if hasattr(manager, "account_name"): - manager.account_name = account_name - - # Execute appropriate action - if action == "stop": - result = manager.stop(region_resources, dry_run) - stats[f"{resource_type}_stopped"] = stats.get( - f"{resource_type}_stopped", 0 - ) + self._safe_get_result(result, "success") - stats[f"{resource_type}_errors"] = stats.get( - f"{resource_type}_errors", 0 - ) + self._safe_get_result(result, "errors") - stats[f"{resource_type}_skipped"] = stats.get( - f"{resource_type}_skipped", 0 - ) + self._safe_get_result(result, "skipped") - else: # action == "start" - result = manager.start(region_resources, dry_run) - stats[f"{resource_type}_started"] = stats.get( - f"{resource_type}_started", 0 - ) + self._safe_get_result(result, "success") - stats[f"{resource_type}_errors"] = stats.get( - f"{resource_type}_errors", 0 - ) + self._safe_get_result(result, "errors") - stats[f"{resource_type}_skipped"] = stats.get( - f"{resource_type}_skipped", 0 - ) + self._safe_get_result(result, "skipped") + # Process resources by region + self._process_resources_by_region( + resource_type, manager_class, resource_list, + credentials, account_id, account_name, + action, dry_run, stats + ) + + def _process_resources_by_region( + self, + resource_type: str, + manager_class: Any, + resource_list: List[Dict[str, Any]], + credentials: Dict[str, str], + account_id: str, + account_name: str, + action: str, + dry_run: bool, + stats: Dict[str, Any] + ) -> None: + """Process resources by region using the appropriate manager.""" + for region in self.region_cache.get(account_id, []): + # Filter resources for this region + region_resources = [r for r in resource_list if r.get("region") == region] + if not region_resources: + continue + + # Create manager instance + manager = manager_class(account_id, region, credentials) + + # Set account name if supported + if hasattr(manager, "account_name"): + manager.account_name = account_name + + # Execute the action + if action == "stop": + result = manager.stop(region_resources, dry_run) + self._update_action_stats(stats, resource_type, "stopped", result) + else: # action == "start" + result = manager.start(region_resources, dry_run) + self._update_action_stats(stats, resource_type, "started", result) + + def _update_action_stats( + self, stats: Dict[str, Any], resource_type: str, action: str, + result: Optional[Dict[str, int]] + ) -> None: + """Update statistics with action results.""" + stats[f"{resource_type}_{action}"] = stats.get(f"{resource_type}_{action}", 0) + self._safe_get_result(result, "success") + stats[f"{resource_type}_errors"] = stats.get(f"{resource_type}_errors", 0) + self._safe_get_result(result, "errors") + stats[f"{resource_type}_skipped"] = stats.get(f"{resource_type}_skipped", 0) + self._safe_get_result(result, "skipped") def _safe_get_result(self, result: Optional[Dict[str, int]], key: str) -> int: """Safely get a result value from a dictionary.""" @@ -620,14 +590,10 @@ def _merge_stats( for key, value in account_stats.items(): if isinstance(value, (int, float)) and key in main_stats: main_stats[key] += value - # Special handling for ECR stats that might - # not be initialized in main_stats yet + # Special handling for ECR stats elif key in [ - "ecr_found", - "ecr_old_images", - "ecr_deleted", - "ecr_skipped", - "ecr_errors", + "ecr_found", "ecr_old_images", "ecr_deleted", + "ecr_skipped", "ecr_errors" ] and isinstance(value, (int, float)): main_stats[key] = main_stats.get(key, 0) + value @@ -636,9 +602,7 @@ def _merge_stats( # 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 - ) + main_stats["rds_engines"][engine] = main_stats["rds_engines"].get(engine, 0) + count # Merge regions_by_account for account_id, regions in account_stats.get("regions_by_account", {}).items(): @@ -647,16 +611,10 @@ def _merge_stats( 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 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)}" - ) + for account_id, account_regions in stats.get("regions_by_account", {}).items(): + logger.debug(f"Account {account_id} regions: {', '.join(account_regions)}") print_resource_summary(stats, action, stats.get("dry_run", False)) @@ -681,46 +639,71 @@ def discover_resources( Returns: Dictionary of resources by type """ + # Initialize params with defaults exclude_accounts = exclude_accounts or [] exclude_resources = exclude_resources or [] exclude_regions = exclude_regions or [] regions = regions or [] - # Log which resource types are being excluded for debugging + # Log resource types being discovered if exclude_resources: resource_types_to_discover = [ rt for rt in self.RESOURCE_TYPES if rt not in exclude_resources ] logger.info(f"Discovering only: {', '.join(resource_types_to_discover)}") - # Get account list + # Get and filter accounts + accounts = self._get_filtered_accounts(exclude_accounts, specific_account) + if not accounts: + logger.error("No accounts available for resource discovery") + return {} + + # Initialize consolidated resources dictionary + all_resources = self._initialize_resource_dict() + + # Process each account + for account in accounts: + account_resources = self._discover_account_resources( + account, regions, exclude_regions, exclude_resources + ) + + # Merge resources into the consolidated dictionary + for resource_type, resources in account_resources.items(): + all_resources.setdefault(resource_type, []).extend(resources) + + # Log summary of discovered resources + for resource_type, resources in all_resources.items(): + if resources: + logger.info(f"Total {resource_type}: {len(resources)}") + + return all_resources + + def _get_filtered_accounts( + self, exclude_accounts: List[str], specific_account: Optional[str] + ) -> List[Dict[str, str]]: + """Get filtered list of accounts based on exclusions and specific account.""" accounts = self._get_accounts() - # Filter accounts if specific_account is provided + # Filter for specific account if provided if specific_account: accounts = [ acc for acc in accounts if acc.get("account_id") == specific_account ] if not accounts: - logger.error( - f"Specified account {specific_account} not found or not accessible" - ) - return {} + logger.error(f"Specified account {specific_account} not found or not accessible") + return [] - # Pre-filter accounts that we can authenticate with + # Pre-filter accounts we can authenticate with accounts = self._pre_validate_accounts(accounts) # Apply account exclusions - accounts = [ + return [ acc for acc in accounts if acc.get("account_id") not in exclude_accounts ] - - if not accounts: - logger.error("No accounts available for resource discovery") - return {} - - # Initialize consolidated resources dictionary - all_resources = { + + def _initialize_resource_dict(self) -> Dict[str, List[Dict[str, Any]]]: + """Initialize empty resource dictionary with all resource types.""" + return { "ec2_instances": [], "rds_instances": [], "eks_clusters": [], @@ -728,63 +711,49 @@ def discover_resources( "ecr_images": [], "ecr_old_images": [], } + + def _discover_account_resources( + self, + account: Dict[str, str], + regions: List[str], + exclude_regions: List[str], + exclude_resources: List[str] + ) -> Dict[str, List[Dict[str, Any]]]: + """Discover resources for a single account.""" + account_id = account.get("account_id") + account_name = ensure_valid_account_name(account_id, account.get("account_name")) + account["account_name"] = account_name # Update in the source - # Process each account - for account in accounts: - account_id = account.get("account_id") - # Ensure account_name is valid and consistent - account_name = ensure_valid_account_name( - account_id, account.get("account_name") - ) - account["account_name"] = account_name # Update in the source - - logger.info( - f"Discovering resources in account {account_name} ({account_id})" - ) - - 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}" - ) - continue - - # Determine regions for this account - account_regions = self._get_account_regions( - account_id, credentials, regions, exclude_regions - ) - - if not account_regions: - logger.warning(f"No valid regions for account {account_id}") - continue - - # Get resources with special handling for EMR - emr_states = self._get_emr_states(exclude_resources) - - # Get all resources for the account - account_resources = get_account_resources( - credentials=credentials, - regions=account_regions, - exclude_resources=exclude_resources, - account_id=account_id, - account_name=account_name, - emr_cluster_states=emr_states, - ) + logger.info(f"Discovering resources in account {account_name} ({account_id})") - # Merge resources into the consolidated dictionary - for resource_type, resources in account_resources.items(): - all_resources.setdefault(resource_type, []).extend(resources) + try: + # 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 {} - except Exception as e: - logger.error( - f"Error discovering resources for account {account_id}: {str(e)}" - ) + # Determine regions + account_regions = self._get_account_regions( + account_id, credentials, regions, exclude_regions + ) + if not account_regions: + logger.warning(f"No valid regions for account {account_id}") + return {} - # Log summary of discovered resources - for resource_type, resources in all_resources.items(): - if resources: - logger.info(f"Total {resource_type}: {len(resources)}") + # Get EMR states + emr_states = self._get_emr_states(exclude_resources) - return all_resources + # Get resources + return get_account_resources( + credentials=credentials, + regions=account_regions, + exclude_resources=exclude_resources, + account_id=account_id, + account_name=account_name, + emr_cluster_states=emr_states, + ) + + except Exception as e: + logger.error(f"Error discovering resources for account {account_id}: {str(e)}") + return {} 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 235f416d..f8b6e6b3 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 @@ -8,7 +8,7 @@ import boto3 from aws_resource_management.utils import ( DEFAULT_REGIONS, - create_boto3_client, + create_client, detect_partition_from_credentials, ensure_valid_account_name, filter_regions_for_partition, @@ -65,15 +65,6 @@ def __init__( aws_session_token=credentials.get("aws_session_token"), ) self.partition = detect_partition_from_credentials(credentials) - - # Get account alias - try: - iam_client = self.session.client("iam") - response = iam_client.list_account_aliases() - self.account_alias = response.get("AccountAliases", [""])[0] if response.get("AccountAliases") else None - except Exception as e: - logger.debug(f"Error getting account alias: {str(e)}") - self.account_alias = None # Log with account name and ID logger.info( @@ -163,7 +154,7 @@ def _discover_ec2( List of EC2 instance dictionaries """ # Create EC2 client using central utility - ec2_client = create_boto3_client("ec2", self.credentials, region) + ec2_client = create_client("ec2", self.credentials, region) # Build filters filters = [] @@ -210,7 +201,6 @@ def process_ec2_page(page: Dict[str, Any]) -> List[Dict[str, Any]]: for instance in instances: instance["accountId"] = self.account_id instance["accountName"] = self.account_name - instance["accountAlias"] = self.account_alias return instances @@ -231,7 +221,7 @@ def _discover_rds( List of RDS instance dictionaries """ # Create RDS client using central utility - rds_client = create_boto3_client("rds", self.credentials, region) + rds_client = create_client("rds", self.credentials, region) try: # Get instances @@ -290,7 +280,6 @@ def _discover_rds( for instance in instances: instance["accountId"] = self.account_id instance["accountName"] = self.account_name - instance["accountAlias"] = self.account_alias return instances @@ -311,7 +300,7 @@ def _discover_emr( List of EMR cluster dictionaries """ # Create EMR client using central utility - emr_client = create_boto3_client("emr", self.credentials, region) + emr_client = create_client("emr", self.credentials, region) try: # Build filter for cluster states @@ -402,7 +391,6 @@ def _discover_emr( for cluster in detailed_clusters: cluster["accountId"] = self.account_id cluster["accountName"] = self.account_name - cluster["accountAlias"] = self.account_alias return detailed_clusters @@ -423,7 +411,7 @@ def _discover_eks( List of EKS cluster dictionaries """ # Create EKS client using central utility - eks_client = create_boto3_client("eks", self.credentials, region) + eks_client = create_client("eks", self.credentials, region) try: # Get clusters @@ -478,7 +466,6 @@ def _discover_eks( for cluster in clusters: cluster["accountId"] = self.account_id cluster["accountName"] = self.account_name - cluster["accountAlias"] = self.account_alias return clusters @@ -499,7 +486,7 @@ def _discover_ecr( List of ECR image dictionaries """ # Create ECR client using central utility - ecr_client = create_boto3_client("ecr", self.credentials, region) + ecr_client = create_client("ecr", self.credentials, region) try: # Get all repositories @@ -578,7 +565,6 @@ def _discover_ecr( # Add account info directly "accountId": self.account_id, "accountName": self.account_name, - "accountAlias": self.account_alias, } all_images.append(image_dict) @@ -632,14 +618,14 @@ def discover_ecr_images( try: for region in regions: try: - # Create ECR manager with provided credentials + # Create ECR manager with corrected parameter signature ecr_manager = ECRManager( - account_id=account_id, region=region, credentials=credentials + account_id=account_id, + region=region, + credentials=credentials, + account_name=account_name, ) - # Set account name - ecr_manager.account_name = account_name - # Discover resources in this region region_images = ecr_manager.discover_resources([region]) all_ecr_images.extend(region_images) @@ -855,9 +841,6 @@ def _get_account_resources_impl( for resource in resource_list: resource["accountId"] = account_id resource["accountName"] = account_name - # Add account alias if it was successfully retrieved - if hasattr(discovery, "account_alias") and discovery.account_alias: - resource["accountAlias"] = discovery.account_alias # After all discoveries are complete, log the counts for each resource type # to ensure they're properly tracked in the logs 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 3e9d9a87..ab5ac906 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 @@ -7,11 +7,9 @@ from typing import Any, Callable, Dict, List, Optional from aws_resource_management.utils import ( - get_account_alias, get_config, get_session_for_account, normalize_resource_state, - paginate_aws_response, ) from botocore.config import Config from botocore.exceptions import ClientError, NoRegionError @@ -32,28 +30,38 @@ class ResourceManager: def __init__( self, account_id: str, - region: str, + region: str = None, credentials: Optional[Dict[str, str]] = None, account_name: Optional[str] = None, - ): + role_name: Optional[str] = None, + regions: Optional[List[str]] = None, + resource_type: Optional[str] = None, + ) -> None: """ - Initialize resource manager. + Initialize the ResourceManager. Args: - account_id: AWS account ID - region: AWS region - credentials: Optional AWS credentials - account_name: Optional AWS account name + account_id: AWS account ID. + region: AWS region (optional). + credentials: AWS credentials dictionary (optional). + account_name: AWS account name (optional). + role_name: AWS role name (optional). + regions: List of AWS regions to manage (optional). + resource_type: Type of resource being managed (optional). """ self.account_id = account_id self.region = region + self.account_name = account_name or account_id + self.role_name = role_name self.credentials = credentials - # Make sure account_name is never None to avoid 'Unknown' in logs - self.account_name = account_name if account_name else account_id - self.account_alias = None # Will be populated when needed - self.session = None - self._clients = {} # Cache for boto3 clients + self.regions = regions or [region] if region else [] + self.resource_type = resource_type or self._RESOURCE_TYPE + self.logger = logging.getLogger(__name__) self.config = get_config() + + # Initialize client cache + self._clients = {} + self.session = None def get_timestamp(self) -> str: """ @@ -63,17 +71,6 @@ def get_timestamp(self) -> str: ISO formatted timestamp string """ return datetime.utcnow().isoformat() - - def get_account_alias(self) -> Optional[str]: - """ - Get the AWS account alias. - - Returns: - Account alias or None if not found - """ - if self.account_alias is None: - self.account_alias = get_account_alias(self.account_id, self.credentials) - return self.account_alias def get_client(self, service_name: str) -> Any: """ @@ -115,21 +112,18 @@ def get_client(self, service_name: str) -> Any: except NoRegionError: logger.error( - f"No region specified for {service_name} client \ - and no default region found" + f"No region specified for {service_name} client and no default region found" ) raise except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "Unknown") logger.error( - f"Error creating {service_name} client in \ - {self.region}: {error_code} - {str(e)}" + f"Error creating {service_name} client in {self.region}: {error_code} - {str(e)}" ) raise except Exception as e: logger.error( - f"Error creating {service_name} client \ - in {self.region}: {str(e)}" + f"Error creating {service_name} client in {self.region}: {str(e)}" ) raise @@ -147,8 +141,7 @@ def _handle_api_error(self, operation: str, error: Exception) -> None: if error_code == "AuthFailure": logger.error( - f"Authentication failure during {operation} in \ - account {self.account_id} " + f"Authentication failure during {operation} in account {self.account_id} " f"region {self.region}. Check IAM permissions." ) elif error_code == "AccessDenied": @@ -209,55 +202,6 @@ def discover_resources(self, regions: List[str]) -> List[Dict[str, Any]]: """ raise NotImplementedError("Subclasses must implement discover_resources method") - def paginate_boto3( - self, client: Any, operation: str, result_key: str, **kwargs - ) -> List[Dict[str, Any]]: - """ - Paginate through AWS API responses. - - Args: - client: Boto3 client - operation: Operation name (e.g., 'describe_instances') - result_key: Key in the response that contains the results - **kwargs: Additional arguments for the operation - - Returns: - Combined list of results from all pages - """ - # Delegate to the centralized utility function - return paginate_aws_response(client, operation, result_key, **kwargs) - - 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 (e.g., 'ec2', 'rds') - region: AWS region - - Returns: - Boto3 client for the specified service in the specified region - """ - try: - # Get an account-specific session with proper role assumption - if not self.session: - self.session = get_session_for_account(self.account_id) - - # Create config with retry settings - boto_config = Config( - region_name=region, - retries={"max_attempts": 3, "mode": "standard"}, - ) - - # Create client with proper configuration - client = self.session.client( - service_name, region_name=region, config=boto_config - ) - return client - except Exception as e: - logger.error(f"Error creating {service_name} client in {region}: {str(e)}") - return None - def log_action( self, resource_id: str, @@ -286,14 +230,13 @@ def log_action( schedule_str = f" [Schedule: {existing_schedule}]" if existing_schedule else "" logger.info( - f"{dry_run_prefix}{action.capitalize()} {self._RESOURCE_TYPE} \ - {resource_id} {name_str} " + f"{dry_run_prefix}{action.capitalize()} {self._RESOURCE_TYPE} {resource_id} {name_str} " f"in {region}{details_str}{schedule_str}" ) # Also log to CSV if tracking is enabled try: - from aws_resource_management.utils.logging_utils import log_action_to_csv + from aws_resource_management.utils import log_action_to_csv log_action_to_csv( account_id=self.account_id, @@ -312,24 +255,6 @@ def log_action( # CSV logging not available pass - def should_exclude(self, resource: Dict[str, Any]) -> bool: - """ - Check if a resource should be excluded based on tags. - - Args: - resource: Resource dictionary with tags - - Returns: - True if the resource should be excluded, False otherwise - """ - tags = resource.get("tags", {}) - exclusion_tag = self.config.get("exclusion_tag") - - if exclusion_tag and exclusion_tag in tags: - return True - - return False - def run_with_retries( self, func: Callable, @@ -386,9 +311,7 @@ def process_resources_in_batches( for i in range(0, len(resources), batch_size): batch = resources[i : i + batch_size] logger.info( - f"{description} batch \ - {i//batch_size + 1}/{(len(resources)-1)//batch_size + 1} \ - ({len(batch)} resources)" + f"{description} batch {i//batch_size + 1}/{(len(resources)-1)//batch_size + 1} ({len(batch)} resources)" ) try: @@ -419,9 +342,6 @@ def normalize_resources(resources: List[Dict[str, Any]]) -> None: if "accountName" not in resource and hasattr(resource, "accountName"): resource["accountName"] = resource.accountName - - if "accountAlias" not in resource and hasattr(resource, "accountAlias"): - resource["accountAlias"] = resource.accountAlias # If still no account name, use account ID as fallback if "accountName" not in resource and "accountId" in resource: 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 d98e0fb1..53f5fa42 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 @@ -8,7 +8,7 @@ import boto3 from aws_resource_management.managers.base import ResourceManager from aws_resource_management.utils import ( - create_boto3_client, + create_client, get_config, parse_iso_datetime, setup_logging, @@ -35,6 +35,7 @@ def __init__( account_id: str, region: str, credentials: Optional[Dict[str, Any]] = None, + account_name: Optional[str] = None, ): """ Initialize ECR manager. @@ -43,10 +44,15 @@ def __init__( account_id: AWS account ID region: AWS region credentials: Optional AWS credentials dictionary + account_name: Optional AWS account name """ - super().__init__(account_id, region) - self.account_name = None # Will be set by the caller if available - self.credentials = credentials + super().__init__( + account_id=account_id, + region=region, + credentials=credentials, + account_name=account_name, + resource_type="ecr" + ) self.age_threshold = config.get( "ecr_age_threshold_days", DEFAULT_AGE_THRESHOLD_DAYS ) @@ -62,7 +68,7 @@ def get_ecr_client(self, region: str = None) -> boto3.client: Boto3 ECR client """ region = region or self.region - return create_boto3_client("ecr", self.credentials, region) + return create_client("ecr", self.credentials, region) def discover_resources(self, regions: List[str]) -> List[Dict[str, Any]]: """ diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/region_manager.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/region_manager.py new file mode 100644 index 00000000..c68311d6 --- /dev/null +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/region_manager.py @@ -0,0 +1,140 @@ +"""Region management utilities for AWS Resource Management.""" + +from typing import Dict, List, Optional + +from aws_resource_management.utils import ( + DEFAULT_PARTITION, + detect_partition_from_credentials, + filter_regions_for_partition, + get_default_regions_for_partition, + get_enabled_regions, + logger, +) + + +class RegionManager: + """Manages region discovery, filtering, and caching.""" + + def __init__( + self, + partition: Optional[str] = None, + discover_regions: bool = False, + ) -> None: + """Initialize the region manager. + + Args: + partition: AWS partition (aws, aws-us-gov, aws-cn) + discover_regions: Whether to discover regions + """ + self.partition = partition + self.discover_regions = discover_regions + self.region_cache = {} + + 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 region discovery is enabled, try to discover regions + if self.discover_regions: + return self._discover_regions( + account_id, credentials, exclude_regions + ) + + # If provided regions, filter them + elif provided_regions: + return self._filter_provided_regions( + provided_regions, exclude_regions, credentials + ) + + # Otherwise use default regions + else: + default_regions = self._get_default_regions(credentials) + logger.info( + f"No regions specified, using defaults: {', '.join(default_regions)}" + ) + return default_regions + + def _discover_regions( + self, account_id: str, credentials: Dict[str, str], exclude_regions: List[str] + ) -> List[str]: + """Discover enabled regions for an account.""" + try: + logger.info(f"Discovering enabled regions for account {account_id}") + discovered_regions = get_enabled_regions(credentials, self.partition) + + # Filter discovered regions for the partition + if self.partition: + discovered_regions = filter_regions_for_partition( + discovered_regions, self.partition + ) + + # Apply exclusions + 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)}" + ) + return [] + + def _filter_provided_regions( + self, provided_regions: List[str], exclude_regions: List[str], + credentials: Dict[str, str] + ) -> List[str]: + """Filter provided regions based on partition and exclusions.""" + if self.partition: + # Use centralized filter + valid_regions = filter_regions_for_partition( + provided_regions, self.partition + ) + filtered_regions = [ + r for r in valid_regions if r not in exclude_regions + ] + else: + 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 or invalid, using defaults: " + f"{', '.join(filtered_regions)}" + ) + + return filtered_regions + + def _get_default_regions( + self, credentials: Optional[Dict[str, str]] = None + ) -> List[str]: + """Get default regions for the current partition.""" + # Try to detect partition from credentials if not explicitly set + 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}") + + # Use GovCloud AWS as the default partition + partition = partition or DEFAULT_PARTITION + return get_default_regions_for_partition(partition) + + def get_all_cached_regions(self) -> List[str]: + """Get all unique regions from the cache.""" + all_regions = set() + for regions in self.region_cache.values(): + all_regions.update(regions) + return sorted(list(all_regions)) diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/stats_manager.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/stats_manager.py new file mode 100644 index 00000000..132c0cdb --- /dev/null +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/stats_manager.py @@ -0,0 +1,76 @@ +"""Statistics management for AWS Resource Management.""" + +from typing import Any, Dict, List, Optional + +from aws_resource_management.reporting import initialize_stats +from aws_resource_management.utils import logger + + +class StatsManager: + """Manages statistics collection and reporting for resource operations.""" + + def __init__(self) -> None: + """Initialize the stats manager with empty statistics.""" + self.stats = initialize_stats() + + def update_resource_count(self, resource_type: str, count: int) -> None: + """Update the count of resources found.""" + stats_key = f"{resource_type}_found" + if stats_key not in self.stats: + self.stats[stats_key] = 0 + self.stats[stats_key] += count + + def update_action_stats(self, resource_type: str, action: str, + result: Optional[Dict[str, int]]) -> None: + """Update statistics with action results.""" + self.stats[f"{resource_type}_{action}"] = self.stats.get( + f"{resource_type}_{action}", 0) + self._safe_get_result(result, "success") + self.stats[f"{resource_type}_errors"] = self.stats.get( + f"{resource_type}_errors", 0) + self._safe_get_result(result, "errors") + self.stats[f"{resource_type}_skipped"] = self.stats.get( + f"{resource_type}_skipped", 0) + self._safe_get_result(result, "skipped") + + def _safe_get_result(self, result: Optional[Dict[str, int]], key: str) -> int: + """Safely get a result value from a dictionary.""" + if result is None: + return 0 + return result.get(key, 0) + + def merge_account_stats(self, 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 self.stats: + self.stats[key] += value + # Special handling for ECR stats + elif key in [ + "ecr_found", "ecr_old_images", "ecr_deleted", + "ecr_skipped", "ecr_errors" + ] and isinstance(value, (int, float)): + self.stats[key] = self.stats.get(key, 0) + value + + # Merge errors list + self.stats["errors"].extend(account_stats.get("errors", [])) + + # Merge RDS engines + for engine, count in account_stats.get("rds_engines", {}).items(): + self.stats["rds_engines"][engine] = self.stats["rds_engines"].get(engine, 0) + count + + # Merge regions_by_account + for account_id, regions in account_stats.get("regions_by_account", {}).items(): + self.stats["regions_by_account"][account_id] = regions + + # Update processed accounts count + self.stats["accounts_processed"] += 1 + + def record_account_skip(self) -> None: + """Record that an account was skipped.""" + self.stats["accounts_skipped"] += 1 + + def add_error(self, error_message: str) -> None: + """Add an error message to the stats.""" + self.stats["errors"].append(error_message) + + def get_stats(self) -> Dict[str, Any]: + """Get the current statistics.""" + return self.stats diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/__init__.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/__init__.py index 90198f18..e1b23a27 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/__init__.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/__init__.py @@ -14,9 +14,8 @@ safe_api_call, ) from aws_resource_management.utils.aws_core_utils import ( - create_boto3_client, + create_client, ensure_valid_account_name, - get_account_alias, get_account_list, get_credentials, get_organization_accounts, diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/api_utils.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/api_utils.py index 860e218d..fd266278 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/api_utils.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/api_utils.py @@ -168,7 +168,7 @@ def normalize_resource_state(resource: Dict[str, Any]) -> None: def add_account_info( - resources: List[Dict[str, Any]], account_id: str, account_name: Optional[str] = None, account_alias: Optional[str] = None + resources: List[Dict[str, Any]], account_id: str, account_name: Optional[str] = None ) -> None: """ Add account information to resources. @@ -177,14 +177,11 @@ def add_account_info( resources: List of resource dictionaries account_id: AWS account ID account_name: AWS account name - account_alias: AWS account alias """ for resource in resources: resource["accountId"] = account_id if account_name: resource["accountName"] = account_name - if account_alias: - resource["accountAlias"] = account_alias def create_boto3_client( diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/aws_core_utils.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/aws_core_utils.py index 3cd384e6..56f287c4 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/aws_core_utils.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/aws_core_utils.py @@ -493,55 +493,6 @@ def ensure_valid_account_name( return account_name -def get_account_alias(account_id: str, credentials: Optional[Dict[str, Any]] = None) -> Optional[str]: - """ - Get the AWS account alias with caching. - - Args: - account_id: AWS account ID - credentials: Optional credentials dictionary - - Returns: - Account alias or None if not found - """ - cache_key = f"alias:{account_id}" - - # 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]["alias"] - - try: - # Get credentials if not provided - if not credentials: - credentials = get_credentials(account_id) - if not credentials: - logger.debug(f"Could not get credentials for account {account_id}") - return None - - # Create IAM client - iam_client = create_boto3_client_from_creds("iam", credentials) - - # Get account alias - response = iam_client.list_account_aliases() - aliases = response.get("AccountAliases", []) - - # Store the first alias or None - alias = aliases[0] if aliases else None - - # Cache the result - with _session_cache_lock: - _session_cache[cache_key] = {"alias": alias, "timestamp": time.time()} - - return alias - except Exception as e: - logger.debug(f"Error getting account alias for {account_id}: {str(e)}") - return None - - def is_pending_subscription(credentials: Dict[str, str], region: str) -> bool: """Check if there's a pending marketplace subscription.""" @@ -569,28 +520,45 @@ def check_subscription_status(): return check_subscription_status() -def create_boto3_client_from_creds( - service: str, credentials: Dict[str, str], region: Optional[str] = None -) -> boto3.client: +def create_client( + service_name: str, + credentials: Optional[Dict[str, str]] = None, + region: Optional[str] = None, + profile_name: Optional[str] = None, +) -> Any: """ - Create a boto3 client from credentials. + Create a boto3 client for the given service. + + This function handles both direct credential passing and session/profile usage. Args: - service: AWS service name - credentials: Dictionary containing AWS credentials - region: AWS region name (optional) + service_name: The name of the AWS service (e.g., 'ec2', 'rds'). + credentials: A dictionary containing AWS credentials ('aws_access_key_id', 'aws_secret_access_key', 'aws_session_token'). + region: The AWS region to use. + profile_name: The name of the AWS profile to use. Returns: - Boto3 client for the specified service + A boto3 client object. """ - 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"), - ) - - -# Replace the lambda with a proper function (fixing E731) -create_boto3_client = create_boto3_client_from_creds + try: + if credentials: + # Use explicit credentials + client = boto3.client( + service_name, + 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"), + region_name=region, + ) + elif profile_name: + # Use a specific profile + session = boto3.Session(profile_name=profile_name) + client = session.client(service_name, region_name=region) + else: + # Use the default session + session = boto3.Session() + client = session.client(service_name, region_name=region) + return client + except Exception as e: + logger.error(f"Error creating {service_name} client: {e}") + raise diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/config_utils.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/config_utils.py index 00d1a22c..bb54efbb 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/config_utils.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/config_utils.py @@ -25,151 +25,161 @@ # Global config cache _config = None +class ConfigManager: + """Manages configuration for resource operations.""" + def __init__(self, profile_name=None, use_profiles=False, + discover_regions=False, partition=None): + self.config = get_config() + self.profile_name = profile_name + self.use_profiles = use_profiles + self.discover_regions = discover_regions + self.partition = partition + + + def get_config() -> Dict[str, Any]: + """ + Get configuration with defaults merged with user settings. + + Returns: + Configuration dictionary + """ + global _config -def get_config() -> Dict[str, Any]: - """ - Get configuration with defaults merged with user settings. - - Returns: - Configuration dictionary - """ - global _config - - # Return cached config if available - if _config is not None: - return _config - - # Start with default config - config = DEFAULT_CONFIG.copy() - - # Look for config files in multiple locations - config_paths = [ - os.path.join( - os.path.dirname(os.path.dirname(os.path.abspath(__file__))), - "..", - "config.yaml", - ), - os.path.expanduser("~/.aws-resource-management/config.yaml"), - ] - - # Try to load and merge configs from files - for path in config_paths: - try: - if os.path.exists(path): - with open(path, "r") as f: - user_config = yaml.safe_load(f) - if user_config and isinstance(user_config, dict): - # Merge with default config - _deep_merge(config, user_config) - except Exception: - # Ignore errors reading config files - pass + # Return cached config if available + if _config is not None: + return _config + + # Start with default config + config = DEFAULT_CONFIG.copy() + + # Look for config files in multiple locations + config_paths = [ + os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), + "..", + "config.yaml", + ), + os.path.expanduser("~/.aws-resource-management/config.yaml"), + ] + + # Try to load and merge configs from files + for path in config_paths: + try: + if os.path.exists(path): + with open(path, "r") as f: + user_config = yaml.safe_load(f) + if user_config and isinstance(user_config, dict): + # Merge with default config + _deep_merge(config, user_config) + except Exception: + # Ignore errors reading config files + pass + + # Cache the config + _config = config + return config - # Cache the config - _config = config - return config + def _deep_merge(base: Dict, update: Dict) -> Dict: + """ + Recursively merge dictionaries. -def _deep_merge(base: Dict, update: Dict) -> Dict: - """ - Recursively merge dictionaries. + Args: + base: Base dictionary to update + update: Dictionary with values to merge - Args: - base: Base dictionary to update - update: Dictionary with values to merge + Returns: + Updated base dictionary + """ + for key, value in update.items(): + if isinstance(value, dict) and key in base and isinstance(base[key], dict): + _deep_merge(base[key], value) + else: + base[key] = value + return base - Returns: - Updated base dictionary - """ - for key, value in update.items(): - if isinstance(value, dict) and key in base and isinstance(base[key], dict): - _deep_merge(base[key], value) - else: - base[key] = value - return base + def get_config_value(key: str, default: Any = None) -> Any: + """ + Get a specific configuration value. -def get_config_value(key: str, default: Any = None) -> Any: - """ - Get a specific configuration value. + Args: + key: Configuration key to retrieve + default: Default value if key not found - Args: - key: Configuration key to retrieve - default: Default value if key not found + Returns: + Configuration value or default + """ + config = get_config() - Returns: - Configuration value or default - """ - config = get_config() + # Handle nested keys with dot notation (e.g., "aws.region") + if "." in key: + parts = key.split(".") + current = config + for part in parts: + if not isinstance(current, dict) or part not in current: + return default + current = current[part] + return current - # Handle nested keys with dot notation (e.g., "aws.region") - if "." in key: - parts = key.split(".") - current = config - for part in parts: - if not isinstance(current, dict) or part not in current: - return default - current = current[part] - return current + return config.get(key, default) - return config.get(key, default) + def save_config(config: Dict[str, Any], config_path: Optional[str] = None) -> bool: + """ + Save configuration to a file. -def save_config(config: Dict[str, Any], config_path: Optional[str] = None) -> bool: - """ - Save configuration to a file. + Args: + config: Configuration dictionary to save + config_path: Path to save config to (default: user config path) - Args: - config: Configuration dictionary to save - config_path: Path to save config to (default: user config path) + Returns: + True if saved successfully, False otherwise + """ + if config_path is None: + config_path = os.path.expanduser("~/.aws-resource-management/config.yaml") - Returns: - True if saved successfully, False otherwise - """ - if config_path is None: - config_path = os.path.expanduser("~/.aws-resource-management/config.yaml") + try: + # Ensure directory exists + os.makedirs(os.path.dirname(config_path), exist_ok=True) - try: - # Ensure directory exists - os.makedirs(os.path.dirname(config_path), exist_ok=True) + # Write config + with open(config_path, "w") as f: + yaml.safe_dump(config, f, default_flow_style=False) - # Write config - with open(config_path, "w") as f: - yaml.safe_dump(config, f, default_flow_style=False) + # Update cache + global _config + _config = config - # Update cache - global _config - _config = config + return True + except Exception: + return False + + + def update_config(key: str, value: Any, config_path: Optional[str] = None) -> bool: + """ + Update a specific configuration value and save. + + Args: + key: Configuration key to update + value: New value + config_path: Path to save config to (default: user config path) + + Returns: + True if updated successfully, False otherwise + """ + config = get_config() + + # Handle nested keys with dot notation + if "." in key: + parts = key.split(".") + current = config + for i, part in enumerate(parts[:-1]): + if part not in current or not isinstance(current[part], dict): + current[part] = {} + current = current[part] + current[parts[-1]] = value + else: + config[key] = value - return True - except Exception: - return False - - -def update_config(key: str, value: Any, config_path: Optional[str] = None) -> bool: - """ - Update a specific configuration value and save. - - Args: - key: Configuration key to update - value: New value - config_path: Path to save config to (default: user config path) - - Returns: - True if updated successfully, False otherwise - """ - config = get_config() - - # Handle nested keys with dot notation - if "." in key: - parts = key.split(".") - current = config - for i, part in enumerate(parts[:-1]): - if part not in current or not isinstance(current[part], dict): - current[part] = {} - current = current[part] - current[parts[-1]] = value - else: - config[key] = value - - return save_config(config, config_path) + return save_config(config, config_path) diff --git a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/region_utils.py b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/region_utils.py index d4c02d73..73bcdca9 100644 --- a/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/region_utils.py +++ b/local-app/python-tools/gfl-resource-actions/aws_resource_management/utils/region_utils.py @@ -74,352 +74,357 @@ # Logger for this module logger = logging.getLogger(__name__) +class RegionManager: + """Handles AWS region management.""" + def __init__(self, partition=None): + self.partition = partition + self.region_cache = {} -def get_partition_for_region(region_name: str) -> str: - """ - Determine AWS partition based on region name. + def get_partition_for_region(region_name: str) -> str: + """ + Determine AWS partition based on region name. - Args: - region_name: AWS region name + Args: + region_name: AWS region name - Returns: - AWS partition name (e.g., 'aws', 'aws-us-gov', 'aws-cn') - """ - # Check common region prefixes first - for prefix, partition in REGION_PREFIX_TO_PARTITION.items(): - if region_name.startswith(prefix): - return partition + Returns: + AWS partition name (e.g., 'aws', 'aws-us-gov', 'aws-cn') + """ + # Check common region prefixes first + for prefix, partition in REGION_PREFIX_TO_PARTITION.items(): + if region_name.startswith(prefix): + return partition - # Default to standard AWS partition - return "aws" + # Default to standard AWS partition + return "aws" -def get_regions_for_partition(partition: str) -> List[str]: - """ - Get the list of regions for a specific partition. + def get_regions_for_partition(partition: str) -> List[str]: + """ + Get the list of regions for a specific partition. - Args: - partition: AWS partition name + Args: + partition: AWS partition name - Returns: - List of region names in the partition - """ - return PARTITION_REGIONS.get(partition, PARTITION_REGIONS["aws"]) + Returns: + List of region names in the partition + """ + return PARTITION_REGIONS.get(partition, PARTITION_REGIONS["aws"]) -def get_default_regions_for_partition(partition: str) -> List[str]: - """ - Get the default regions commonly used for a partition. + def get_default_regions_for_partition(partition: str) -> List[str]: + """ + Get the default regions commonly used for a partition. - Args: - partition: AWS partition name + Args: + partition: AWS partition name - Returns: - List of default region names - """ - return DEFAULT_REGIONS.get(partition, DEFAULT_REGIONS["aws"]) + Returns: + List of default region names + """ + return DEFAULT_REGIONS.get(partition, DEFAULT_REGIONS["aws"]) -def is_region_in_partition(region: str, partition: str) -> bool: - """ - Check if a region belongs to the specified partition. + def is_region_in_partition(region: str, partition: str) -> bool: + """ + Check if a region belongs to the specified partition. - Args: - region: AWS region name - partition: AWS partition name + Args: + region: AWS region name + partition: AWS partition name - Returns: - True if the region is in the partition, False otherwise - """ - return get_partition_for_region(region) == partition + Returns: + True if the region is in the partition, False otherwise + """ + return get_partition_for_region(region) == partition -def is_valid_region(region: str) -> bool: - """ - Check if a region is known to be valid without making an API call. + def is_valid_region(region: str) -> bool: + """ + Check if a region is known to be valid without making an API call. - Args: - region: Region name to validate + Args: + region: Region name to validate + + Returns: + True if the region is valid, False otherwise + """ + if not isinstance(region, str): + return False + + # First check if it's in our known regions list + if region in ALL_KNOWN_REGIONS: + return True + + # Check if it follows expected naming patterns for newer regions + # Commercial regions follow patterns like us-east-1, eu-west-2 + if ( + region.startswith(("us-", "eu-", "ap-", "sa-", "ca-", "me-", "af-")) + and len(region.split("-")) >= 3 + ): + return True + + # GovCloud regions + if region.startswith("us-gov-") and len(region.split("-")) >= 3: + return True + + # China regions + if region.startswith("cn-") and len(region.split("-")) >= 3: + return True - Returns: - True if the region is valid, False otherwise - """ - if not isinstance(region, str): return False - # First check if it's in our known regions list - if region in ALL_KNOWN_REGIONS: - return True - - # Check if it follows expected naming patterns for newer regions - # Commercial regions follow patterns like us-east-1, eu-west-2 - if ( - region.startswith(("us-", "eu-", "ap-", "sa-", "ca-", "me-", "af-")) - and len(region.split("-")) >= 3 - ): - return True - - # GovCloud regions - if region.startswith("us-gov-") and len(region.split("-")) >= 3: - return True - - # China regions - if region.startswith("cn-") and len(region.split("-")) >= 3: - return True - - return False - - -def filter_regions_for_partition(regions: List[str], partition: str) -> List[str]: - """ - Filter regions to include only those belonging to the specified partition. - - Args: - regions: List of region names - partition: AWS partition name - - Returns: - List of valid region names within the partition - """ - # First validate the regions and filter out invalid ones - valid_regions = [r for r in regions if is_valid_region(r)] - - # Then filter for the specific partition - partition_regions = [ - r for r in valid_regions if is_region_in_partition(r, partition) - ] - - return partition_regions - - -def detect_partition_from_credentials(credentials: Dict[str, str]) -> str: - """ - Detect partition from credentials with minimal API calls. - - Args: - credentials: Dictionary containing AWS credentials - - Returns: - AWS partition name - """ - if not credentials: - return DEFAULT_PARTITION # 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( - indicator in access_key or indicator in session_token - for indicator in ["gov", "usgovcloud"] - ): - return "aws-us-gov" - - # Check for China indicators - if any( - indicator in access_key or indicator in session_token - for indicator in ["cn-", "china"] - ): - return "aws-cn" - - # Try one API call to most likely partition based on environment - try: - 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 + + def filter_regions_for_partition(regions: List[str], partition: str) -> List[str]: + """ + Filter regions to include only those belonging to the specified partition. + + Args: + regions: List of region names + partition: AWS partition name + + Returns: + List of valid region names within the partition + """ + # First validate the regions and filter out invalid ones + valid_regions = [r for r in regions if is_valid_region(r)] + + # Then filter for the specific partition + partition_regions = [ + r for r in valid_regions if is_region_in_partition(r, partition) + ] + + return partition_regions + + + def detect_partition_from_credentials(credentials: Dict[str, str]) -> str: + """ + Detect partition from credentials with minimal API calls. + + Args: + credentials: Dictionary containing AWS credentials + + Returns: + AWS partition name + """ + if not credentials: + return DEFAULT_PARTITION # 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( + indicator in access_key or indicator in session_token + for indicator in ["gov", "usgovcloud"] + ): + return "aws-us-gov" + + # Check for China indicators + if any( + indicator in access_key or indicator in session_token + for indicator in ["cn-", "china"] + ): + return "aws-cn" + + # Try one API call to most likely partition based on environment try: boto3.client( "sts", - region_name="us-east-1", + 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" + return "aws-us-gov" except Exception: - # Default to the environment's default partition - return DEFAULT_PARTITION - - -def get_all_regions(partition: Optional[str] = None) -> List[str]: - """ - Get all AWS regions, using cache to minimize API calls. - - Args: - partition: Optional partition to filter regions by - - Returns: - List of region names - """ - # 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"] - ] + # 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 the environment's default partition + return DEFAULT_PARTITION + + + def get_all_regions(partition: Optional[str] = None) -> List[str]: + """ + Get all AWS regions, using cache to minimize API calls. + + Args: + partition: Optional partition to filter regions by + + Returns: + List of region names + """ + # 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"] + ] + + # 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 isinstance(r, str) 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.warning(f"Error getting regions: {str(e)}") + return get_regions_for_partition(partition or DEFAULT_PARTITION) + + + def get_enabled_regions( + credentials: Dict[str, str], partition: Optional[str] = None + ) -> List[str]: + """ + Get enabled regions with minimized API calls using cache. + + Args: + credentials: Dictionary containing AWS credentials + partition: Optional partition to filter regions by + + Returns: + List of enabled region names + """ + # 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}" - # Cache the full list - _region_cache["all_regions"]["all"] = all_regions - _region_cache["timestamp"] = current_time + # Check cache + if ( + cache_key in _region_cache["enabled_regions"] + and time.time() - _region_cache["timestamp"] < CACHE_EXPIRY + ): + return _region_cache["enabled_regions"][cache_key] - # If partition specified, filter and cache that too - if partition: - filtered_regions = [ - r - for r in all_regions - if isinstance(r, str) and get_partition_for_region(r) == partition + # Make a single API call to describe_regions rather than checking each region + 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"), + ) + regions = [ + region["RegionName"] + for region in session.client( + "ec2", region_name="us-east-1" + ).describe_regions()["Regions"] ] - _region_cache["all_regions"][partition] = filtered_regions - return filtered_regions - - return all_regions - except Exception as e: - logger.warning(f"Error getting regions: {str(e)}") - return get_regions_for_partition(partition or DEFAULT_PARTITION) - - -def get_enabled_regions( - credentials: Dict[str, str], partition: Optional[str] = None -) -> List[str]: - """ - Get enabled regions with minimized API calls using cache. - - Args: - credentials: Dictionary containing AWS credentials - partition: Optional partition to filter regions by - - Returns: - List of enabled region names - """ - # 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}" - - # Check cache - if ( - cache_key in _region_cache["enabled_regions"] - and time.time() - _region_cache["timestamp"] < CACHE_EXPIRY - ): - return _region_cache["enabled_regions"][cache_key] - - # Make a single API call to describe_regions rather than checking each region - 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"), - ) - regions = [ - region["RegionName"] - for region in session.client( + + # Cache the result + _region_cache["enabled_regions"][cache_key] = regions + _region_cache["timestamp"] = time.time() + + return regions + except Exception: + # Fall back to default regions + return get_default_regions_for_partition(partition) + + + def list_enabled_regions( + session: boto3.Session, exclude_regions: Optional[List[str]] = None + ) -> List[str]: + """ + List all enabled regions for the account, respecting partition. + + Args: + session: Boto3 session + exclude_regions: List of regions to exclude + + Returns: + List of enabled region names + """ + if exclude_regions is None: + exclude_regions = [] + + try: + ec2 = session.client( "ec2", region_name="us-east-1" - ).describe_regions()["Regions"] - ] + ) # Use US East 1 for global operations + regions_response = ec2.describe_regions( + AllRegions=False + ) # Only get enabled regions + regions = [ + r["RegionName"] + for r in regions_response["Regions"] + if r["RegionName"] not in exclude_regions + ] + return regions + except Exception as e: + logger.warning(f"Could not retrieve regions, using defaults: {str(e)}") - # Cache the result - _region_cache["enabled_regions"][cache_key] = regions - _region_cache["timestamp"] = time.time() - - return regions - except Exception: - # Fall back to default regions - return get_default_regions_for_partition(partition) - - -def list_enabled_regions( - session: boto3.Session, exclude_regions: Optional[List[str]] = None -) -> List[str]: - """ - List all enabled regions for the account, respecting partition. - - Args: - session: Boto3 session - exclude_regions: List of regions to exclude - - Returns: - List of enabled region names - """ - if exclude_regions is None: - exclude_regions = [] - - try: - ec2 = session.client( - "ec2", region_name="us-east-1" - ) # Use US East 1 for global operations - regions_response = ec2.describe_regions( - AllRegions=False - ) # Only get enabled regions - regions = [ - r["RegionName"] - for r in regions_response["Regions"] - if r["RegionName"] not in exclude_regions - ] - return regions - except Exception as e: - logger.warning(f"Could not retrieve regions, using defaults: {str(e)}") - - # For GovCloud, provide sensible defaults - if session.region_name and session.region_name.startswith("us-gov-"): - return ["us-gov-east-1", "us-gov-west-1"] - - # Default to common commercial regions - return ["us-east-1", "us-east-2", "us-west-1", "us-west-2"] - - -def discover_account_regions(client: Any, include_disabled: bool = False) -> List[str]: - """ - Discover regions enabled for an account using EC2 client. - - Args: - client: EC2 client to use for region discovery - include_disabled: Whether to include disabled regions - - Returns: - List of region names - """ - try: - # Call describe_regions API - if include_disabled: - response = client.describe_regions(AllRegions=True) - else: - response = client.describe_regions() - - # Extract region names - regions = [region["RegionName"] for region in response.get("Regions", [])] - return regions - - except Exception as e: - logger.error(f"Error discovering account regions: {str(e)}") - return [] + # For GovCloud, provide sensible defaults + if session.region_name and session.region_name.startswith("us-gov-"): + return ["us-gov-east-1", "us-gov-west-1"] + + # Default to common commercial regions + return ["us-east-1", "us-east-2", "us-west-1", "us-west-2"] + + + def discover_account_regions(client: Any, include_disabled: bool = False) -> List[str]: + """ + Discover regions enabled for an account using EC2 client. + + Args: + client: EC2 client to use for region discovery + include_disabled: Whether to include disabled regions + + Returns: + List of region names + """ + try: + # Call describe_regions API + if include_disabled: + response = client.describe_regions(AllRegions=True) + else: + response = client.describe_regions() + + # Extract region names + regions = [region["RegionName"] for region in response.get("Regions", [])] + return regions + + except Exception as e: + logger.error(f"Error discovering account regions: {str(e)}") + return []