lambdagent/agentpaas/api/middleware/auth.py

"""
api.middleware.auth — API Key authentication middleware.

Security: Only SHA-256 hash stored. Key shown once at creation.
Format: ap_xxxxxxxxxxxxxxxx (32 char random, prefix ap_)
"""
from __future__ import annotations
import hashlib
import time
from typing import Optional

from agentpaas.db.session import get_db
from agentpaas.db.models import now_utc


class AuthError(Exception):
    def __init__(self, message: str, code: int = 401):
        self.message = message
        self.code = code
        super().__init__(message)


class TenantContext:
    """Injected into request state after authentication."""
    def __init__(self, tenant_id: str, tenant_name: str, plan: str,
                 user_id: str = "", role: str = "developer",
                 scopes: list = None, rate_limit: int = 60):
        self.tenant_id = tenant_id
        self.tenant_name = tenant_name
        self.plan = plan
        self.user_id = user_id
        self.role = role
        self.scopes = scopes or []
        self.rate_limit = rate_limit


def _legacy_hash_key(key: str) -> str:
    """Legacy SHA-256 hash (for migration)."""
    return hashlib.sha256(key.encode()).hexdigest()


def hash_key(key: str) -> str:
    """SEC-20: PBKDF2 hash for API keys (migration-compatible).

    New keys use PBKDF2. Legacy SHA-256 hashes are still accepted
    for backward compatibility but should be migrated.
    """
    # Use PBKDF2 with fixed salt derived from key prefix
    # Fixed salt ensures same key always produces same hash (needed for lookup)
    salt = hashlib.sha256(key[:8].encode()).digest()[:16]
    dk = hashlib.pbkdf2_hmac('sha256', key.encode(), salt, 100000)
    return dk.hex()


async def authenticate(api_key: str) -> TenantContext:
    """
    Authenticate via API Key.
    Returns TenantContext with tenant info and permissions.
    """
    if not api_key:
        raise AuthError("Missing API key", 401)

    # Strip "Bearer " prefix if present
    if api_key.startswith("Bearer "):
        api_key = api_key[7:]

    db = get_db()

    # Try PBKDF2 hash first, fall back to legacy SHA-256
    key_hash = hash_key(api_key)
    record = db.fetchone(
        "SELECT k.*, t.name as tenant_name, t.plan, t.status as tenant_status "
        "FROM api_keys k JOIN tenants t ON k.tenant_id = t.id "
        "WHERE k.key_hash = ? AND k.status = 'active'",
        (key_hash,)
    )
    if not record:
        # Try legacy hash for migration
        legacy_hash = _legacy_hash_key(api_key)
        record = db.fetchone(
            "SELECT k.*, t.name as tenant_name, t.plan, t.status as tenant_status "
            "FROM api_keys k JOIN tenants t ON k.tenant_id = t.id "
            "WHERE k.key_hash = ? AND k.status = 'active'",
            (legacy_hash,)
        )
        if record:
            # Auto-migrate to PBKDF2
            db.execute("UPDATE api_keys SET key_hash = ? WHERE id = ?",
                       (key_hash, record["id"]))
            db.commit()

    if not record:
        raise AuthError("Invalid API key", 401)

    if record["tenant_status"] != "active":
        raise AuthError("Tenant suspended", 403)

    if record.get("expires_at") and record["expires_at"] < now_utc():
        raise AuthError("API key expired", 401)

    # Update last_used_at
    db.execute("UPDATE api_keys SET last_used_at = ? WHERE id = ?",
               (now_utc(), record["id"]))
    db.commit()

    import json
    scopes = json.loads(record.get("scopes", "[]"))

    return TenantContext(
        tenant_id=record["tenant_id"],
        tenant_name=record["tenant_name"],
        plan=record["plan"],
        user_id=record.get("user_id", ""),
        scopes=scopes,
        rate_limit=record.get("rate_limit", 60),
    )