From f5fdbf1755ee46066fd6c51c9e585cd45ab7e673 Mon Sep 17 00:00:00 2001 From: "Matthew C. Morgan" Date: Thu, 24 Apr 2025 18:23:18 -0400 Subject: [PATCH] add account alias and imagelastPulledTime --- .../aws_resource_management/core.py | 12 ++- .../aws_resource_management/discovery.py | 17 ++++ .../aws_resource_management/managers/base.py | 16 ++++ .../aws_resource_management/managers/ecr.py | 2 + .../aws_resource_management/utils/__init__.py | 1 + .../utils/api_utils.py | 5 +- .../utils/aws_core_utils.py | 86 +++++++++++-------- 7 files changed, 100 insertions(+), 39 deletions(-) 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 c238b260..afacdf7a 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 @@ -45,7 +45,7 @@ def __init__( self.discover_regions = discover_regions self.partition = partition self.region_cache = {} - self.MAX_CONCURRENT_ACCOUNTS = 3 # Limit concurrent account processing + self.MAX_CONCURRENT_ACCOUNTS = 20 # Limit concurrent account processing # Get resource registry self.resource_registry = get_registry() @@ -213,6 +213,15 @@ 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 @@ -257,6 +266,7 @@ def _process_account_safely( 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, ) 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 ee42c594..235f416d 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 @@ -65,6 +65,15 @@ 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( @@ -201,6 +210,7 @@ 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 @@ -280,6 +290,7 @@ def _discover_rds( for instance in instances: instance["accountId"] = self.account_id instance["accountName"] = self.account_name + instance["accountAlias"] = self.account_alias return instances @@ -391,6 +402,7 @@ 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 @@ -466,6 +478,7 @@ def _discover_eks( for cluster in clusters: cluster["accountId"] = self.account_id cluster["accountName"] = self.account_name + cluster["accountAlias"] = self.account_alias return clusters @@ -565,6 +578,7 @@ def _discover_ecr( # Add account info directly "accountId": self.account_id, "accountName": self.account_name, + "accountAlias": self.account_alias, } all_images.append(image_dict) @@ -841,6 +855,9 @@ 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 e4ad45ca..3e9d9a87 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,6 +7,7 @@ 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, @@ -49,6 +50,7 @@ def __init__( 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.config = get_config() @@ -61,6 +63,17 @@ 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: """ @@ -406,6 +419,9 @@ 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 cc20d297..d98e0fb1 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 @@ -137,11 +137,13 @@ def _get_repository_images( "repositoryName": repo_name, "tags": image.get("imageTags", []), "imagePushedAt": image.get("imagePushedAt"), + "lastRecordedPullTime": image.get("lastRecordedPullTime"), "imageSizeInBytes": image.get("imageSizeInBytes", 0), "status": "AVAILABLE", # Add account info directly "accountId": self.account_id, "accountName": self.account_name or self.account_id, + "accountAlias": self.get_account_alias() or self.account_id, } images.append(image_dict) 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 09b5e9e7..90198f18 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 @@ -16,6 +16,7 @@ from aws_resource_management.utils.aws_core_utils import ( create_boto3_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 fd266278..860e218d 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 + resources: List[Dict[str, Any]], account_id: str, account_name: Optional[str] = None, account_alias: Optional[str] = None ) -> None: """ Add account information to resources. @@ -177,11 +177,14 @@ 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 da9e3c8c..3cd384e6 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 @@ -230,43 +230,6 @@ def _assume_role_for_account(account_id: str) -> Optional[Dict[str, Any]]: return None -def _try_assume_role( - account_id: str, role_name: str, session: boto3.Session -) -> Optional[Dict[str, str]]: - """ - Try to assume a specific role in an account. - - Args: - account_id: AWS account ID - role_name: Role name to assume - session: boto3 session to use - - Returns: - Dictionary with credential information or None on failure - """ - try: - sts_client = session.client("sts") - role_arn = f"arn:aws:iam::{account_id}:role/{role_name}" - - # Log attempt at debug level - logger.debug(f"Attempting to assume role {role_arn} for account {account_id}") - - response = sts_client.assume_role( - RoleArn=role_arn, RoleSessionName="ResourceManagementSession" - ) - credentials = response["Credentials"] - return { - "aws_access_key_id": credentials["AccessKeyId"], - "aws_secret_access_key": credentials["SecretAccessKey"], - "aws_session_token": credentials["SessionToken"], - } - except botocore.exceptions.ClientError as e: - logger.error( - f"Failed to assume role {role_name} for account {account_id}: {str(e)}" - ) - return None - - # --------------------------------------------------------------------------- # Simplified session management # --------------------------------------------------------------------------- @@ -530,6 +493,55 @@ 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."""