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

runbooks/inventory/organizations_discovery.py

@@ -24,47 +24,53 @@ from typing import Dict, List, Optional, Set, Tuple
 
 import boto3
 from botocore.exceptions import ClientError, NoCredentialsError
-from rich.console import Console
 from rich.panel import Panel
-from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+from rich.progress import BarColumn, SpinnerColumn, TextColumn, TimeElapsedColumn
 from rich.status import Status
 from rich.table import Table
 
-# Initialize Rich console
-console = Console()
+# Test Mode Support: Use centralized rich_utils for console and Progress (test-aware)
+# Issue: Rich Console/Progress write to StringIO buffer that Click CliRunner closes, causing ValueError
+# Solution: rich_utils provides MockConsole/MockProgress in test mode, Rich components in production
+import os
+import sys
+import re
 
 from ..utils.logger import configure_logger
 from ..common.performance_optimization_engine import get_optimization_engine
+from ..common.rich_utils import console, Progress
 
 logger = configure_logger(__name__)
 
 # Global Organizations cache to prevent duplicate API calls across all instances
-_GLOBAL_ORGS_CACHE = {
-    'data': None,
-    'timestamp': None,
-    'ttl_minutes': 30
-}
+_GLOBAL_ORGS_CACHE = {"data": None, "timestamp": None, "ttl_minutes": 30}
+
 
 def _get_global_organizations_cache():
     """Get cached Organizations data if valid (module-level cache)."""
-    if not _GLOBAL_ORGS_CACHE['timestamp']:
+    if not _GLOBAL_ORGS_CACHE["timestamp"]:
         return None
-    
-    cache_age_minutes = (datetime.now(timezone.utc) - _GLOBAL_ORGS_CACHE['timestamp']).total_seconds() / 60
-    if cache_age_minutes < _GLOBAL_ORGS_CACHE['ttl_minutes']:
+
+    cache_age_minutes = (datetime.now(timezone.utc) - _GLOBAL_ORGS_CACHE["timestamp"]).total_seconds() / 60
+    if cache_age_minutes < _GLOBAL_ORGS_CACHE["ttl_minutes"]:
         console.print("[blue]🚀 Global performance optimization: Using cached Organizations data[/blue]")
-        return _GLOBAL_ORGS_CACHE['data']
+        return _GLOBAL_ORGS_CACHE["data"]
     return None
 
+
 def _set_global_organizations_cache(data):
     """Cache Organizations data globally (module-level cache)."""
-    _GLOBAL_ORGS_CACHE['data'] = data
-    _GLOBAL_ORGS_CACHE['timestamp'] = datetime.now(timezone.utc)
-    accounts_count = len(data.get('accounts', {}).get('discovered_accounts', [])) if data else 0
-    console.print(f"[green]✅ Global Organizations cache: {accounts_count} accounts (TTL: {_GLOBAL_ORGS_CACHE['ttl_minutes']}min)[/green]")
+    _GLOBAL_ORGS_CACHE["data"] = data
+    _GLOBAL_ORGS_CACHE["timestamp"] = datetime.now(timezone.utc)
+    accounts_count = len(data.get("accounts", {}).get("discovered_accounts", [])) if data else 0
+    console.print(
+        f"[green]✅ Global Organizations cache: {accounts_count} accounts (TTL: {_GLOBAL_ORGS_CACHE['ttl_minutes']}min)[/green]"
+    )
+
 
 # Universal AWS Environment Profile Support (Compatible with ANY AWS Setup)
 import os
+
 ENTERPRISE_PROFILES = {
     "BILLING_PROFILE": os.getenv("BILLING_PROFILE", "default"),  # Universal compatibility
     "MANAGEMENT_PROFILE": os.getenv("MANAGEMENT_PROFILE", "default"),  # Works with any profile
@@ -246,40 +252,44 @@ class EnhancedOrganizationsDiscovery:
         """Check if Organizations cache is still valid."""
         if not self._organizations_cache_timestamp:
             return False
-        
+
         from datetime import datetime, timedelta
+
         cache_age_minutes = (datetime.now() - self._organizations_cache_timestamp).total_seconds() / 60
         return cache_age_minutes < self._cache_ttl_minutes
 
     async def discover_all_accounts(self) -> Dict:
         """
         Cached wrapper for Organizations discovery to prevent duplicate API calls.
-        
-        This method implements both global and instance-level caching to avoid the 
-        performance penalty of duplicate Organizations API calls when multiple 
+
+        This method implements both global and instance-level caching to avoid the
+        performance penalty of duplicate Organizations API calls when multiple
         components need the same account data.
         """
         # Check global cache first (shared across all instances)
         global_cached_result = _get_global_organizations_cache()
         if global_cached_result:
             return global_cached_result
-        
+
         # Check instance cache
         if self._is_organizations_cache_valid() and self._organizations_cache:
             console.print("[blue]🚀 Performance optimization: Using cached Organizations data[/blue]")
             return self._organizations_cache
-        
+
         # Cache miss - perform discovery
         console.print("[cyan]🔍 Performing Organizations discovery (cache miss)[/cyan]")
         results = await self.discover_organization_structure()
-        
+
         # Cache the results
-        if results and results.get('accounts'):
+        if results and results.get("accounts"):
             self._organizations_cache = results
             from datetime import datetime
+
             self._organizations_cache_timestamp = datetime.now()
-            console.print(f"[green]✅ Cached Organizations data: {len(results.get('accounts', {}).get('discovered_accounts', []))} accounts (TTL: {self._cache_ttl_minutes}min)[/green]")
-        
+            console.print(
+                f"[green]✅ Cached Organizations data: {len(results.get('accounts', {}).get('discovered_accounts', []))} accounts (TTL: {self._cache_ttl_minutes}min)[/green]"
+            )
+
         return results
 
     def initialize_sessions(self) -> Dict[str, str]:
@@ -417,13 +427,13 @@ class EnhancedOrganizationsDiscovery:
         """
         # Get performance optimization engine
         optimization_engine = get_optimization_engine(
-            max_workers=self.max_workers,
-            cache_ttl_minutes=30,
-            memory_limit_mb=2048
+            max_workers=self.max_workers, cache_ttl_minutes=30, memory_limit_mb=2048
         )
-        
+
         # Use optimized discovery with performance monitoring
-        with optimization_engine.optimize_operation("organization_structure_discovery", self.performance_target_seconds):
+        with optimization_engine.optimize_operation(
+            "organization_structure_discovery", self.performance_target_seconds
+        ):
             return await self._discover_organization_structure_optimized(optimization_engine)
 
     async def _discover_organization_structure_optimized(self, optimization_engine) -> Dict:
@@ -436,7 +446,7 @@ class EnhancedOrganizationsDiscovery:
         )
 
         logger.info("🏢 Starting optimized organization structure discovery with SRE automation patterns")
-        
+
         # Check global cache first to prevent duplicate calls
         cached_result = _get_global_organizations_cache()
         if cached_result:
@@ -446,7 +456,7 @@ class EnhancedOrganizationsDiscovery:
             self.discovery_metrics["end_time"] = self.current_benchmark.end_time
             self.discovery_metrics["duration_seconds"] = self.current_benchmark.duration_seconds
             return cached_result
-            
+
         self.discovery_metrics["start_time"] = self.current_benchmark.start_time
 
         with Status("Initializing enterprise discovery...", console=console, spinner="dots"):
@@ -561,10 +571,10 @@ class EnhancedOrganizationsDiscovery:
                     "performance_benchmark": performance_benchmark_dict,
                     "timestamp": datetime.now().isoformat(),
                 }
-                
+
                 # Cache the successful result to prevent duplicate calls
                 _set_global_organizations_cache(discovery_result)
-                
+
                 return discovery_result
 
             except Exception as e:
@@ -606,7 +616,7 @@ class EnhancedOrganizationsDiscovery:
     async def _discover_accounts_optimized(self, optimization_engine) -> Dict:
         """
         Optimized account discovery using performance optimization engine
-        
+
         Addresses: Organization Discovery Performance (52.3s -> <30s target)
         Features:
         - Intelligent caching with TTL management
@@ -615,46 +625,44 @@ class EnhancedOrganizationsDiscovery:
         - Memory-efficient processing
         """
         logger.info("📊 Discovering organization accounts with SRE optimization patterns")
-        
+
         # Use optimization engine for discovery
         optimized_discover_accounts = optimization_engine.optimize_organization_discovery(
-            management_profile=self.management_profile,
-            use_parallel_processing=True,
-            batch_size=20
+            management_profile=self.management_profile, use_parallel_processing=True, batch_size=20
         )
-        
+
         # Execute optimized discovery
         try:
             result = optimized_discover_accounts()
-            
+
             # Convert to expected format
-            accounts_data = result.get('accounts', [])
-            
+            accounts_data = result.get("accounts", [])
+
             # Create AWSAccount objects for compatibility
             for account_data in accounts_data:
                 account = AWSAccount(
                     account_id=account_data["Id"],
-                    name=account_data["Name"], 
+                    name=account_data["Name"],
                     email=account_data["Email"],
                     status=account_data["Status"],
                     joined_method=account_data["JoinedMethod"],
                     joined_timestamp=account_data.get("JoinedTimestamp"),
-                    tags=account_data.get("Tags", {})
+                    tags=account_data.get("Tags", {}),
                 )
                 self.accounts_cache[account.account_id] = account
-            
+
             # Update metrics
             self.discovery_metrics["accounts_discovered"] = len(accounts_data)
-            
+
             # Enhanced account categorization
             active_accounts = [a for a in accounts_data if a.get("Status") == "ACTIVE"]
-            suspended_accounts = [a for a in accounts_data if a.get("Status") == "SUSPENDED"] 
+            suspended_accounts = [a for a in accounts_data if a.get("Status") == "SUSPENDED"]
             closed_accounts = [a for a in accounts_data if a.get("Status") == "CLOSED"]
-            
-            optimization_info = result.get('optimizations_applied', [])
+
+            optimization_info = result.get("optimizations_applied", [])
             logger.info(f"✅ Optimized discovery: {len(accounts_data)} accounts ({len(active_accounts)} active)")
             logger.info(f"🚀 Optimizations applied: {', '.join(optimization_info)}")
-            
+
             return {
                 "total_accounts": len(accounts_data),
                 "active_accounts": len(active_accounts),
@@ -665,7 +673,7 @@ class EnhancedOrganizationsDiscovery:
                 "profile_used": "management",
                 "optimizations_applied": optimization_info,
             }
-            
+
         except Exception as e:
             logger.error(f"Optimized account discovery failed: {e}")
             # Fallback to original method

runbooks/inventory/recover_cfn_stack_ids.py

@@ -10,7 +10,7 @@ import Inventory_Modules
 import simplejson as json
 from account_class import aws_acct_access
 from ArgumentsClass import CommonArguments
-from colorama import Fore, init
+from runbooks.common.rich_utils import console
 from Inventory_Modules import get_credentials_for_accounts_in_org
 
 """
@@ -18,7 +18,6 @@ This script was created to help solve a testing problem for the "move_stack_inst
 Originally, that script didn't have built-in recovery, so we needed this script to "recover" those stack-instance ids that might have been lost during the move_stack_instances.py run. However, that script now has built-in recovery, so this script isn't really needed. However, it can still be used to find any stack-instances that have been orphaned from their original stack-set, if that happens. 
 """
 
-init()
 __version__ = "2024.05.18"
 
 
@@ -70,9 +69,9 @@ def setup_auth_and_regions(fProfile: str = None) -> (object, list, list):
     ChildAccounts = Inventory_Modules.RemoveCoreAccounts(ChildAccounts, pAccountsToSkip)
     AccountList = [account["AccountId"] for account in ChildAccounts]
 
-    print(f"You asked to find stacks with this fragment list: {Fore.RED}{pFragments}{Fore.RESET}")
-    print(f"\t\tin these accounts: {Fore.RED}{AccountList}{Fore.RESET}")
-    print(f"\t\tin these regions: {Fore.RED}{RegionList}{Fore.RESET}")
+    print(f"You asked to find stacks with this fragment list: [red]{pFragments}")
+    print(f"\t\tin these accounts: [red]{AccountList}")
+    print(f"\t\tin these regions: [red]{RegionList}")
     return aws_acct, AccountList, RegionList
 
 
@@ -129,7 +128,7 @@ if __name__ == "__main__":
                 Stacks = Inventory_Modules.find_stacks2(cred, region, pFragments)
                 logging.warning(f"Account: {cred['AccountId']} | Region: {region} | Found {len(Stacks)} Stacks")
                 print(
-                    f"{ERASE_LINE}{Fore.RED}Account: {cred['AccountId']} Region: {region} Found {len(Stacks)} Stacks{Fore.RESET} ({item_counter} of {len(AccountList) * len(RegionList)})",
+                    f"{ERASE_LINE}[red]Account: {cred['AccountId']} Region: {region} Found {len(Stacks)} Stacks ({item_counter} of {len(AccountList) * len(RegionList)})",
                     end="\r",
                 )
             except Exception as my_Error:
@@ -162,7 +161,7 @@ if __name__ == "__main__":
         lAccountsAndRegions.append((StacksFound[i]["Account"], StacksFound[i]["Region"]))
     print(ERASE_LINE)
     print(
-        f"{Fore.RED}Looked through {len(StacksFound)} Stacks across {len(AccountList)} accounts across {len(RegionList)} regions{Fore.RESET}"
+        f"[red]Looked through {len(StacksFound)} Stacks across {len(AccountList)} accounts across {len(RegionList)} regions"
     )
     print()
     if args.loglevel < 21:  # INFO level
@@ -198,7 +197,7 @@ if __name__ == "__main__":
 
     if pTiming:
         print(ERASE_LINE)
-        print(f"{Fore.GREEN}This script took {time() - begin_time:.2f} seconds{Fore.RESET}")
+        print(f"[green]This script took {time() - begin_time:.2f} seconds")
 
     print()
     print("Thanks for using this script...")

runbooks/inventory/requirements.txt

@@ -1,7 +1,6 @@
 boto3>=1.26.84
 botocore>=1.29.84
 colorama>=0.4.6
-prettytable>=3.6.0
 simplejson>=3.18.3
 urllib3>=1.26.14
 randominfo>=2.0.2

runbooks/inventory/rich_inventory_display.py

@@ -21,9 +21,9 @@ from datetime import datetime
 from typing import Any, Dict, List, Optional
 
 from rich import box
-from rich.console import Console
 from rich.panel import Panel
-from rich.progress import BarColumn, Progress, SpinnerColumn, TaskProgressColumn, TextColumn
+from rich.progress import BarColumn, SpinnerColumn, TaskProgressColumn, TextColumn
+from runbooks.common.rich_utils import Progress
 from rich.table import Table
 from rich.text import Text
 from rich.tree import Tree
@@ -349,18 +349,18 @@ def display_account_tree(accounts_data: Dict[str, Dict]) -> None:
 
 
 def display_results_rich(
-    results_list: List[Dict[str, Any]], 
-    fdisplay_dict: Dict[str, Dict], 
-    defaultAction: Any = None, 
-    file_to_save: Optional[str] = None, 
+    results_list: List[Dict[str, Any]],
+    fdisplay_dict: Dict[str, Dict],
+    defaultAction: Any = None,
+    file_to_save: Optional[str] = None,
     subdisplay: bool = False,
-    title: str = "Inventory Results"
+    title: str = "Inventory Results",
 ) -> None:
     """
     Rich CLI replacement for legacy display_results function.
-    
+
     Provides backwards-compatible interface while using Rich formatting.
-    
+
     Args:
         results_list: List of dictionaries with resource data
         fdisplay_dict: Display configuration dictionary with format:
@@ -369,7 +369,7 @@ def display_results_rich(
         file_to_save: Optional filename to save results
         subdisplay: Whether this is a sub-display (affects formatting)
         title: Title for the table display
-    
+
     Example:
         display_dict = {
             'AccountId': {'DisplayOrder': 1, 'Heading': 'Account'},
@@ -379,59 +379,60 @@ def display_results_rich(
         display_results_rich(instance_data, display_dict, title="EC2 Instances")
     """
     from datetime import datetime
-    
+
     if not results_list:
         print_info("ℹ️ No results to display")
         return
-        
+
     # Sort display fields by DisplayOrder
-    sorted_fields = sorted(fdisplay_dict.items(), key=lambda x: x[1].get('DisplayOrder', 999))
-    
+    sorted_fields = sorted(fdisplay_dict.items(), key=lambda x: x[1].get("DisplayOrder", 999))
+
     # Create Rich table
     table = create_table(
         title=f"📊 {title}",
-        caption=f"Found {len(results_list)} results • {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+        caption=f"Found {len(results_list)} results • {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
     )
-    
+
     # Add columns based on display dictionary
     for field_name, field_config in sorted_fields:
-        heading = field_config.get('Heading', field_name)
+        heading = field_config.get("Heading", field_name)
         table.add_column(heading, style="cyan" if "Id" in field_name else "white", no_wrap=True)
-    
+
     # Add rows
     for result in results_list[:100]:  # Limit to first 100 for performance
         row_data = []
-        
+
         for field_name, field_config in sorted_fields:
             # Apply condition filter if specified
-            condition = field_config.get('Condition', [])
+            condition = field_config.get("Condition", [])
             if condition:
                 value = result.get(field_name, defaultAction)
                 if value not in condition:
                     continue
-                    
+
             # Get field value with default fallback
             value = result.get(field_name, defaultAction)
             if value is None:
                 value = "N/A"
-            
+
             # Format value as string
             row_data.append(str(value)[:50])  # Truncate long values
-        
+
         table.add_row(*row_data)
-    
+
     # Display the table
     console.print(table)
-    
+
     # Show truncation notice if needed
     if len(results_list) > 100:
         console.print(f"[dim]Showing first 100 results. Total found: {len(results_list)}[/dim]")
-    
+
     # Save to file if requested
     if file_to_save:
         try:
             import json
-            with open(file_to_save, 'w') as f:
+
+            with open(file_to_save, "w") as f:
                 json.dump(results_list, f, indent=2, default=str)
             print_success(f"💾 Results saved to: {file_to_save}")
         except Exception as e:
@@ -441,7 +442,7 @@ def display_results_rich(
 def display_progress_rich(current: int, total: int, description: str = "Processing") -> None:
     """
     Rich CLI replacement for legacy progress indicators.
-    
+
     Args:
         current: Current progress count
         total: Total items to process
@@ -454,21 +455,21 @@ def display_progress_rich(current: int, total: int, description: str = "Processi
 def print_colorized_rich(text: str, color: str = "white") -> None:
     """
     Rich CLI replacement for colorama print statements.
-    
+
     Args:
         text: Text to print
         color: Color name (red, green, yellow, cyan, blue, white)
     """
     color_map = {
         "red": "red",
-        "green": "green", 
+        "green": "green",
         "yellow": "yellow",
         "cyan": "cyan",
         "blue": "blue",
         "white": "white",
-        "magenta": "magenta"
+        "magenta": "magenta",
     }
-    
+
     rich_color = color_map.get(color.lower(), "white")
     console.print(f"[{rich_color}]{text}[/{rich_color}]")
 
@@ -476,13 +477,13 @@ def print_colorized_rich(text: str, color: str = "white") -> None:
 # Export public functions
 __all__ = [
     "display_inventory_header",
-    "create_inventory_progress", 
+    "create_inventory_progress",
     "display_ec2_inventory_results",
     "display_generic_inventory_results",
     "display_inventory_error",
     "display_multi_resource_summary",
     "display_account_tree",
     "display_results_rich",
-    "display_progress_rich", 
+    "display_progress_rich",
     "print_colorized_rich",
 ]

runbooks/inventory/run_on_multi_accounts.py

@@ -7,14 +7,12 @@ import Inventory_Modules
 from account_class import aws_acct_access
 from ArgumentsClass import CommonArguments
 from botocore.exceptions import ClientError
-from colorama import Fore, init
-from rich.console import Console
+from runbooks.common.rich_utils import console
 from rich.panel import Panel
 
-# Initialize Rich console
-console = Console()
+# Initialize Rich console with test mode support
+from runbooks.common.rich_utils import console
 
-init()
 __version__ = "2023.05.04"
 
 parser = CommonArguments()