api/auth.py

"""
api/auth.py — API key authentication and role-based authorization.

Roles:
  admin           — all permissions
  compliance_lead — query, view_results, manage_alerts
  reader          — query, view_results
  auditor         — view_audit only

API keys are loaded from the env var COMPLIANCE_API_KEYS in the format:
  key1:role1:tenant1,key2:role2:tenant2

Default (for local dev) is set below.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass
from typing import Optional

from fastapi import HTTPException, status
from fastapi.security import APIKeyHeader

log = logging.getLogger(__name__)

# ─── Security scheme ──────────────────────────────────────────────────────────

API_KEY_HEADER = APIKeyHeader(name="X-API-Key", auto_error=False)

# ─── Role definitions ─────────────────────────────────────────────────────────

ROLE_PERMISSIONS: dict[str, set[str]] = {
    "admin": {"query", "view_results", "manage_alerts", "view_audit", "admin"},
    "compliance_lead": {"query", "view_results", "manage_alerts"},
    "reader": {"query", "view_results"},
    "auditor": {"view_audit"},
}

VALID_ROLES = set(ROLE_PERMISSIONS.keys())

# ─── User model ───────────────────────────────────────────────────────────────


@dataclass
class User:
    api_key: str
    role: str
    tenant_id: str

    @property
    def permissions(self) -> set[str]:
        return ROLE_PERMISSIONS.get(self.role, set())

    def has_permission(self, perm: str) -> bool:
        return perm in self.permissions


# ─── API key store ────────────────────────────────────────────────────────────

_DEFAULT_KEYS = (
    "admin-key-001:admin:tenant-default,"
    "reader-key-001:reader:tenant-default"
)

# Mutable runtime dict — admin can add keys at runtime via POST /admin/api-keys
_API_KEY_STORE: dict[str, tuple[str, str]] = {}  # key -> (role, tenant_id)


def _load_api_keys() -> None:
    """Parse COMPLIANCE_API_KEYS env var and populate _API_KEY_STORE."""
    raw = os.environ.get("COMPLIANCE_API_KEYS", _DEFAULT_KEYS)
    for entry in raw.split(","):
        entry = entry.strip()
        if not entry:
            continue
        parts = entry.split(":")
        if len(parts) != 3:
            log.warning("Skipping malformed API key entry (expected key:role:tenant): %r", entry)
            continue
        key, role, tenant = parts[0].strip(), parts[1].strip(), parts[2].strip()
        if role not in VALID_ROLES:
            log.warning("Skipping API key with unknown role %r", role)
            continue
        _API_KEY_STORE[key] = (role, tenant)
    log.info("Loaded %d API keys", len(_API_KEY_STORE))


# Load on module import
_load_api_keys()


def add_api_key(key: str, role: str, tenant_id: str) -> None:
    """Add or update an API key at runtime (admin action)."""
    if role not in VALID_ROLES:
        raise ValueError(f"Unknown role {role!r}. Valid roles: {sorted(VALID_ROLES)}")
    _API_KEY_STORE[key] = (role, tenant_id)
    log.info("API key added/updated: role=%s tenant=%s", role, tenant_id)


def list_api_keys() -> list[dict]:
    """Return sanitized list of keys (masked) for admin inspection."""
    result = []
    for key, (role, tenant) in _API_KEY_STORE.items():
        masked = key[:4] + "****" + key[-4:] if len(key) >= 8 else "****"
        result.append({"key_preview": masked, "role": role, "tenant_id": tenant})
    return result


# ─── FastAPI dependency ───────────────────────────────────────────────────────


def get_current_user(api_key: Optional[str] = None) -> User:
    """
    Validate X-API-Key header against _API_KEY_STORE.

    Call this with Security(get_current_user) where api_key comes from
    API_KEY_HEADER, or pass directly for testing.
    """
    if not api_key:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing X-API-Key header",
        )
    entry = _API_KEY_STORE.get(api_key)
    if entry is None:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Invalid API key",
        )
    role, tenant_id = entry
    return User(api_key=api_key, role=role, tenant_id=tenant_id)


def require_role(*roles: str):
    """
    FastAPI dependency factory — raises 403 if the user's role is not in roles.

    Usage:
        @app.get("/admin/...", dependencies=[Depends(require_role("admin"))])
    """
    from fastapi import Depends, Security

    def _dependency(api_key: Optional[str] = Security(API_KEY_HEADER)) -> User:
        user = get_current_user(api_key)
        if user.role not in roles:
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=f"Role '{user.role}' is not authorized. Required: {list(roles)}",
            )
        return user

    return Depends(_dependency)


def require_permission(permission: str):
    """
    FastAPI dependency factory — raises 403 if the user lacks the permission.

    Usage:
        @app.get("/audit", dependencies=[Depends(require_permission("view_audit"))])
    """
    from fastapi import Depends, Security

    def _dependency(api_key: Optional[str] = Security(API_KEY_HEADER)) -> User:
        user = get_current_user(api_key)
        if not user.has_permission(permission):
            raise HTTPException(
                status_code=status.HTTP_403_FORBIDDEN,
                detail=f"Permission '{permission}' required. Your role '{user.role}' does not have it.",
            )
        return user

    return Depends(_dependency)