agentic-mobile-control/packages/cloud-platform/cloud/auth.py

from __future__ import annotations

from dataclasses import dataclass, field
from hashlib import sha256
from hmac import compare_digest
from collections.abc import Iterable
from typing import TYPE_CHECKING, Protocol, runtime_checkable

if TYPE_CHECKING:
    from cloud.repository import CloudRepository


TASKS_SUBMIT_SCOPE = "tasks:submit"
TASKS_READ_SCOPE = "tasks:read"
POOL_READ_SCOPE = "pool:read"
PLUGINS_READ_SCOPE = "plugins:read"
PLUGINS_ADMIN_SCOPE = "plugins:admin"
USERS_ADMIN_SCOPE = "users:admin"


@dataclass(frozen=True)
class Principal:
    id: str = "anonymous"
    scopes: frozenset[str] = field(default_factory=frozenset)
    host_id: str | None = None
    session_id: str | None = None
    must_change_password: bool = False

    def has_scope(self, scope: str) -> bool:
        return "*" in self.scopes or scope in self.scopes

    def require_host(self, requested_host_id: str) -> None:
        if self.host_id is None:
            raise HostPrincipalRequiredError("a host-bound principal is required")
        if self.host_id != requested_host_id:
            raise HostIdentityMismatchError(
                "authenticated host cannot act for the requested host"
            )


class HostAuthorizationError(PermissionError):
    pass


class HostPrincipalRequiredError(HostAuthorizationError):
    pass


class HostIdentityMismatchError(HostAuthorizationError):
    pass


@runtime_checkable
class AuthProvider(Protocol):
    def authenticate(self, request: object) -> Principal | None: ...


class NullAuthProvider:
    """Explicit insecure-development provider with unrestricted scope."""

    def authenticate(self, request: object) -> Principal | None:
        return Principal(scopes=frozenset({"*"}))


@dataclass(frozen=True)
class BearerCredential:
    principal_id: str
    token: str = field(repr=False)
    scopes: frozenset[str] = field(default_factory=frozenset)
    host_id: str | None = None

    def __post_init__(self) -> None:
        if not self.principal_id.strip():
            raise ValueError("bearer credential principal_id must not be empty")
        if not self.token:
            raise ValueError("bearer credential token must not be empty")
        if self.host_id is not None and not self.host_id.strip():
            raise ValueError("bearer credential host_id must not be empty")


@dataclass(frozen=True)
class _StoredBearerCredential:
    principal: Principal
    token_digest: bytes = field(repr=False)


@dataclass(frozen=True)
class EnrollmentCredential:
    principal_id: str
    token: str = field(repr=False)

    def __post_init__(self) -> None:
        if not self.principal_id.strip():
            raise ValueError("enrollment credential principal_id must not be empty")
        if not self.token:
            raise ValueError("enrollment credential token must not be empty")


@dataclass(frozen=True)
class EnrollmentPrincipal:
    id: str
    token_digest: str | None = field(default=None, repr=False)


@runtime_checkable
class EnrollmentAuthProvider(Protocol):
    def authenticate(self, request: object) -> EnrollmentPrincipal | None: ...


class ConfiguredBearerAuthProvider:
    """Authenticate configured bearer tokens without exposing credential text."""

    def __init__(self, credentials: list[BearerCredential]) -> None:
        self._credentials = tuple(
            _StoredBearerCredential(
                principal=Principal(
                    id=credential.principal_id,
                    scopes=frozenset(credential.scopes),
                    host_id=credential.host_id,
                ),
                token_digest=_token_digest(credential.token),
            )
            for credential in credentials
        )

    def authenticate(self, request: object) -> Principal | None:
        token = _extract_bearer_token(request)
        if token is None:
            return None

        candidate_digest = _token_digest(token)
        matched_principal: Principal | None = None
        for credential in self._credentials:
            if compare_digest(candidate_digest, credential.token_digest):
                matched_principal = credential.principal
        return matched_principal


class ConfiguredEnrollmentTokenProvider:
    def __init__(self, credentials: Iterable[EnrollmentCredential]) -> None:
        self._credentials = tuple(
            EnrollmentPrincipal(
                id=credential.principal_id,
                token_digest=digest_token(credential.token),
            )
            for credential in credentials
        )

    def authenticate(self, request: object) -> EnrollmentPrincipal | None:
        candidate = bearer_token_digest(request)
        if candidate is None:
            return None
        candidate_bytes = bytes.fromhex(candidate)
        matched: EnrollmentPrincipal | None = None
        for credential in self._credentials:
            if compare_digest(
                candidate_bytes,
                bytes.fromhex(credential.token_digest),
            ):
                matched = credential
        return matched


class SelfServiceEnrollmentAuthProvider:
    """Unconditionally authorizes Host enrollment with no pre-issued token."""

    def authenticate(self, request: object) -> EnrollmentPrincipal | None:
        return EnrollmentPrincipal(id="self-service", token_digest=None)


class ChainedEnrollmentAuthProvider:
    def __init__(self, providers: Iterable[EnrollmentAuthProvider]) -> None:
        self.providers = tuple(providers)

    def authenticate(self, request: object) -> EnrollmentPrincipal | None:
        for provider in self.providers:
            principal = provider.authenticate(request)
            if principal is not None:
                return principal
        return None


class RepositoryHostAuthProvider:
    def __init__(self, repository: CloudRepository) -> None:
        self.repository = repository

    def authenticate(self, request: object) -> Principal | None:
        credential_digest = bearer_token_digest(request)
        if credential_digest is None:
            return None
        host_id = self.repository.authenticate_enrolled_host(credential_digest)
        if host_id is None:
            return None
        return Principal(id=f"enrolled-host:{host_id}", host_id=host_id)


class UserSessionAuthProvider:
    """Resolve the opaque browser session cookie to an existing Principal."""

    def __init__(self, user_auth_service: object) -> None:
        self.user_auth_service = user_auth_service

    def authenticate(self, request: object) -> Principal | None:
        from cloud.user_auth import USER_SESSION_COOKIE

        token = _extract_cookie(request, USER_SESSION_COOKIE)
        if token is None:
            return None
        authenticated = self.user_auth_service.authenticate_session(token)  # type: ignore[attr-defined]
        if authenticated is None:
            return None
        user = authenticated.user
        return Principal(
            id=f"user:{user.id}",
            scopes=user.scopes,
            session_id=authenticated.session.id,
            must_change_password=user.must_change_password,
        )


class ChainedAuthProvider:
    def __init__(self, providers: Iterable[AuthProvider]) -> None:
        self.providers = tuple(providers)

    def authenticate(self, request: object) -> Principal | None:
        for provider in self.providers:
            principal = provider.authenticate(request)
            if principal is not None:
                return principal
        return None


def create_auth_provider(
    credentials: Iterable[BearerCredential],
    *,
    allow_insecure_anonymous: bool,
) -> AuthProvider:
    configured = list(credentials)
    if configured:
        return ConfiguredBearerAuthProvider(configured)
    if allow_insecure_anonymous:
        return NullAuthProvider()
    return ConfiguredBearerAuthProvider([])


def _extract_bearer_token(request: object) -> str | None:
    headers = getattr(request, "headers", None)
    if headers is None:
        return None
    authorization = headers.get("authorization")
    if not isinstance(authorization, str):
        return None
    scheme, separator, token = authorization.partition(" ")
    if not separator or scheme.lower() != "bearer" or not token:
        return None
    return token


def _extract_cookie(request: object, name: str) -> str | None:
    cookies = getattr(request, "cookies", None)
    if cookies is None:
        return None
    value = cookies.get(name)
    return value if isinstance(value, str) and value else None


def bearer_token_digest(request: object) -> str | None:
    token = _extract_bearer_token(request)
    return digest_token(token) if token is not None else None


def digest_token(token: str) -> str:
    return _token_digest(token).hex()


def _token_digest(token: str) -> bytes:
    return sha256(token.encode("utf-8")).digest()