app/security/service.py

from __future__ import annotations

import asyncio
import hashlib
import hmac
import re
import secrets
import time
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone

from sqlalchemy import func, select, update

from app.core.config import Settings
from app.security.context import AuthContext
from app.security.database import SecurityDatabase
from app.security.errors import (
    APIKeyConflictError,
    APIKeyNotFoundError,
    ForbiddenError,
    UnauthorizedError,
)
from app.security.models import APIKey
from app.security.schemas import APIKeyCreate, APIKeyPatch
from app.security.scopes import configured_roles, effective_scopes
from app.security.tenancy import TenantService

KEY_PATTERN = re.compile(r"^mp_([a-z][a-z0-9]{1,15})_([A-Za-z0-9_-]{43,})$")
HASH_PATTERN = re.compile(r"^[a-f0-9]{64}$")


def utcnow() -> datetime:
    return datetime.now(timezone.utc)


def aware(value: datetime | None) -> datetime | None:
    if value is None:
        return None
    return (
        value.replace(tzinfo=timezone.utc)
        if value.tzinfo is None
        else value.astimezone(timezone.utc)
    )


@dataclass(frozen=True, slots=True)
class KeyMaterial:
    api_key: str
    key_prefix: str
    key_hash: str
    environment: str


class APIKeyService:
    """Creates and validates opaque API keys without retaining plaintext secrets."""

    def __init__(
        self, database: SecurityDatabase, settings: Settings, tenants: TenantService
    ) -> None:
        self.database = database
        self.settings = settings
        self.tenants = tenants
        self.roles = configured_roles(settings.auth_role_scopes)
        self._last_used_cache: dict[str, float] = {}
        self._last_used_lock = asyncio.Lock()

    @staticmethod
    def generate_material(environment: str) -> KeyMaterial:
        normalized = environment.strip().lower()
        if normalized not in {"live", "test"}:
            raise ValueError("API key environment must be live or test")
        secret = secrets.token_urlsafe(32)
        api_key = f"mp_{normalized}_{secret}"
        return KeyMaterial(
            api_key=api_key,
            key_prefix=f"mp_{normalized}_{secret[:8]}",
            key_hash=hashlib.sha256(api_key.encode("utf-8")).hexdigest(),
            environment=normalized,
        )

    @staticmethod
    def hash_key(api_key: str) -> str:
        return hashlib.sha256(api_key.encode("utf-8")).hexdigest()

    @staticmethod
    def parse_key(api_key: str) -> tuple[str, str]:
        match = KEY_PATTERN.fullmatch(api_key)
        if not match:
            raise UnauthorizedError
        environment, secret = match.groups()
        return environment, f"mp_{environment}_{secret[:8]}"

    async def ensure_bootstrap_admin(self) -> None:
        key_hash = self.settings.auth_bootstrap_key_hash.strip().lower()
        key_prefix = self.settings.auth_bootstrap_key_prefix.strip()
        if not key_hash and not key_prefix:
            return
        if not HASH_PATTERN.fullmatch(key_hash) or not re.fullmatch(
            r"mp_[a-z][a-z0-9]{1,15}_[A-Za-z0-9_-]{8}", key_prefix
        ):
            raise ValueError("Bootstrap key hash or prefix is malformed")
        prefix_environment = key_prefix.split("_", 2)[1]
        if prefix_environment != self.settings.auth_bootstrap_environment:
            raise ValueError("Bootstrap key prefix and environment do not match")
        async with self.database.session() as session:
            count = await session.scalar(select(func.count()).select_from(APIKey))
            if count:
                return
            session.add(
                APIKey(
                    name=self.settings.auth_bootstrap_key_name,
                    key_prefix=key_prefix,
                    key_hash=key_hash,
                    environment=self.settings.auth_bootstrap_environment,
                    status="active",
                    role="admin",
                    scopes=["admin"],
                    created_by="bootstrap",
                    notes="Hash-only bootstrap administrator",
                    requests_per_minute=self.settings.auth_default_requests_per_minute,
                    concurrent_jobs=self.settings.auth_default_concurrent_jobs,
                    uploads_per_hour=self.settings.auth_default_uploads_per_hour,
                    processing_bytes_per_day=(self.settings.auth_default_processing_bytes_per_day),
                )
            )
            await session.commit()
        await self.tenants.ensure_all_api_key_principals()

    async def create(
        self,
        payload: APIKeyCreate,
        *,
        created_by: str | None,
        workspace_id: str | None = None,
        user_id: str | None = None,
    ) -> tuple[APIKey, str]:
        role = payload.role.strip().lower() if payload.role else None
        if role and role not in self.roles:
            raise APIKeyConflictError(f"Unknown role '{role}'")
        material = self.generate_material(payload.environment)
        expires_at = aware(payload.expires_at)
        if payload.expires_in_seconds is not None:
            expires_at = utcnow() + timedelta(seconds=payload.expires_in_seconds)
        record = APIKey(
            name=payload.name.strip(),
            key_prefix=material.key_prefix,
            key_hash=material.key_hash,
            environment=material.environment,
            status="active",
            role=role,
            scopes=payload.scopes,
            expires_at=expires_at,
            created_by=created_by,
            notes=payload.notes,
            requests_per_minute=(
                payload.requests_per_minute or self.settings.auth_default_requests_per_minute
            ),
            concurrent_jobs=(payload.concurrent_jobs or self.settings.auth_default_concurrent_jobs),
            uploads_per_hour=(
                payload.uploads_per_hour or self.settings.auth_default_uploads_per_hour
            ),
            processing_bytes_per_day=(
                payload.processing_bytes_per_day
                or self.settings.auth_default_processing_bytes_per_day
            ),
        )
        async with self.database.session() as session:
            session.add(record)
            await session.commit()
            await session.refresh(record)
        if workspace_id and user_id:
            await self.tenants.bind_api_key(
                api_key_id=record.id, workspace_id=workspace_id, user_id=user_id
            )
        else:
            await self.tenants.resolve_api_key(record.id)
        return record, material.api_key

    async def authenticate(self, api_key: str) -> AuthContext:
        environment, key_prefix = self.parse_key(api_key)
        supplied_hash = self.hash_key(api_key)
        async with self.database.session() as session:
            candidates = list(
                (await session.scalars(select(APIKey).where(APIKey.key_prefix == key_prefix))).all()
            )
        record: APIKey | None = None
        for candidate in candidates:
            if hmac.compare_digest(candidate.key_hash, supplied_hash):
                record = candidate
        now = utcnow()
        if record is None or record.environment != environment:
            raise UnauthorizedError
        expires_at = aware(record.expires_at)
        grace_expires_at = aware(record.grace_expires_at)
        if record.status == "rotating":
            if grace_expires_at is None or grace_expires_at <= now:
                await self._finalize_rotation(record.id)
                raise UnauthorizedError
        elif record.status != "active":
            raise UnauthorizedError
        if expires_at is not None and expires_at <= now:
            raise UnauthorizedError
        scopes = effective_scopes(record.role, record.scopes or [], self.settings.auth_role_scopes)
        principal = await self.tenants.resolve_api_key(record.id)
        return AuthContext(
            api_key_id=record.id,
            key_name=record.name,
            key_prefix=record.key_prefix,
            environment=record.environment,
            role=record.role,
            scopes=scopes,
            requests_per_minute=record.requests_per_minute,
            concurrent_jobs=record.concurrent_jobs,
            uploads_per_hour=record.uploads_per_hour,
            processing_bytes_per_day=record.processing_bytes_per_day,
            expires_at=expires_at,
            workspace_id=principal.workspace_id,
            user_id=principal.user_id,
            membership_id=principal.membership_id,
            membership_role=principal.membership_role,
        )

    @staticmethod
    def authorize(context: AuthContext, required_scope: str | None) -> None:
        if required_scope is not None and not context.allows(required_scope):
            raise ForbiddenError

    async def mark_used(self, context: AuthContext) -> None:
        await self._touch_last_used(context.api_key_id)

    async def list(self, *, offset: int = 0, limit: int = 100) -> tuple[list[APIKey], int]:
        await self._finalize_expired_rotations()
        async with self.database.session() as session:
            total = int(await session.scalar(select(func.count()).select_from(APIKey)) or 0)
            records = list(
                (
                    await session.scalars(
                        select(APIKey)
                        .order_by(APIKey.created_at.desc())
                        .offset(offset)
                        .limit(limit)
                    )
                ).all()
            )
        return records, total

    async def get(self, key_id: str) -> APIKey:
        await self._finalize_expired_rotations()
        async with self.database.session() as session:
            record = await session.get(APIKey, key_id)
            if record is None:
                raise APIKeyNotFoundError
            return record

    async def patch(self, key_id: str, payload: APIKeyPatch) -> APIKey:
        async with self.database.session() as session:
            record = await session.get(APIKey, key_id)
            if record is None:
                raise APIKeyNotFoundError
            fields = payload.model_fields_set
            if payload.name is not None:
                record.name = payload.name.strip()
            if "role" in fields:
                role = payload.role.strip().lower() if payload.role else None
                if role and role not in self.roles:
                    raise APIKeyConflictError(f"Unknown role '{role}'")
                record.role = role
            if payload.scopes is not None:
                record.scopes = payload.scopes
            if payload.clear_expiration:
                record.expires_at = None
            elif "expires_at" in fields:
                record.expires_at = aware(payload.expires_at)
            if "notes" in fields:
                record.notes = payload.notes
            for field in (
                "requests_per_minute",
                "concurrent_jobs",
                "uploads_per_hour",
                "processing_bytes_per_day",
            ):
                value = getattr(payload, field)
                if value is not None:
                    setattr(record, field, value)
            await session.commit()
            await session.refresh(record)
            return record

    async def set_status(self, key_id: str, status: str) -> APIKey:
        if status not in {"active", "disabled", "revoked"}:
            raise APIKeyConflictError("Unsupported API key status transition")
        async with self.database.session() as session:
            record = await session.get(APIKey, key_id)
            if record is None:
                raise APIKeyNotFoundError
            expires_at = aware(record.expires_at)
            if record.status == "revoked" and status != "revoked":
                raise APIKeyConflictError("Revoked API keys cannot be changed")
            if status == "active" and record.status != "disabled":
                raise APIKeyConflictError("Only disabled API keys can be enabled")
            if status == "active" and expires_at is not None and expires_at <= utcnow():
                raise APIKeyConflictError("Expired API keys cannot be enabled")
            if status == "disabled" and record.status != "active":
                raise APIKeyConflictError("Only active API keys can be disabled")
            record.status = status
            if status != "rotating":
                record.grace_expires_at = None
            await session.commit()
            await session.refresh(record)
            return record

    async def rotate(
        self, key_id: str, grace_period_seconds: int, *, created_by: str | None
    ) -> tuple[APIKey, str]:
        principal = await self.tenants.resolve_api_key(key_id)
        async with self.database.session() as session:
            old = await session.get(APIKey, key_id)
            if old is None:
                raise APIKeyNotFoundError
            if old.status != "active":
                raise APIKeyConflictError("Only active API keys can be rotated")
            old_expires_at = aware(old.expires_at)
            if old_expires_at is not None and old_expires_at <= utcnow():
                raise APIKeyConflictError("Expired API keys cannot be rotated")
            material = self.generate_material(old.environment)
            replacement = APIKey(
                name=old.name,
                key_prefix=material.key_prefix,
                key_hash=material.key_hash,
                environment=old.environment,
                status="active",
                role=old.role,
                scopes=list(old.scopes or []),
                expires_at=old.expires_at,
                created_by=created_by,
                notes=old.notes,
                rotated_from_id=old.id,
                requests_per_minute=old.requests_per_minute,
                concurrent_jobs=old.concurrent_jobs,
                uploads_per_hour=old.uploads_per_hour,
                processing_bytes_per_day=old.processing_bytes_per_day,
            )
            session.add(replacement)
            old.status = "rotating" if grace_period_seconds else "revoked"
            old.grace_expires_at = (
                utcnow() + timedelta(seconds=grace_period_seconds) if grace_period_seconds else None
            )
            await session.commit()
            await session.refresh(replacement)
        await self.tenants.bind_api_key(
            api_key_id=replacement.id,
            workspace_id=principal.workspace_id,
            user_id=principal.user_id,
        )
        return replacement, material.api_key

    async def _touch_last_used(self, key_id: str) -> None:
        interval = self.settings.auth_last_used_update_seconds
        current = time.monotonic()
        async with self._last_used_lock:
            previous = self._last_used_cache.get(key_id)
            if previous is not None and current - previous < interval:
                return
            self._last_used_cache[key_id] = current
        async with self.database.session() as session:
            record = await session.get(APIKey, key_id)
            if record is not None:
                record.last_used_at = utcnow()
                await session.commit()

    async def _finalize_rotation(self, key_id: str) -> None:
        async with self.database.session() as session:
            await session.execute(
                update(APIKey)
                .where(APIKey.id == key_id, APIKey.status == "rotating")
                .values(status="revoked", grace_expires_at=None)
            )
            await session.commit()

    async def _finalize_expired_rotations(self) -> None:
        async with self.database.session() as session:
            await session.execute(
                update(APIKey)
                .where(
                    APIKey.status == "rotating",
                    APIKey.grace_expires_at.is_not(None),
                    APIKey.grace_expires_at <= utcnow(),
                )
                .values(status="revoked", grace_expires_at=None)
            )
            await session.commit()