runbooks 1.1.4__py3-none-any.whl → 1.1.6__py3-none-any.whl

runbooks/common/aws_profile_manager.py

@@ -24,6 +24,7 @@ from rich.console import Console
 
 from runbooks.common.rich_utils import console, print_error, print_success, print_warning, print_info
 
+
 class AWSProfileManager:
     """
     Universal AWS Profile Manager for CloudOps v1.1.x compatibility.
@@ -58,14 +59,14 @@ class AWSProfileManager:
             return user_profile
 
         # Tier 2: Environment variable
-        env_profile = os.getenv('AWS_PROFILE')
+        env_profile = os.getenv("AWS_PROFILE")
         if env_profile:
             return env_profile
 
         # Tier 3: Default (no explicit profile)
         return None
 
-    def get_session(self, region: str = 'us-east-1') -> boto3.Session:
+    def get_session(self, region: str = "us-east-1") -> boto3.Session:
         """
         Get boto3 session with resolved profile.
 
@@ -89,7 +90,7 @@ class AWSProfileManager:
                     print_info("Using default AWS credentials")
 
                 # Validate credentials by getting caller identity
-                sts = self.session.client('sts')
+                sts = self.session.client("sts")
                 identity = sts.get_caller_identity()
                 print_success(f"✅ Authenticated as: {identity.get('Arn', 'Unknown')}")
 
@@ -108,7 +109,7 @@ class AWSProfileManager:
 
         return self.session
 
-    def get_account_id(self, region: str = 'us-east-1') -> str:
+    def get_account_id(self, region: str = "us-east-1") -> str:
         """
         Get current AWS account ID dynamically.
 
@@ -123,9 +124,9 @@ class AWSProfileManager:
         if cache_key not in self._account_cache:
             try:
                 session = self.get_session(region)
-                sts = session.client('sts')
+                sts = session.client("sts")
                 identity = sts.get_caller_identity()
-                account_id = identity['Account']
+                account_id = identity["Account"]
                 self._account_cache[cache_key] = account_id
                 print_info(f"Current account ID: {account_id}")
 
@@ -136,7 +137,7 @@ class AWSProfileManager:
 
         return self._account_cache[cache_key]
 
-    def discover_organization_accounts(self, region: str = 'us-east-1') -> List[Dict[str, Any]]:
+    def discover_organization_accounts(self, region: str = "us-east-1") -> List[Dict[str, Any]]:
         """
         Discover all accounts in AWS Organizations (if available).
 
@@ -148,7 +149,7 @@ class AWSProfileManager:
         """
         try:
             session = self.get_session(region)
-            org_client = session.client('organizations')
+            org_client = session.client("organizations")
 
             # Get organization information
             try:
@@ -160,17 +161,19 @@ class AWSProfileManager:
 
             # List all accounts
             accounts = []
-            paginator = org_client.get_paginator('list_accounts')
+            paginator = org_client.get_paginator("list_accounts")
 
             for page in paginator.paginate():
-                for account in page['Accounts']:
-                    accounts.append({
-                        'Id': account['Id'],
-                        'Name': account['Name'],
-                        'Email': account['Email'],
-                        'Status': account['Status'],
-                        'JoinedMethod': account.get('JoinedMethod', 'UNKNOWN')
-                    })
+                for account in page["Accounts"]:
+                    accounts.append(
+                        {
+                            "Id": account["Id"],
+                            "Name": account["Name"],
+                            "Email": account["Email"],
+                            "Status": account["Status"],
+                            "JoinedMethod": account.get("JoinedMethod", "UNKNOWN"),
+                        }
+                    )
 
             print_success(f"✅ Discovered {len(accounts)} organization accounts")
             return accounts
@@ -190,7 +193,7 @@ class AWSProfileManager:
             Dict mapping service names to access status
         """
         if not required_services:
-            required_services = ['sts', 'ce', 'ec2', 's3']
+            required_services = ["sts", "ce", "ec2", "s3"]
 
         access_status = {}
         session = self.get_session()
@@ -200,21 +203,22 @@ class AWSProfileManager:
                 client = session.client(service)
 
                 # Service-specific health checks
-                if service == 'sts':
+                if service == "sts":
                     client.get_caller_identity()
-                elif service == 'ce':
+                elif service == "ce":
                     # Test Cost Explorer access
                     from datetime import datetime, timedelta
-                    end_date = datetime.now().strftime('%Y-%m-%d')
-                    start_date = (datetime.now() - timedelta(days=7)).strftime('%Y-%m-%d')
+
+                    end_date = datetime.now().strftime("%Y-%m-%d")
+                    start_date = (datetime.now() - timedelta(days=7)).strftime("%Y-%m-%d")
                     client.get_cost_and_usage(
-                        TimePeriod={'Start': start_date, 'End': end_date},
-                        Granularity='MONTHLY',
-                        Metrics=['UnblendedCost']
+                        TimePeriod={"Start": start_date, "End": end_date},
+                        Granularity="MONTHLY",
+                        Metrics=["UnblendedCost"],
                     )
-                elif service == 'ec2':
+                elif service == "ec2":
                     client.describe_regions(MaxResults=1)
-                elif service == 's3':
+                elif service == "s3":
                     client.list_buckets()
 
                 access_status[service] = True
@@ -232,8 +236,8 @@ class AWSProfileManager:
             import configparser
             import os.path
 
-            credentials_path = os.path.expanduser('~/.aws/credentials')
-            config_path = os.path.expanduser('~/.aws/config')
+            credentials_path = os.path.expanduser("~/.aws/credentials")
+            config_path = os.path.expanduser("~/.aws/config")
 
             profiles = set()
 
@@ -248,10 +252,10 @@ class AWSProfileManager:
                 config_config = configparser.ConfigParser()
                 config_config.read(config_path)
                 for section in config_config.sections():
-                    if section.startswith('profile '):
+                    if section.startswith("profile "):
                         profiles.add(section[8:])  # Remove 'profile ' prefix
-                    elif section == 'default':
-                        profiles.add('default')
+                    elif section == "default":
+                        profiles.add("default")
 
             if profiles:
                 for profile in sorted(profiles):
@@ -263,7 +267,7 @@ class AWSProfileManager:
             print_warning(f"⚠️ Could not list profiles: {e}")
 
     @classmethod
-    def create_mock_account_context(cls, mock_account_id: str = "123456789012") -> 'AWSProfileManager':
+    def create_mock_account_context(cls, mock_account_id: str = "123456789012") -> "AWSProfileManager":
         """
         Create mock profile manager for testing scenarios.
 
@@ -274,7 +278,7 @@ class AWSProfileManager:
             ProfileManager configured for testing
         """
         manager = cls()
-        manager._account_cache['mock'] = mock_account_id
+        manager._account_cache["mock"] = mock_account_id
         return manager
 
     def get_profile_display_name(self) -> str:
@@ -294,7 +298,7 @@ class AWSProfileManager:
 
 
 # Global convenience functions for backward compatibility
-def get_current_account_id(profile: Optional[str] = None, region: str = 'us-east-1') -> str:
+def get_current_account_id(profile: Optional[str] = None, region: str = "us-east-1") -> str:
     """
     Convenience function to get current account ID.
 
@@ -334,4 +338,4 @@ def validate_profile_or_exit(profile: Optional[str] = None, required_services: L
     except (ProfileNotFound, NoCredentialsError):
         print_error("❌ Cannot proceed without valid AWS credentials")
         print_info("Please configure AWS credentials and try again.")
-        exit(1)
+        exit(1)

runbooks/common/aws_utils.py

@@ -31,67 +31,67 @@ from runbooks.common.rich_utils import console
 class AWSProfileSanitizer:
     """
     Enterprise-grade AWS profile name sanitization for security logging.
-    
+
     Prevents AWS account ID exposure in logs while maintaining audit trail integrity.
     Following FAANG security-as-code principles for sensitive identifier protection.
     """
-    
+
     # Pattern to detect AWS account IDs in profile names
-    ACCOUNT_ID_PATTERN = re.compile(r'\b\d{12}\b')
-    
+    ACCOUNT_ID_PATTERN = re.compile(r"\b\d{12}\b")
+
     # Pattern to detect enterprise profile patterns
-    ENTERPRISE_PROFILE_PATTERN = re.compile(r'(ams|aws)-.*-ReadOnlyAccess-(\d{12})')
-    
+    ENTERPRISE_PROFILE_PATTERN = re.compile(r"(ams|aws)-.*-ReadOnlyAccess-(\d{12})")
+
     @classmethod
     def sanitize_profile_name(cls, profile_name: str, mask_style: str = "***masked***") -> str:
         """
         Sanitize AWS profile name by masking account IDs for secure logging.
-        
+
         Replaces 12-digit AWS account IDs with masked values to prevent account enumeration
         while preserving profile identification capabilities for audit purposes.
-        
+
         Args:
             profile_name: Original AWS profile name
             mask_style: Masking pattern for account IDs (default: ***masked***)
-            
+
         Returns:
             Sanitized profile name with masked account IDs
-            
+
         Example:
             'my-billing-profile-123456789012' → 'my-billing-profile-***masked***'
         """
         if not profile_name:
             return profile_name
-            
+
         # Check for enterprise pattern first (more specific)
         if cls.ENTERPRISE_PROFILE_PATTERN.match(profile_name):
-            return cls.ENTERPRISE_PROFILE_PATTERN.sub(r'\1-masked-ReadOnlyAccess-***masked***', profile_name)
-        
+            return cls.ENTERPRISE_PROFILE_PATTERN.sub(r"\1-masked-ReadOnlyAccess-***masked***", profile_name)
+
         # General account ID masking
         return cls.ACCOUNT_ID_PATTERN.sub(mask_style, profile_name)
-    
+
     @classmethod
     def sanitize_profile_list(cls, profiles: List[str]) -> List[str]:
         """
         Sanitize a list of AWS profile names for secure logging.
-        
+
         Args:
             profiles: List of AWS profile names
-            
+
         Returns:
             List of sanitized profile names
         """
         return [cls.sanitize_profile_name(profile) for profile in profiles]
-    
+
     @classmethod
     def create_secure_log_context(cls, profile: str, operation: str) -> Dict[str, str]:
         """
         Create secure logging context with sanitized profile information.
-        
+
         Args:
             profile: AWS profile name
             operation: Operation being performed
-            
+
         Returns:
             Dictionary with sanitized context for secure logging
         """
@@ -99,14 +99,14 @@ class AWSProfileSanitizer:
             "operation": operation,
             "profile_sanitized": cls.sanitize_profile_name(profile),
             "profile_type": cls._classify_profile_type(profile),
-            "timestamp": datetime.utcnow().isoformat()
+            "timestamp": datetime.utcnow().isoformat(),
         }
-    
+
     @classmethod
     def _classify_profile_type(cls, profile_name: str) -> str:
         """Classify profile type for enhanced logging context."""
         profile_lower = profile_name.lower()
-        
+
         if "billing" in profile_lower:
             return "billing"
         elif "management" in profile_lower:
@@ -122,109 +122,104 @@ class AWSProfileSanitizer:
 class AWSTokenManager:
     """
     Proactive AWS token management with security-focused error handling.
-    
+
     Implements proactive token refresh, retry logic, and enhanced error messaging
     to reduce authentication timing exposure and improve operational security.
     """
-    
+
     # Token refresh thresholds
     TOKEN_REFRESH_THRESHOLD_MINUTES = 15  # Refresh if expires within 15 minutes
     MAX_RETRY_ATTEMPTS = 3
     RETRY_BACKOFF_BASE = 2  # Exponential backoff base (seconds)
-    
+
     def __init__(self, profile_name: str):
         """Initialize token manager for specific AWS profile."""
         self.profile_name = profile_name
         self.sanitized_profile = AWSProfileSanitizer.sanitize_profile_name(profile_name)
         self._session = None
         self._last_refresh_check = None
-        
+
     def get_secure_session(self, force_refresh: bool = False) -> boto3.Session:
         """
         Get AWS session with proactive token refresh and security enhancements.
-        
+
         Implements:
         - Proactive token expiration checking
         - Silent token refresh before expiration
         - Exponential backoff retry logic
         - Security-aware error messages
-        
+
         Args:
             force_refresh: Force token refresh regardless of expiration status
-            
+
         Returns:
             Boto3 session with valid credentials
-            
+
         Raises:
             SecurityError: For authentication security issues
             TokenRefreshError: For token refresh failures
         """
         current_time = datetime.utcnow()
-        
+
         # Check if proactive refresh is needed
-        if (force_refresh or 
-            self._session is None or 
-            self._needs_token_refresh(current_time)):
-            
+        if force_refresh or self._session is None or self._needs_token_refresh(current_time):
             self._session = self._refresh_session_with_retry()
             self._last_refresh_check = current_time
-            
+
             # Log secure refresh event
-            console.log(
-                f"[dim green]✅ Token refresh completed for profile: {self.sanitized_profile}[/]"
-            )
-        
+            console.log(f"[dim green]✅ Token refresh completed for profile: {self.sanitized_profile}[/]")
+
         return self._session
-    
+
     def _needs_token_refresh(self, current_time: datetime) -> bool:
         """Check if proactive token refresh is needed."""
         if self._last_refresh_check is None:
             return True
-            
+
         # Check every 5 minutes to avoid excessive API calls
         if (current_time - self._last_refresh_check) < timedelta(minutes=5):
             return False
-            
+
         try:
             # Test session validity with STS call
             if self._session:
-                sts_client = self._session.client('sts')
+                sts_client = self._session.client("sts")
                 sts_client.get_caller_identity()
                 return False  # Session still valid
         except ClientError as e:
-            error_code = e.response.get('Error', {}).get('Code', '')
-            if error_code in ['ExpiredToken', 'InvalidToken', 'TokenRefreshRequired']:
+            error_code = e.response.get("Error", {}).get("Code", "")
+            if error_code in ["ExpiredToken", "InvalidToken", "TokenRefreshRequired"]:
                 return True
-                
+
         return False
-    
+
     def _refresh_session_with_retry(self) -> boto3.Session:
         """Refresh session with exponential backoff retry logic."""
         last_exception = None
-        
+
         for attempt in range(self.MAX_RETRY_ATTEMPTS):
             try:
                 # Create new session
                 session = boto3.Session(profile_name=self.profile_name)
-                
+
                 # Validate session with STS call
-                sts_client = session.client('sts')
+                sts_client = session.client("sts")
                 caller_identity = sts_client.get_caller_identity()
-                
+
                 # Log successful refresh (with sanitized profile)
                 console.log(
                     f"[dim cyan]🔄 Session validated for {self.sanitized_profile} "
                     f"(attempt {attempt + 1}/{self.MAX_RETRY_ATTEMPTS})[/]"
                 )
-                
+
                 return session
-                
+
             except (ClientError, NoCredentialsError, TokenRetrievalError) as e:
                 last_exception = e
-                
+
                 if attempt < self.MAX_RETRY_ATTEMPTS - 1:
                     # Wait with exponential backoff
-                    wait_time = self.RETRY_BACKOFF_BASE ** attempt
+                    wait_time = self.RETRY_BACKOFF_BASE**attempt
                     console.log(
                         f"[yellow]⏳ Token refresh attempt {attempt + 1} failed, "
                         f"retrying in {wait_time}s for {self.sanitized_profile}[/]"
@@ -233,17 +228,17 @@ class AWSTokenManager:
                 else:
                     # Final attempt failed, provide enhanced guidance
                     self._handle_token_refresh_failure(last_exception)
-        
+
         # If we get here, all attempts failed
         raise TokenRefreshError(
             f"Failed to refresh AWS session for profile {self.sanitized_profile} "
             f"after {self.MAX_RETRY_ATTEMPTS} attempts"
         )
-    
+
     def _handle_token_refresh_failure(self, error: Exception) -> None:
         """Provide enhanced guidance for token refresh failures."""
         error_str = str(error)
-        
+
         # Determine error type for appropriate guidance
         if "ExpiredToken" in error_str or "InvalidToken" in error_str:
             console.log(
@@ -275,59 +270,59 @@ class AWSTokenManager:
 
 class SecurityError(Exception):
     """Raised for AWS security-related errors."""
+
     pass
 
 
 class TokenRefreshError(Exception):
     """Raised for AWS token refresh failures."""
+
     pass
 
 
 def create_secure_aws_session(profile_name: str, operation_context: str = "aws_operation") -> boto3.Session:
     """
     Create secure AWS session with enterprise security enhancements.
-    
+
     This is the primary entry point for secure AWS session creation across
     all CloudOps modules. Implements:
-    
+
     - Profile name sanitization for secure logging
     - Proactive token refresh
     - Enhanced error handling
     - Security audit trail
-    
+
     Args:
         profile_name: AWS profile name
         operation_context: Description of the operation for audit logging
-        
+
     Returns:
         Secure boto3 session with valid credentials
-        
+
     Raises:
         SecurityError: For security-related authentication issues
         TokenRefreshError: For token refresh failures
-        
+
     Example:
         session = create_secure_aws_session("my-billing-profile-123456789012", "cost_analysis")
     """
     # Create secure logging context
     log_context = AWSProfileSanitizer.create_secure_log_context(profile_name, operation_context)
-    
+
     console.log(
         f"[dim cyan]🔐 Initiating secure AWS session for {log_context['profile_sanitized']} "
         f"({log_context['profile_type']} profile)[/]"
     )
-    
+
     try:
         # Initialize token manager and get secure session
         token_manager = AWSTokenManager(profile_name)
         session = token_manager.get_secure_session()
-        
-        console.log(
-            f"[dim green]✅ Secure session established for {log_context['profile_sanitized']}[/]"
-        )
-        
+
+        console.log(f"[dim green]✅ Secure session established for {log_context['profile_sanitized']}[/]")
+
         return session
-        
+
     except Exception as e:
         console.log(
             f"[red]❌ Failed to create secure session for {log_context['profile_sanitized']}: {str(e)[:100]}[/]"
@@ -338,10 +333,10 @@ def create_secure_aws_session(profile_name: str, operation_context: str = "aws_o
 def sanitize_aws_error_message(error_message: str) -> str:
     """
     Sanitize AWS error messages to remove sensitive account information.
-    
+
     Args:
         error_message: Original AWS error message
-        
+
     Returns:
         Sanitized error message with account IDs masked
     """
@@ -351,10 +346,10 @@ def sanitize_aws_error_message(error_message: str) -> str:
 def get_profile_classification(profile_name: str) -> Dict[str, str]:
     """
     Get security classification information for AWS profile.
-    
+
     Args:
         profile_name: AWS profile name
-        
+
     Returns:
         Dictionary with profile security classification
     """
@@ -363,5 +358,5 @@ def get_profile_classification(profile_name: str) -> Dict[str, str]:
         "original": profile_name,
         "sanitized": sanitizer.sanitize_profile_name(profile_name),
         "type": sanitizer._classify_profile_type(profile_name),
-        "risk_level": "high" if "admin" in profile_name.lower() else "medium"
-    }
+        "risk_level": "high" if "admin" in profile_name.lower() else "medium",
+    }