aws-inventory-manager 0.17.12__py3-none-any.whl

src/aws/rate_limiter.py

@@ -0,0 +1,177 @@
+"""Rate limiter utility using token bucket algorithm."""
+
+import logging
+import time
+from threading import Lock
+from typing import Dict, Optional
+
+logger = logging.getLogger(__name__)
+
+
+# Service-specific rate limits (calls per second)
+# Based on AWS API throttling limits
+SERVICE_RATE_LIMITS: Dict[str, float] = {
+    "iam": 5.0,  # IAM has strict rate limits (global service)
+    "cloudformation": 2.0,  # CloudFormation is particularly slow
+    "sts": 10.0,  # STS is also rate-limited
+    "default": 10.0,  # Conservative default for other services
+}
+
+
+class RateLimiter:
+    """Token bucket rate limiter for controlling API call frequency.
+
+    This prevents hitting AWS API rate limits by throttling client-side
+    before the request is even made.
+    """
+
+    def __init__(self, rate: float):
+        """Initialize rate limiter.
+
+        Args:
+            rate: Maximum number of calls per second
+        """
+        self.rate = rate
+        self.tokens = rate
+        self.last_update = time.time()
+        self.lock = Lock()
+
+        logger.debug(f"Initialized rate limiter with rate {rate} calls/sec")
+
+    def acquire(self, blocking: bool = True) -> bool:
+        """Acquire permission to make an API call.
+
+        This method will block until a token is available (if blocking=True)
+        or return immediately (if blocking=False).
+
+        Args:
+            blocking: If True, wait until token available. If False, return immediately.
+
+        Returns:
+            True if token acquired, False if blocking=False and no token available
+        """
+        with self.lock:
+            now = time.time()
+            elapsed = now - self.last_update
+
+            # Refill tokens based on elapsed time
+            self.tokens = min(self.rate, self.tokens + elapsed * self.rate)
+            self.last_update = now
+
+            if self.tokens >= 1:
+                # Token available
+                self.tokens -= 1
+                return True
+            else:
+                # No token available
+                if not blocking:
+                    return False
+
+                # Calculate how long to sleep
+                sleep_time = (1 - self.tokens) / self.rate
+                logger.debug(f"Rate limiter sleeping for {sleep_time:.3f}s")
+
+        # Sleep outside the lock to allow other threads
+        time.sleep(sleep_time)
+
+        # Acquire token after sleeping
+        with self.lock:
+            self.tokens = 0
+            self.last_update = time.time()
+            return True
+
+    def try_acquire(self) -> bool:
+        """Try to acquire a token without blocking.
+
+        Returns:
+            True if token acquired, False otherwise
+        """
+        return self.acquire(blocking=False)
+
+
+class ServiceRateLimiter:
+    """Manages rate limiters for different AWS services."""
+
+    def __init__(self, rate_limits: Optional[Dict[str, float]] = None):
+        """Initialize service rate limiter.
+
+        Args:
+            rate_limits: Dictionary mapping service names to rates (optional)
+        """
+        self.rate_limits = rate_limits or SERVICE_RATE_LIMITS
+        self._limiters: Dict[str, RateLimiter] = {}
+        self._lock = Lock()
+
+    def get_limiter(self, service_name: str) -> RateLimiter:
+        """Get or create a rate limiter for a service.
+
+        Args:
+            service_name: AWS service name (e.g., 'iam', 'ec2')
+
+        Returns:
+            RateLimiter instance for the service
+        """
+        with self._lock:
+            if service_name not in self._limiters:
+                rate = self.rate_limits.get(service_name, self.rate_limits["default"])
+                self._limiters[service_name] = RateLimiter(rate)
+                logger.debug(f"Created rate limiter for {service_name} ({rate} calls/sec)")
+
+            return self._limiters[service_name]
+
+    def acquire(self, service_name: str, blocking: bool = True) -> bool:
+        """Acquire permission to call an AWS service API.
+
+        Args:
+            service_name: AWS service name
+            blocking: Whether to block until token available
+
+        Returns:
+            True if token acquired
+        """
+        limiter = self.get_limiter(service_name)
+        return limiter.acquire(blocking=blocking)
+
+    def try_acquire(self, service_name: str) -> bool:
+        """Try to acquire permission without blocking.
+
+        Args:
+            service_name: AWS service name
+
+        Returns:
+            True if token acquired, False otherwise
+        """
+        return self.acquire(service_name, blocking=False)
+
+
+# Global service rate limiter instance
+_global_limiter: Optional[ServiceRateLimiter] = None
+
+
+def get_global_rate_limiter() -> ServiceRateLimiter:
+    """Get the global service rate limiter instance.
+
+    Returns:
+        Global ServiceRateLimiter instance
+    """
+    global _global_limiter
+    if _global_limiter is None:
+        _global_limiter = ServiceRateLimiter()
+    return _global_limiter
+
+
+def rate_limited_call(service_name: str, func, *args, **kwargs):  # type: ignore[no-untyped-def]
+    """Execute a function with rate limiting applied.
+
+    Args:
+        service_name: AWS service name for rate limiting
+        func: Function to call
+        *args: Positional arguments to pass to func
+        **kwargs: Keyword arguments to pass to func
+
+    Returns:
+        Result of func(*args, **kwargs)
+    """
+    limiter = get_global_rate_limiter()
+    limiter.acquire(service_name)
+    return func(*args, **kwargs)

src/cli/__init__.py

@@ -0,0 +1,12 @@
+"""CLI module for AWS Baseline Snapshot tool."""
+
+__all__ = ["app", "cli_main"]
+
+
+def __getattr__(name: str):
+    """Lazy import to avoid RuntimeWarning when running with python -m."""
+    if name in ("app", "cli_main"):
+        from .main import app, cli_main
+
+        return app if name == "app" else cli_main
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

src/cli/config.py

@@ -0,0 +1,130 @@
+"""Configuration loader for CLI."""
+
+import logging
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+logger = logging.getLogger(__name__)
+
+
+class Config:
+    """Application configuration manager."""
+
+    def __init__(self) -> None:
+        """Initialize configuration with defaults."""
+        self.snapshot_dir: str = ".snapshots"
+        self.storage_path: Optional[str] = None  # Runtime override for snapshot storage path
+        self.regions: List[str] = []  # Empty means all enabled regions
+        self.resource_types: List[str] = []  # Empty means all supported types
+        self.aws_profile: Optional[str] = None
+        self.parallel_workers: int = 10
+        self.auto_compress_mb: int = 10
+        self.log_level: str = "INFO"
+
+    @classmethod
+    def load(cls, config_file: Optional[str] = None) -> "Config":
+        """Load configuration from file and environment variables.
+
+        Priority (highest to lowest):
+        1. Environment variables
+        2. Config file
+        3. Defaults
+
+        Args:
+            config_file: Path to config file (optional)
+
+        Returns:
+            Config instance
+        """
+        config = cls()
+
+        # Try to load from config file
+        config_path = None
+        if config_file:
+            config_path = Path(config_file)
+        else:
+            # Try current directory first
+            config_path = Path(".aws-baseline.yaml")
+            if not config_path.exists():
+                # Try home directory
+                config_path = Path.home() / ".aws-baseline.yaml"
+
+        if config_path and config_path.exists():
+            config._load_from_file(config_path)
+
+        # Override with environment variables
+        config._load_from_env()
+
+        return config
+
+    def _load_from_file(self, config_path: Path) -> None:
+        """Load configuration from YAML file.
+
+        Args:
+            config_path: Path to config file
+        """
+        try:
+            with open(config_path, "r") as f:
+                data = yaml.safe_load(f)
+
+            if not data:
+                return
+
+            self.snapshot_dir = data.get("snapshot_dir", self.snapshot_dir)
+            self.regions = data.get("regions", self.regions)
+            self.aws_profile = data.get("aws_profile", self.aws_profile)
+            self.parallel_workers = data.get("parallel_workers", self.parallel_workers)
+            self.auto_compress_mb = data.get("auto_compress_mb", self.auto_compress_mb)
+
+            # Handle resource_types (include/exclude)
+            resource_types_config = data.get("resource_types", {})
+            if isinstance(resource_types_config, dict):
+                self.resource_types = resource_types_config.get("include", [])
+            elif isinstance(resource_types_config, list):
+                self.resource_types = resource_types_config
+
+            logger.info(f"Loaded configuration from {config_path}")
+
+        except Exception as e:
+            logger.warning(f"Could not load config file {config_path}: {e}")
+
+    def _load_from_env(self) -> None:
+        """Load configuration from environment variables."""
+        # AWS_BASELINE_SNAPSHOT_DIR
+        snapshot_dir = os.getenv("AWS_BASELINE_SNAPSHOT_DIR")
+        if snapshot_dir:
+            self.snapshot_dir = snapshot_dir
+
+        # AWS_BASELINE_LOG_LEVEL
+        log_level = os.getenv("AWS_BASELINE_LOG_LEVEL")
+        if log_level:
+            self.log_level = log_level
+
+        # AWS_PROFILE
+        aws_profile = os.getenv("AWS_PROFILE")
+        if aws_profile:
+            self.aws_profile = aws_profile
+
+        # AWS_REGION (single region from env)
+        aws_region = os.getenv("AWS_REGION")
+        if aws_region:
+            self.regions = [aws_region]
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert configuration to dictionary.
+
+        Returns:
+            Configuration as dictionary
+        """
+        return {
+            "snapshot_dir": self.snapshot_dir,
+            "regions": self.regions,
+            "resource_types": self.resource_types,
+            "aws_profile": self.aws_profile,
+            "parallel_workers": self.parallel_workers,
+            "auto_compress_mb": self.auto_compress_mb,
+            "log_level": self.log_level,
+        }