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

runbooks/common/profile_utils.py

@@ -17,14 +17,12 @@ _session_cache: Dict[str, boto3.Session] = {}  # Cache AWS sessions
 
 # Timeout configuration for AWS API calls (prevents execution flow hangs)
 _AWS_CLIENT_CONFIG = Config(
-    connect_timeout=30,    # Connection timeout: 30 seconds
-    read_timeout=60,       # Read timeout: 60 seconds
-    retries={
-        'max_attempts': 3,
-        'mode': 'adaptive'
-    }
+    connect_timeout=30,  # Connection timeout: 30 seconds
+    read_timeout=60,  # Read timeout: 60 seconds
+    retries={"max_attempts": 3, "mode": "adaptive"},
 )
 
+
 def _get_session_id() -> str:
     """Generate consistent session ID for cache scoping"""
     global _session_id
@@ -32,10 +30,9 @@ def _get_session_id() -> str:
         _session_id = f"session_{int(time.time())}"
     return _session_id
 
+
 def get_profile_for_operation(
-    operation_type: str,
-    user_specified_profile: Optional[str] = None,
-    profiles: Optional[List[str]] = None
+    operation_type: str, user_specified_profile: Optional[str] = None, profiles: Optional[List[str]] = None
 ) -> str:
     """
     Enhanced profile resolution with intelligent caching and enterprise logging optimization.
@@ -65,9 +62,7 @@ def get_profile_for_operation(
     profile_cache_key = f"{_get_session_id()}:{user_specified_profile or 'default'}"
 
     # Return cached result if still valid and within TTL
-    if (profile_cache_key in _profile_cache and
-        _cache_timestamp and
-        current_time - _cache_timestamp < _cache_ttl):
+    if profile_cache_key in _profile_cache and _cache_timestamp and current_time - _cache_timestamp < _cache_ttl:
         return _profile_cache[profile_cache_key]
 
     # Update cache timestamp only when cache is actually refreshed
@@ -125,6 +120,7 @@ def get_profile_for_operation(
         console.log("[yellow]Please run: aws configure sso or aws configure[/]")
         raise SystemExit(1)
 
+
 def validate_profile_access(profile_name: str, operation_description: str = "") -> bool:
     """
     Validate that the specified profile has proper AWS access with caching.
@@ -141,14 +137,12 @@ def validate_profile_access(profile_name: str, operation_description: str = "")
     current_time = time.time()
     cache_key = f"validation:{profile_name}"
 
-    if (cache_key in _validation_cache and
-        _cache_timestamp and
-        current_time - _cache_timestamp < _cache_ttl):
+    if cache_key in _validation_cache and _cache_timestamp and current_time - _cache_timestamp < _cache_ttl:
         return _validation_cache[cache_key]
 
     try:
         session = boto3.Session(profile_name=profile_name)
-        sts_client = session.client('sts')
+        sts_client = session.client("sts")
         sts_client.get_caller_identity()
 
         # Cache successful validation
@@ -160,6 +154,7 @@ def validate_profile_access(profile_name: str, operation_description: str = "")
         console.log(f"[yellow]Profile {profile_name} validation failed: {e}[/]")
         return False
 
+
 def get_account_id_from_profile(profile_name: str) -> Optional[str]:
     """
     Extract account ID from AWS profile.
@@ -172,12 +167,87 @@ def get_account_id_from_profile(profile_name: str) -> Optional[str]:
     """
     try:
         session = boto3.Session(profile_name=profile_name)
-        sts_client = session.client('sts')
+        sts_client = session.client("sts")
         response = sts_client.get_caller_identity()
-        return response.get('Account')
+        return response.get("Account")
     except Exception:
         return None
 
+
+def auto_discover_enterprise_profiles() -> Dict[str, Optional[str]]:
+    """
+    Auto-discover enterprise AWS SSO profiles for streamlined initialization.
+
+    Searches for profiles matching common enterprise naming patterns:
+    - *Billing* or *billing* for BILLING_PROFILE
+    - *Management* or *management* for MANAGEMENT_PROFILE
+    - *Ops* or *ops* for CENTRALISED_OPS_PROFILE
+    - Single account profiles for SINGLE_AWS_PROFILE
+
+    Returns:
+        Dict mapping profile types to discovered profile names
+    """
+    available_profiles = boto3.Session().available_profiles
+    discovered = {"billing": None, "management": None, "centralised_ops": None, "single_aws": None}
+
+    # Search patterns for enterprise profiles
+    for profile in available_profiles:
+        profile_lower = profile.lower()
+
+        # Billing profile detection
+        if ("billing" in profile_lower or "cost" in profile_lower) and not discovered["billing"]:
+            discovered["billing"] = profile
+
+        # Management profile detection
+        elif ("management" in profile_lower or "admin" in profile_lower) and not discovered["management"]:
+            discovered["management"] = profile
+
+        # Operations profile detection
+        elif (
+            "ops" in profile_lower or "operational" in profile_lower or "centralised" in profile_lower
+        ) and not discovered["centralised_ops"]:
+            discovered["centralised_ops"] = profile
+
+        # Single account detection (typically shorter names or containing 'single')
+        elif ("single" in profile_lower or len(profile) < 20) and not discovered["single_aws"]:
+            discovered["single_aws"] = profile
+
+    # Log discovered profiles for transparency
+    for profile_type, profile_name in discovered.items():
+        if profile_name:
+            console.log(f"[green]✅ Auto-discovered {profile_type}: {profile_name}[/green]")
+        else:
+            console.log(f"[yellow]⚠️ No profile found for {profile_type}[/yellow]")
+
+    return discovered
+
+
+def setup_enterprise_environment_variables(discovered_profiles: Optional[Dict[str, Optional[str]]] = None) -> None:
+    """
+    Setup enterprise environment variables from discovered profiles.
+
+    Args:
+        discovered_profiles: Optional pre-discovered profiles dict
+    """
+    if not discovered_profiles:
+        discovered_profiles = auto_discover_enterprise_profiles()
+
+    # Set environment variables if not already set
+    env_mappings = {
+        "BILLING_PROFILE": discovered_profiles.get("billing"),
+        "MANAGEMENT_PROFILE": discovered_profiles.get("management"),
+        "CENTRALISED_OPS_PROFILE": discovered_profiles.get("centralised_ops"),
+        "SINGLE_AWS_PROFILE": discovered_profiles.get("single_aws"),
+    }
+
+    for env_var, profile_name in env_mappings.items():
+        if profile_name and not os.getenv(env_var):
+            os.environ[env_var] = profile_name
+            console.log(f"[blue]📋 Set {env_var}={profile_name}[/blue]")
+        elif os.getenv(env_var):
+            console.log(f"[dim]Using existing {env_var}={os.getenv(env_var)}[/dim]")
+
+
 def create_cost_session(profile_name: Optional[str] = None) -> boto3.Session:
     """
     Create AWS session optimized for cost operations (Cost Explorer) with token error handling.
@@ -203,7 +273,7 @@ def create_cost_session(profile_name: Optional[str] = None) -> boto3.Session:
         # Quick validation that the cached session still works
         try:
             # Test with a minimal STS call to check if credentials are valid (with timeout)
-            sts_client = cached_session.client('sts', config=_AWS_CLIENT_CONFIG)
+            sts_client = cached_session.client("sts", config=_AWS_CLIENT_CONFIG)
             sts_client.get_caller_identity()
             return cached_session
         except (TokenRetrievalError, NoCredentialsError):
@@ -214,7 +284,7 @@ def create_cost_session(profile_name: Optional[str] = None) -> boto3.Session:
     try:
         session = boto3.Session(profile_name=cost_profile)
         # Test the session to ensure credentials are valid (with timeout)
-        sts_client = session.client('sts', config=_AWS_CLIENT_CONFIG)
+        sts_client = session.client("sts", config=_AWS_CLIENT_CONFIG)
         sts_client.get_caller_identity()
 
         # Cache only if session works
@@ -250,6 +320,7 @@ def create_cost_session(profile_name: Optional[str] = None) -> boto3.Session:
         console.log(f"[dim]Error details: {str(e)}[/]")
         raise SystemExit(1)
 
+
 def create_management_session(profile_name: Optional[str] = None) -> boto3.Session:
     """
     Create AWS session optimized for management operations (Organizations) with token error handling.
@@ -275,7 +346,7 @@ def create_management_session(profile_name: Optional[str] = None) -> boto3.Sessi
         # Quick validation that the cached session still works
         try:
             # Test with a minimal STS call to check if credentials are valid (with timeout)
-            sts_client = cached_session.client('sts', config=_AWS_CLIENT_CONFIG)
+            sts_client = cached_session.client("sts", config=_AWS_CLIENT_CONFIG)
             sts_client.get_caller_identity()
             return cached_session
         except (TokenRetrievalError, NoCredentialsError):
@@ -286,7 +357,7 @@ def create_management_session(profile_name: Optional[str] = None) -> boto3.Sessi
     try:
         session = boto3.Session(profile_name=mgmt_profile)
         # Test the session to ensure credentials are valid (with timeout)
-        sts_client = session.client('sts', config=_AWS_CLIENT_CONFIG)
+        sts_client = session.client("sts", config=_AWS_CLIENT_CONFIG)
         sts_client.get_caller_identity()
 
         # Cache only if session works
@@ -322,6 +393,7 @@ def create_management_session(profile_name: Optional[str] = None) -> boto3.Sessi
         console.log(f"[dim]Error details: {str(e)}[/]")
         raise SystemExit(1)
 
+
 def create_operational_session(profile_name: Optional[str] = None) -> boto3.Session:
     """
     Create AWS session optimized for operational tasks (EC2, S3, etc).
@@ -343,6 +415,7 @@ def create_operational_session(profile_name: Optional[str] = None) -> boto3.Sess
     _session_cache[session_key] = session
     return session
 
+
 def get_current_profile_info(profile_name: Optional[str] = None) -> Dict[str, Any]:
     """
     Get current profile information including account ID and region.
@@ -355,28 +428,26 @@ def get_current_profile_info(profile_name: Optional[str] = None) -> Dict[str, An
     """
     try:
         session = boto3.Session(profile_name=profile_name)
-        sts_client = session.client('sts', config=_AWS_CLIENT_CONFIG)
+        sts_client = session.client("sts", config=_AWS_CLIENT_CONFIG)
         identity = sts_client.get_caller_identity()
 
         return {
-            'profile_name': profile_name or 'default',
-            'account_id': identity.get('Account'),
-            'user_arn': identity.get('Arn'),
-            'region': session.region_name or 'us-east-1'
+            "profile_name": profile_name or "default",
+            "account_id": identity.get("Account"),
+            "user_arn": identity.get("Arn"),
+            "region": session.region_name or "us-east-1",
         }
     except Exception as e:
         return {
-            'profile_name': profile_name or 'default',
-            'error': str(e),
-            'account_id': None,
-            'user_arn': None,
-            'region': None
+            "profile_name": profile_name or "default",
+            "error": str(e),
+            "account_id": None,
+            "user_arn": None,
+            "region": None,
         }
 
-def resolve_profile_for_operation_silent(
-    operation_type: str,
-    user_specified_profile: Optional[str] = None
-) -> str:
+
+def resolve_profile_for_operation_silent(operation_type: str, user_specified_profile: Optional[str] = None) -> str:
     """
     Silent version of profile resolution without logging.
 
@@ -404,6 +475,7 @@ def resolve_profile_for_operation_silent(
 
     return user_specified_profile or "default"
 
+
 def list_available_profiles() -> List[str]:
     """
     Get list of all available AWS profiles.
@@ -413,6 +485,7 @@ def list_available_profiles() -> List[str]:
     """
     return boto3.Session().available_profiles
 
+
 def clear_profile_cache() -> None:
     """Clear the profile cache for testing or troubleshooting."""
     global _profile_cache, _validation_cache, _session_cache, _cache_timestamp, _session_id
@@ -422,6 +495,7 @@ def clear_profile_cache() -> None:
     _cache_timestamp = None
     _session_id = None
 
+
 def create_timeout_protected_client(session: boto3.Session, service_name: str, region_name: Optional[str] = None):
     """
     Create AWS service client with timeout protection to prevent execution flow hangs.
@@ -442,4 +516,4 @@ def create_timeout_protected_client(session: boto3.Session, service_name: str, r
         ce_client = create_timeout_protected_client(session, 'ce', 'us-east-1')
         ec2_client = create_timeout_protected_client(session, 'ec2', region_name)
     """
-    return session.client(service_name, region_name=region_name, config=_AWS_CLIENT_CONFIG)
+    return session.client(service_name, region_name=region_name, config=_AWS_CLIENT_CONFIG)