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

runbooks/vpc/cross_account_session.py

@@ -24,8 +24,9 @@ logger = logging.getLogger(__name__)
 @dataclass
 class AccountSession:
     """Represents a cross-account session with metadata"""
+
     account_id: str
-    account_name: Optional[str] 
+    account_name: Optional[str]
     session: boto3.Session
     status: str
     error_message: Optional[str] = None
@@ -34,68 +35,67 @@ class AccountSession:
 def create_multi_profile_sessions(profiles: List[str]) -> List[AccountSession]:
     """
     Create sessions using direct profile access for organizations without cross-account roles.
-    
+
     This is an alternative approach when OrganizationAccountAccessRole is not available.
     Uses environment variables like CENTRALISED_OPS_PROFILE, BILLING_PROFILE, etc.
-    
+
     Args:
         profiles: List of AWS profile names to use
-        
+
     Returns:
         List of AccountSession objects with successful and failed sessions
     """
     import os
-    
+
     print_info(f"🌐 Creating sessions for {len(profiles)} profiles")
     account_sessions = []
-    
+
     for profile_name in profiles:
         try:
             # Validate profile exists and is accessible
             session = boto3.Session(profile_name=profile_name)
-            sts_client = session.client('sts')
+            sts_client = session.client("sts")
             identity = sts_client.get_caller_identity()
-            
-            account_id = identity['Account']
-            
+
+            account_id = identity["Account"]
+
             # Try to get account name from Organizations if possible
             account_name = profile_name  # Default to profile name
             try:
-                orgs_client = session.client('organizations')
+                orgs_client = session.client("organizations")
                 account_info = orgs_client.describe_account(AccountId=account_id)
-                account_name = account_info['Account']['Name']
+                account_name = account_info["Account"]["Name"]
             except:
                 pass  # Use profile name as fallback
-            
-            account_sessions.append(AccountSession(
-                account_id=account_id,
-                account_name=account_name,
-                session=session,
-                status="success"
-            ))
-            
+
+            account_sessions.append(
+                AccountSession(account_id=account_id, account_name=account_name, session=session, status="success")
+            )
+
             print_success(f"✅ Session created for {account_id} using profile {profile_name}")
-            
+
         except Exception as e:
             print_warning(f"⚠️ Failed to create session for profile {profile_name}: {e}")
-            account_sessions.append(AccountSession(
-                account_id=profile_name,  # Use profile name as ID when we can't get real ID
-                account_name=profile_name,
-                session=None,
-                status="failed",
-                error_message=str(e)
-            ))
-    
+            account_sessions.append(
+                AccountSession(
+                    account_id=profile_name,  # Use profile name as ID when we can't get real ID
+                    account_name=profile_name,
+                    session=None,
+                    status="failed",
+                    error_message=str(e),
+                )
+            )
+
     return account_sessions
 
 
 class CrossAccountSessionManager:
     """
     Enterprise cross-account session manager using STS AssumeRole pattern.
-    
+
     This replaces the broken profile@accountId format with proper STS AssumeRole
     for multi-account VPC discovery across Landing Zone accounts.
-    
+
     Key Features:
     - Uses CENTRALISED_OPS_PROFILE as base session for assuming roles
     - Standard OrganizationAccountAccessRole assumption
@@ -103,91 +103,94 @@ class CrossAccountSessionManager:
     - Comprehensive error handling and graceful degradation
     - Compatible with existing VPC module architecture
     """
-    
+
     def __init__(self, base_profile: str, role_name: str = "OrganizationAccountAccessRole"):
         """
         Initialize cross-account session manager.
-        
+
         Args:
             base_profile: Base profile (e.g., CENTRALISED_OPS_PROFILE) for assuming roles
             role_name: IAM role name to assume in target accounts
         """
         self.base_profile = base_profile
         self.role_name = role_name
-        
+
         # Use management session for cross-account role assumptions
         # Management account has the trust relationships for OrganizationAccountAccessRole
-        self.session = create_management_session(profile=base_profile)
-        
+        self.session = create_management_session(profile_name=base_profile)
+
         print_info(f"🔐 Cross-account session manager initialized with {base_profile}")
-        
+
     def create_cross_account_sessions(
-        self, 
-        accounts: List[Dict[str, str]], 
-        max_workers: int = 10
+        self, accounts: List[Dict[str, str]], max_workers: int = 10
     ) -> List[AccountSession]:
         """
         Create cross-account sessions using STS AssumeRole pattern.
-        
+
         Args:
             accounts: List of account dictionaries from Organizations API
             max_workers: Maximum parallel workers for session creation
-            
+
         Returns:
             List of AccountSession objects with successful and failed sessions
         """
         print_info(f"🌐 Creating cross-account sessions for {len(accounts)} accounts")
-        
+
         account_sessions = []
-        
+
         # Filter active accounts only
         active_accounts = [acc for acc in accounts if acc.get("status") == "ACTIVE"]
         print_info(f"📋 Processing {len(active_accounts)} active accounts")
-        
+
         # Create sessions in parallel for performance
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
             future_to_account = {
-                executor.submit(self._create_account_session, account): account
-                for account in active_accounts
+                executor.submit(self._create_account_session, account): account for account in active_accounts
             }
-            
+
             for future in as_completed(future_to_account):
                 account = future_to_account[future]
                 try:
                     account_session = future.result()
                     account_sessions.append(account_session)
-                    
+
                     if account_session.status == "success":
                         print_success(f"✅ Session created for {account_session.account_id}")
                     else:
-                        print_warning(f"⚠️  Session failed for {account_session.account_id}: {account_session.error_message}")
-                        
+                        print_warning(
+                            f"⚠️  Session failed for {account_session.account_id}: {account_session.error_message}"
+                        )
+
                 except Exception as e:
                     print_error(f"❌ Unexpected error creating session for {account['id']}: {e}")
-                    account_sessions.append(AccountSession(
-                        account_id=account["id"],
-                        account_name=account.get("name"),
-                        session=None,
-                        status="error",
-                        error_message=str(e)
-                    ))
-        
+                    account_sessions.append(
+                        AccountSession(
+                            account_id=account["id"],
+                            account_name=account.get("name"),
+                            session=None,
+                            status="error",
+                            error_message=str(e),
+                        )
+                    )
+
         successful_sessions = [s for s in account_sessions if s.status == "success"]
         failed_sessions = [s for s in account_sessions if s.status != "success"]
-        
-        print_info(f"🎯 Session creation complete: {len(successful_sessions)} successful, {len(failed_sessions)} failed")
-        
+
+        print_info(
+            f"🎯 Session creation complete: {len(successful_sessions)} successful, {len(failed_sessions)} failed"
+        )
+
         return account_sessions
-    
+
     def _create_account_session(self, account: Dict[str, str]) -> AccountSession:
         """
         Create a session for a single account using STS AssumeRole.
-        
+
         This is the core implementation of the enterprise pattern.
         """
         account_id = account["id"]
         account_name = account.get("name", f"Account-{account_id}")
-        
+
         # Try multiple role patterns for different organization setups - universal compatibility
         role_patterns = [
             self.role_name,  # Default: OrganizationAccountAccessRole
@@ -199,87 +202,80 @@ class CrossAccountSessionManager:
             "CrossAccountRole",  # Generic cross-account role
             "AssumeRole",  # Generic assume role
         ]
-        
+
         for role_name in role_patterns:
             try:
                 # Step 1: Assume role in target account using STS
-                sts_client = self.session.client('sts')
+                sts_client = self.session.client("sts")
                 assumed_role = sts_client.assume_role(
                     RoleArn=f"arn:aws:iam::{account_id}:role/{role_name}",
-                    RoleSessionName=f"VPCDiscovery-{account_id[:12]}"
+                    RoleSessionName=f"VPCDiscovery-{account_id[:12]}",
                 )
-                
+
                 # If successful, continue with this role
-                
+
                 # Step 2: Create session with assumed role credentials
                 assumed_session = boto3.Session(
-                    aws_access_key_id=assumed_role['Credentials']['AccessKeyId'],
-                    aws_secret_access_key=assumed_role['Credentials']['SecretAccessKey'],
-                    aws_session_token=assumed_role['Credentials']['SessionToken']
+                    aws_access_key_id=assumed_role["Credentials"]["AccessKeyId"],
+                    aws_secret_access_key=assumed_role["Credentials"]["SecretAccessKey"],
+                    aws_session_token=assumed_role["Credentials"]["SessionToken"],
                 )
-                
+
                 # Step 3: Validate session with basic STS call
-                assumed_sts = assumed_session.client('sts')
+                assumed_sts = assumed_session.client("sts")
                 identity = assumed_sts.get_caller_identity()
-                
-                logger.debug(f"Successfully assumed role {role_name} in account {account_id}, identity: {identity['Arn']}")
-                
+
+                logger.debug(
+                    f"Successfully assumed role {role_name} in account {account_id}, identity: {identity['Arn']}"
+                )
+
                 return AccountSession(
-                    account_id=account_id,
-                    account_name=account_name,
-                    session=assumed_session,
-                    status="success"
+                    account_id=account_id, account_name=account_name, session=assumed_session, status="success"
                 )
-                
+
             except ClientError as e:
                 # Continue to next role pattern
                 continue
-        
+
         # If no role patterns worked, return failure
         error_msg = f"Unable to assume any role pattern in {account_id} - tried: {', '.join(role_patterns)}"
         logger.warning(f"Failed to create session for {account_id}: {error_msg}")
-        
+
         return AccountSession(
-            account_id=account_id,
-            account_name=account_name,
-            session=None,
-            status="failed",
-            error_message=error_msg
+            account_id=account_id, account_name=account_name, session=None, status="failed", error_message=error_msg
         )
-    
+
     def get_successful_sessions(self, account_sessions: List[AccountSession]) -> List[AccountSession]:
         """Get only successful account sessions for VPC discovery."""
         successful = [s for s in account_sessions if s.status == "success"]
         print_info(f"🎯 {len(successful)} accounts ready for VPC discovery")
         return successful
-    
+
     def get_session_summary(self, account_sessions: List[AccountSession]) -> Dict[str, int]:
         """Get summary statistics for session creation."""
         summary = {
             "total": len(account_sessions),
             "successful": len([s for s in account_sessions if s.status == "success"]),
             "failed": len([s for s in account_sessions if s.status == "failed"]),
-            "errors": len([s for s in account_sessions if s.status == "error"])
+            "errors": len([s for s in account_sessions if s.status == "error"]),
         }
         return summary
 
 
 def create_cross_account_vpc_sessions(
-    accounts: List[Dict[str, str]], 
-    base_profile: str,
-    role_name: str = "OrganizationAccountAccessRole"
+    accounts: List[Dict[str, str]], base_profile: str, role_name: str = "OrganizationAccountAccessRole"
 ) -> List[AccountSession]:
     """
     Convenience function to create cross-account VPC sessions.
-    
+
     This is the main entry point for VPC modules to replace the broken
     profile@accountId pattern with proper STS AssumeRole.
-    
+
     Args:
         accounts: List of organization accounts from get_organization_accounts
         base_profile: Base profile for assuming roles (CENTRALISED_OPS_PROFILE)
         role_name: IAM role name to assume
-        
+
     Returns:
         List of AccountSession objects ready for VPC discovery
     """
@@ -289,24 +285,23 @@ def create_cross_account_vpc_sessions(
 
 # Compatibility functions for existing VPC module integration
 def convert_accounts_to_sessions(
-    accounts: List[Dict[str, str]], 
-    base_profile: str
+    accounts: List[Dict[str, str]], base_profile: str
 ) -> Tuple[List[AccountSession], Dict[str, Dict[str, str]]]:
     """
     Convert organization accounts to cross-account sessions.
-    
+
     This replaces the broken convert_accounts_to_profiles function
     with proper STS AssumeRole session creation.
-    
+
     Returns:
         Tuple of (successful_sessions, account_metadata)
     """
     account_sessions = create_cross_account_vpc_sessions(accounts, base_profile)
     successful_sessions = [s for s in account_sessions if s.status == "success"]
-    
+
     # Create account metadata dict for compatibility
     account_metadata = {}
     for account in accounts:
         account_metadata[account["id"]] = account
-    
-    return successful_sessions, account_metadata
+
+    return successful_sessions, account_metadata