argus-cloud-optimizer 0.2.0__py3-none-any.whl

core/secrets.py

@@ -0,0 +1,145 @@
+"""
+Secret manager integration — resolve secret references in environment variables.
+
+If an env var's value matches a secret reference pattern, the real value is
+fetched from the corresponding cloud secret manager and the env var is updated
+in-place before the config layer reads it.
+
+Supported patterns:
+    arn:aws:secretsmanager:<region>:<account>:secret:<name>
+        → AWS Secrets Manager
+    gcp-secret://<project>/<secret-name>[/<version>]
+        → GCP Secret Manager
+    akv://<vault-name>/<secret-name>
+        → Azure Key Vault
+
+Call ``resolve_secrets()`` once at startup, before ``validate_environment()``.
+It is safe to call when no secret references exist — it's a no-op.
+
+This module lives in core/ but imports cloud SDKs lazily inside the resolver
+functions (only when a matching reference is found). If the required SDK is
+not installed, a clear error is raised.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+_SECRET_VARS = (
+    "ANTHROPIC_API_KEY",
+    "SLACK_WEBHOOK_URL",
+    "TEAMS_WEBHOOK_URL",
+    "WEBHOOK_URL",
+    "AZURE_OPENAI_API_KEY",
+    "AZURE_OPENAI_ENDPOINT",
+)
+
+_AWS_ARN_PATTERN = re.compile(
+    r"^arn:aws:secretsmanager:[\w-]+:\d+:secret:.+"
+)
+_GCP_PATTERN = re.compile(r"^gcp-secret://([^/]+)/([^/]+)(?:/([^/]+))?$")
+_AKV_PATTERN = re.compile(r"^akv://([^/]+)/(.+)$")
+
+
+def resolve_secrets() -> None:
+    """
+    Scan secret-eligible env vars for reference patterns and resolve them.
+
+    Updates ``os.environ`` in-place so downstream code (config layer,
+    validation, providers) sees the real values transparently.
+    """
+    for var in _SECRET_VARS:
+        value = os.environ.get(var, "")
+        if not value:
+            continue
+
+        resolved = _try_resolve(var, value)
+        if resolved is not None:
+            os.environ[var] = resolved
+            logger.info("secret_resolved", var=var)
+
+
+def _try_resolve(var: str, value: str) -> str | None:
+    """Return the resolved secret value, or None if the value is not a reference."""
+    if _AWS_ARN_PATTERN.match(value):
+        return _resolve_aws(var, value)
+    match = _GCP_PATTERN.match(value)
+    if match:
+        project, name, version = match.groups()
+        return _resolve_gcp(var, project, name, version or "latest")
+    match = _AKV_PATTERN.match(value)
+    if match:
+        vault, name = match.groups()
+        return _resolve_azure(var, vault, name)
+    return None
+
+
+def _resolve_aws(var: str, arn: str) -> str:
+    try:
+        import boto3
+        from botocore.exceptions import ClientError
+    except ImportError:
+        raise ImportError(
+            f"{var} references AWS Secrets Manager ({arn}) but boto3 is not "
+            "installed. Install with: pip install boto3"
+        ) from None
+
+    region = arn.split(":")[3]
+    client = boto3.client("secretsmanager", region_name=region)
+    try:
+        resp = client.get_secret_value(SecretId=arn)
+    except ClientError as exc:
+        raise RuntimeError(
+            f"Failed to resolve {var} from AWS Secrets Manager: {exc}"
+        ) from exc
+    return resp["SecretString"]
+
+
+def _resolve_gcp(var: str, project: str, name: str, version: str) -> str:
+    try:
+        from google.cloud import secretmanager
+    except ImportError:
+        raise ImportError(
+            f"{var} references GCP Secret Manager "
+            f"(gcp-secret://{project}/{name}/{version}) but "
+            "google-cloud-secret-manager is not installed. "
+            "Install with: pip install google-cloud-secret-manager"
+        ) from None
+
+    client = secretmanager.SecretManagerServiceClient()
+    resource = f"projects/{project}/secrets/{name}/versions/{version}"
+    try:
+        resp = client.access_secret_version(request={"name": resource})
+    except Exception as exc:
+        raise RuntimeError(
+            f"Failed to resolve {var} from GCP Secret Manager: {exc}"
+        ) from exc
+    return resp.payload.data.decode("utf-8")
+
+
+def _resolve_azure(var: str, vault_name: str, secret_name: str) -> str:
+    try:
+        from azure.identity import DefaultAzureCredential
+        from azure.keyvault.secrets import SecretClient
+    except ImportError:
+        raise ImportError(
+            f"{var} references Azure Key Vault "
+            f"(akv://{vault_name}/{secret_name}) but azure-keyvault-secrets "
+            "is not installed. Install with: pip install azure-keyvault-secrets "
+            "azure-identity"
+        ) from None
+
+    vault_url = f"https://{vault_name}.vault.azure.net"
+    client = SecretClient(vault_url=vault_url, credential=DefaultAzureCredential())
+    try:
+        secret = client.get_secret(secret_name)
+    except Exception as exc:
+        raise RuntimeError(
+            f"Failed to resolve {var} from Azure Key Vault: {exc}"
+        ) from exc
+    return secret.value

core/token_tracker.py

@@ -0,0 +1,97 @@
+"""
+LLM token and cost tracking with hard budget enforcement.
+
+Tracks cumulative input/output tokens across agent iterations and estimates
+USD cost using per-provider pricing. When ``LLM_BUDGET_USD`` is exceeded,
+raises ``BudgetExceededError`` so the agent loop can abort gracefully and
+still return partial findings.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import structlog
+
+logger = structlog.get_logger(__name__)
+
+# Per-million-token pricing (input, output) by provider.
+# Updated 2025-05 — check provider pricing pages for current rates.
+_PRICING: dict[str, tuple[float, float]] = {
+    "anthropic": (3.0, 15.0),
+    "bedrock": (3.0, 15.0),
+    "vertexai": (1.25, 5.0),
+    "azure_openai": (2.50, 10.0),
+}
+
+_DEFAULT_PRICING = (3.0, 15.0)
+
+
+class BudgetExceededError(Exception):
+    """Raised when cumulative LLM cost exceeds the configured budget."""
+
+    def __init__(self, spent_usd: float, budget_usd: float) -> None:
+        self.spent_usd = spent_usd
+        self.budget_usd = budget_usd
+        super().__init__(
+            f"LLM budget exceeded: ${spent_usd:.4f} spent "
+            f"(budget: ${budget_usd:.2f})"
+        )
+
+
+@dataclass
+class TokenTracker:
+    """Accumulates token usage and enforces a hard USD budget."""
+
+    budget_usd: float
+    provider: str = "anthropic"
+
+    total_input_tokens: int = field(default=0, init=False)
+    total_output_tokens: int = field(default=0, init=False)
+    iteration_count: int = field(default=0, init=False)
+    _per_iteration: list[dict[str, int]] = field(default_factory=list, init=False)
+
+    def record(self, input_tokens: int, output_tokens: int) -> None:
+        """
+        Record tokens from one AI call and check the budget.
+
+        Raises ``BudgetExceededError`` if cumulative cost exceeds budget.
+        """
+        self.total_input_tokens += input_tokens
+        self.total_output_tokens += output_tokens
+        self.iteration_count += 1
+        self._per_iteration.append(
+            {"input": input_tokens, "output": output_tokens}
+        )
+
+        spent = self.estimated_cost_usd
+        logger.info(
+            "token_usage",
+            iteration=self.iteration_count,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            cumulative_input=self.total_input_tokens,
+            cumulative_output=self.total_output_tokens,
+            spent_usd=round(spent, 4),
+            budget_usd=self.budget_usd,
+        )
+
+        if self.budget_usd > 0 and spent > self.budget_usd:
+            raise BudgetExceededError(round(spent, 4), self.budget_usd)
+
+    @property
+    def estimated_cost_usd(self) -> float:
+        input_rate, output_rate = _PRICING.get(self.provider, _DEFAULT_PRICING)
+        cost = (self.total_input_tokens / 1_000_000 * input_rate) + (
+            self.total_output_tokens / 1_000_000 * output_rate
+        )
+        return round(cost, 4)
+
+    def summary(self) -> dict[str, float | int]:
+        return {
+            "total_input_tokens": self.total_input_tokens,
+            "total_output_tokens": self.total_output_tokens,
+            "iterations": self.iteration_count,
+            "estimated_cost_usd": self.estimated_cost_usd,
+            "budget_usd": self.budget_usd,
+        }

core/validation.py

@@ -0,0 +1,214 @@
+"""
+Startup environment validation.
+
+Called once at the top of each entrypoint before any cloud API calls are made.
+Raises ConfigurationError with a clear, actionable message if required env vars
+are missing or malformed. This prevents wasting a 15-minute scan that would
+fail at the very end due to a missing credential.
+
+No cloud SDK imports here — this is core/ and must stay cloud-free.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import urllib.parse
+
+
+class ConfigurationError(Exception):
+    """Raised when the environment is misconfigured at startup."""
+
+
+def validate_environment(cloud: str) -> None:
+    """
+    Validate all required environment variables for the given cloud.
+
+    Args:
+        cloud: "aws" | "gcp" | "azure"
+
+    Raises:
+        ConfigurationError: with a human-readable message describing every
+            problem found (not just the first one).
+    """
+    errors: list[str] = []
+
+    _check_ai_provider(errors)
+    _check_slack(errors)
+
+    if cloud == "aws":
+        _check_aws(errors)
+    elif cloud == "gcp":
+        _check_gcp(errors)
+    elif cloud == "azure":
+        _check_azure(errors)
+
+    if errors:
+        bullet_list = "\n".join(f"  • {e}" for e in errors)
+        raise ConfigurationError(
+            f"Argus cannot start — {len(errors)} configuration error(s) found:\n"
+            f"{bullet_list}\n\n"
+            "Fix the above and re-deploy or re-run."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Shared checks
+# ---------------------------------------------------------------------------
+
+
+def _check_ai_provider(errors: list[str]) -> None:
+    provider = os.environ.get("AI_PROVIDER", "").strip().lower()
+
+    # Each cloud has its own default, so an empty AI_PROVIDER is fine —
+    # the entrypoint will pick the cloud-native default.
+    if not provider:
+        return
+
+    known = {"anthropic", "bedrock", "vertexai", "azure_openai"}
+    if provider not in known:
+        errors.append(
+            f"AI_PROVIDER={provider!r} is not recognised. "
+            f"Valid values: {', '.join(sorted(known))}."
+        )
+        return  # No point checking credentials for an unknown provider.
+
+    if provider == "anthropic":
+        key = os.environ.get("ANTHROPIC_API_KEY", "").strip()
+        if not key:
+            errors.append(
+                "AI_PROVIDER=anthropic requires ANTHROPIC_API_KEY to be set. "
+                "Get a key at https://console.anthropic.com/settings/api-keys"
+            )
+        elif not key.startswith("sk-ant-"):
+            errors.append(
+                "ANTHROPIC_API_KEY looks malformed "
+                "(expected it to start with 'sk-ant-'). "
+                "Check the key in the Anthropic console."
+            )
+
+    if provider == "azure_openai":
+        endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT", "").strip()
+        if not endpoint:
+            errors.append(
+                "AI_PROVIDER=azure_openai requires AZURE_OPENAI_ENDPOINT to be set. "
+                "Example: https://<resource>.openai.azure.com/"
+            )
+        elif not _is_https_url(endpoint):
+            errors.append(
+                f"AZURE_OPENAI_ENDPOINT={endpoint!r} is not a valid HTTPS URL."
+            )
+
+
+def _check_slack(errors: list[str]) -> None:
+    dry_run = os.environ.get("DRY_RUN", "false").lower() in ("true", "1", "yes")
+    if dry_run:
+        return  # Webhook not needed in dry-run mode.
+
+    url = os.environ.get("SLACK_WEBHOOK_URL", "").strip()
+    if not url:
+        errors.append(
+            "SLACK_WEBHOOK_URL is not set. "
+            "Create an incoming webhook at https://api.slack.com/apps "
+            "or set DRY_RUN=true to skip Slack delivery."
+        )
+        return
+
+    if not _is_https_url(url):
+        errors.append(
+            f"SLACK_WEBHOOK_URL={url!r} is not a valid HTTPS URL. "
+            "Expected format: https://hooks.slack.com/services/..."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Cloud-specific checks
+# ---------------------------------------------------------------------------
+
+
+def _check_aws(errors: list[str]) -> None:
+    accounts_mode = os.environ.get("ACCOUNTS_MODE", "single").lower()
+
+    if accounts_mode not in ("single", "multi"):
+        errors.append(
+            f"ACCOUNTS_MODE={accounts_mode!r} is not valid. " "Use 'single' or 'multi'."
+        )
+
+    if accounts_mode == "multi":
+        raw = os.environ.get("ACCOUNTS_CONFIG", "").strip()
+        if not raw:
+            errors.append(
+                "ACCOUNTS_MODE=multi requires ACCOUNTS_CONFIG to be set. "
+                "Set it to a JSON array of account objects: "
+                '[{"id":"123456789012","name":"prod","role_arn":"arn:aws:iam::..."}]'
+            )
+            return
+
+        try:
+            accounts = json.loads(raw)
+        except json.JSONDecodeError as exc:
+            errors.append(
+                f"ACCOUNTS_CONFIG is not valid JSON: {exc}. "
+                "Expected a JSON array of account objects."
+            )
+            return
+
+        if not isinstance(accounts, list) or len(accounts) == 0:
+            errors.append(
+                "ACCOUNTS_CONFIG must be a non-empty JSON array of account objects."
+            )
+            return
+
+        for i, acct in enumerate(accounts):
+            if not isinstance(acct, dict):
+                errors.append(
+                    f"ACCOUNTS_CONFIG[{i}] is not an object — each account must be "
+                    '{"id": "...", "name": "...", "role_arn": "..."}.'
+                )
+                continue
+            missing = [f for f in ("id", "role_arn") if not acct.get(f)]
+            if missing:
+                name = acct.get("name", f"index {i}")
+                errors.append(
+                    f"ACCOUNTS_CONFIG account '{name}' is missing required "
+                    f"field(s): {', '.join(missing)}."
+                )
+
+
+def _check_gcp(errors: list[str]) -> None:
+    project_id = os.environ.get("GCP_PROJECT_ID", "").strip()
+    if not project_id:
+        errors.append(
+            "GCP_PROJECT_ID is required for GCP scans. "
+            "Set it to your GCP project ID (e.g. my-project-123)."
+        )
+
+
+def _check_azure(errors: list[str]) -> None:
+    raw = os.environ.get("AZURE_SUBSCRIPTION_IDS", "").strip()
+    if not raw:
+        errors.append(
+            "AZURE_SUBSCRIPTION_IDS is required for Azure scans. "
+            "Set it to one or more subscription IDs separated by commas."
+        )
+        return
+
+    subscription_ids = [s.strip() for s in raw.split(",") if s.strip()]
+    if not subscription_ids:
+        errors.append(
+            "AZURE_SUBSCRIPTION_IDS is set but contains no valid IDs. "
+            "Expected one or more GUIDs separated by commas."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _is_https_url(value: str) -> bool:
+    try:
+        parsed = urllib.parse.urlparse(value)
+        return parsed.scheme == "https" and bool(parsed.netloc)
+    except ValueError:
+        return False