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

src/aws/client.py

@@ -0,0 +1,128 @@
+"""Boto3 client wrapper with retry configuration and error handling."""
+
+import logging
+from typing import Any, Optional
+
+import boto3
+from botocore.config import Config
+from botocore.exceptions import ClientError, NoCredentialsError
+
+logger = logging.getLogger(__name__)
+
+
+# Aggressive retry configuration for batch operations
+DEFAULT_RETRY_CONFIG = Config(
+    retries={"max_attempts": 10, "mode": "adaptive"},  # Adaptive retry mode (boto3 1.16+)
+    max_pool_connections=50,
+    connect_timeout=60,
+    read_timeout=60,
+)
+
+
+def create_boto_client(
+    service_name: str,
+    region_name: str = "us-east-1",
+    profile_name: Optional[str] = None,
+    retry_config: Optional[Config] = None,
+) -> Any:
+    """Create boto3 client with retry configuration and error handling.
+
+    Args:
+        service_name: AWS service name (e.g., 'ec2', 'iam', 'lambda')
+        region_name: AWS region (default: 'us-east-1')
+        profile_name: AWS profile name from ~/.aws/config (optional)
+        retry_config: Custom botocore Config object (optional)
+
+    Returns:
+        Configured boto3 client
+
+    Raises:
+        NoCredentialsError: If AWS credentials are not found
+        ClientError: If client creation fails
+    """
+    if retry_config is None:
+        retry_config = DEFAULT_RETRY_CONFIG
+
+    try:
+        # Create session (with or without profile)
+        if profile_name:
+            session = boto3.Session(profile_name=profile_name)
+            client = session.client(service_name, region_name=region_name, config=retry_config)
+        else:
+            client = boto3.client(service_name, region_name=region_name, config=retry_config)
+
+        logger.debug(f"Created {service_name} client for region {region_name}")
+        return client
+
+    except NoCredentialsError:
+        logger.error("AWS credentials not found")
+        raise
+    except ClientError as e:
+        logger.error(f"Failed to create {service_name} client: {e}")
+        raise
+    except Exception as e:
+        logger.error(f"Unexpected error creating {service_name} client: {e}")
+        raise
+
+
+def get_enabled_regions(profile_name: Optional[str] = None) -> list[str]:  # type: ignore
+    """Get list of enabled AWS regions for the account.
+
+    Args:
+        profile_name: AWS profile name (optional)
+
+    Returns:
+        List of enabled region names
+    """
+    try:
+        ec2_client = create_boto_client("ec2", region_name="us-east-1", profile_name=profile_name)
+        response = ec2_client.describe_regions(AllRegions=False)  # Only enabled regions
+        regions = [region["RegionName"] for region in response["Regions"]]
+        logger.info(f"Found {len(regions)} enabled regions")
+        return regions
+    except Exception as e:
+        logger.warning(f"Could not retrieve enabled regions: {e}")
+        # Return default set of common regions
+        return [
+            "us-east-1",
+            "us-east-2",
+            "us-west-1",
+            "us-west-2",
+            "eu-west-1",
+            "eu-central-1",
+            "ap-southeast-1",
+            "ap-northeast-1",
+        ]
+
+
+def test_client_connection(client: Any) -> bool:
+    """Test if a boto3 client can connect to AWS.
+
+    Args:
+        client: Boto3 client instance
+
+    Returns:
+        True if connection successful, False otherwise
+    """
+    try:
+        # Different services have different test methods
+        service_name = client._service_model.service_name
+
+        if service_name == "sts":
+            client.get_caller_identity()
+        elif service_name == "ec2":
+            client.describe_regions(MaxResults=1)
+        elif service_name == "iam":
+            client.list_users(MaxItems=1)
+        elif service_name == "lambda":
+            client.list_functions(MaxItems=1)
+        elif service_name == "s3":
+            client.list_buckets()
+        else:
+            # Generic test - just try to get service metadata
+            client._make_api_call("ListObjects", {})
+
+        return True
+    except Exception as e:
+        logger.debug(f"Client connection test failed: {e}")
+        return False

src/aws/credentials.py

@@ -0,0 +1,191 @@
+"""AWS credential validation and permission checking."""
+
+import logging
+from typing import Any, Dict, List, Optional
+
+import boto3
+from botocore.exceptions import ClientError, NoCredentialsError, PartialCredentialsError
+
+logger = logging.getLogger(__name__)
+
+
+class CredentialValidationError(Exception):
+    """Raised when AWS credentials are invalid or missing."""
+
+    pass
+
+
+def validate_credentials(profile_name: Optional[str] = None) -> Dict[str, Any]:
+    """Validate AWS credentials and return caller identity information.
+
+    Args:
+        profile_name: AWS profile name from ~/.aws/config (optional)
+
+    Returns:
+        Dictionary with account ID, user ID, and ARN
+
+    Raises:
+        CredentialValidationError: If credentials are invalid or missing
+    """
+    try:
+        if profile_name:
+            session = boto3.Session(profile_name=profile_name)
+            sts_client = session.client("sts")
+        else:
+            sts_client = boto3.client("sts")
+
+        # Get caller identity to validate credentials
+        identity = sts_client.get_caller_identity()
+
+        logger.debug(f"Validated credentials for account {identity['Account']}")
+
+        return {
+            "account_id": identity["Account"],
+            "user_id": identity["UserId"],
+            "arn": identity["Arn"],
+        }
+
+    except NoCredentialsError:
+        error_msg = (
+            "AWS credentials not found. Please configure credentials using one of these methods:\n"
+            "  1. Run: aws configure\n"
+            "  2. Set environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY\n"
+            "  3. Use --profile option with a configured profile\n\n"
+            "For more info: https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html"
+        )
+        logger.error("No AWS credentials found")
+        raise CredentialValidationError(error_msg)
+
+    except PartialCredentialsError as e:
+        error_msg = f"Incomplete AWS credentials: {e}"
+        logger.error(error_msg)
+        raise CredentialValidationError(error_msg)
+
+    except ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code", "Unknown")
+        if error_code == "InvalidClientTokenId":
+            error_msg = "AWS credentials are invalid. Please check your access key ID."
+        elif error_code == "SignatureDoesNotMatch":
+            error_msg = "AWS credentials signature mismatch. Please check your secret access key."
+        elif error_code == "ExpiredToken":
+            error_msg = "AWS credentials have expired. Please refresh your temporary credentials."
+        else:
+            error_msg = f"AWS credential validation failed: {e}"
+
+        logger.error(error_msg)
+        raise CredentialValidationError(error_msg)
+
+    except Exception as e:
+        error_msg = f"Unexpected error validating credentials: {e}"
+        logger.error(error_msg)
+        raise CredentialValidationError(error_msg)
+
+
+def check_required_permissions(
+    profile_name: Optional[str] = None, required_actions: Optional[List[str]] = None
+) -> Dict[str, bool]:
+    """Check if credentials have required IAM permissions.
+
+    Note: This is a best-effort check using IAM policy simulation.
+    Some permissions may not be accurately detected.
+
+    Args:
+        profile_name: AWS profile name (optional)
+        required_actions: List of IAM actions to check (e.g., ['ec2:DescribeInstances'])
+
+    Returns:
+        Dictionary mapping action names to permission status (True/False)
+    """
+    if required_actions is None:
+        # Default minimum required permissions for snapshot operations
+        required_actions = [
+            "ec2:DescribeInstances",
+            "ec2:DescribeRegions",
+            "iam:ListRoles",
+            "lambda:ListFunctions",
+            "s3:ListAllMyBuckets",
+        ]
+
+    try:
+        # Get caller identity first
+        identity = validate_credentials(profile_name)
+
+        if profile_name:
+            session = boto3.Session(profile_name=profile_name)
+            iam_client = session.client("iam")
+        else:
+            iam_client = boto3.client("iam")
+
+        results = {}
+
+        # Try to simulate policy for each action
+        for action in required_actions:
+            try:
+                response = iam_client.simulate_principal_policy(
+                    PolicySourceArn=identity["arn"],
+                    ActionNames=[action],
+                )
+
+                # Check if action is allowed
+                eval_results = response.get("EvaluationResults", [])
+                if eval_results:
+                    decision = eval_results[0].get("EvalDecision", "deny")
+                    results[action] = decision.lower() == "allowed"
+                else:
+                    results[action] = False
+
+            except ClientError as e:
+                # If simulation fails (e.g., lack of iam:SimulatePrincipalPolicy permission),
+                # we can't determine the permission status
+                logger.debug(f"Could not check permission for {action}: {e}")
+                results[action] = None  # Unknown
+
+        return results
+
+    except Exception as e:
+        logger.warning(f"Permission check failed: {e}")
+        return {action: None for action in required_actions}  # type: ignore
+
+
+def get_account_id(profile_name: Optional[str] = None) -> str:
+    """Get AWS account ID for the current credentials.
+
+    Args:
+        profile_name: AWS profile name (optional)
+
+    Returns:
+        12-digit AWS account ID
+
+    Raises:
+        CredentialValidationError: If credentials are invalid
+    """
+    identity = validate_credentials(profile_name)
+    return identity["account_id"]
+
+
+def get_credential_summary(profile_name: Optional[str] = None) -> str:
+    """Get a human-readable summary of current AWS credentials.
+
+    Args:
+        profile_name: AWS profile name (optional)
+
+    Returns:
+        Formatted string with credential information
+    """
+    try:
+        identity = validate_credentials(profile_name)
+
+        summary = f"""
+AWS Credentials Valid
+Account ID: {identity['account_id']}
+User/Role: {identity['arn'].split('/')[-1]}
+ARN: {identity['arn']}
+"""
+
+        if profile_name:
+            summary += f"Profile: {profile_name}\n"
+
+        return summary.strip()
+
+    except CredentialValidationError as e:
+        return f"Credential Validation Failed:\n{e}"

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,5 @@
+"""CLI module for AWS Baseline Snapshot tool."""
+
+from .main import app, cli_main
+
+__all__ = ["app", "cli_main"]

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, 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,
+        }