LangBot/src/langbot/pkg/cloud/support_admin.py

from __future__ import annotations

import dataclasses
import datetime
import hashlib
import re
import time
import typing

import jwt
from sqlalchemy.exc import IntegrityError

from ..entity.persistence.support_admin import SupportAdminTemporarySession
from ..workspace.errors import WorkspaceError

if typing.TYPE_CHECKING:
    from ..core.app import Application


SUPPORT_ADMIN_TOKEN_TYP = 'langbot-support-admin+jwt'
SUPPORT_ADMIN_TOKEN_KIND = 'support_admin.session'
SUPPORT_ADMIN_EFFECTIVE_ROLE = 'owner'
SUPPORT_ADMIN_MAX_TOKEN_SECONDS = 300
_SHA256_HEX = re.compile(r'^[0-9a-f]{64}$')


class SupportAdminSessionError(ValueError):
    """Raised when a support-admin session or token is not admissible."""


class SupportAdminReplayError(SupportAdminSessionError):
    """Raised when a launch grant JTI has already been consumed."""


@dataclasses.dataclass(frozen=True, slots=True)
class IssuedSupportAdminSession:
    token: str
    grant_jti_hash: str
    workspace_uuid: str
    actor_account_uuid: str
    issued_at: datetime.datetime
    expires_at: datetime.datetime


@dataclasses.dataclass(frozen=True, slots=True)
class SupportAdminSessionIdentity:
    grant_jti_hash: str
    workspace_uuid: str
    actor_account_uuid: str
    instance_uuid: str
    placement_generation: int


def hash_grant_jti(jti: str) -> str:
    return hashlib.sha256(jti.encode('utf-8')).hexdigest()


class SupportAdminSessionService:
    """Issue and validate temporary Workspace-scoped support-admin sessions."""

    def __init__(
        self,
        ap: Application,
        *,
        wall_time: typing.Callable[[], float] = time.time,
    ) -> None:
        self.ap = ap
        self._wall_time = wall_time

    async def consume_launch_grant(
        self,
        *,
        grant_jti_hash: str,
        workspace_uuid: str,
        actor_account_uuid: str,
    ) -> IssuedSupportAdminSession:
        self._validate_grant_hash(grant_jti_hash)
        if not workspace_uuid or not actor_account_uuid:
            raise SupportAdminSessionError('Support admin session requires an actor and Workspace')

        issued_at = self._utcnow()
        expires_at = issued_at + datetime.timedelta(seconds=SUPPORT_ADMIN_MAX_TOKEN_SECONDS)
        tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
        if not callable(tenant_uow):
            raise SupportAdminSessionError('Support admin sessions require tenant persistence')

        try:
            async with tenant_uow(workspace_uuid) as uow:
                await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
                uow.session.add(
                    SupportAdminTemporarySession(
                        grant_jti_hash=grant_jti_hash,
                        workspace_uuid=workspace_uuid,
                        actor_account_uuid=actor_account_uuid,
                        issued_at=issued_at,
                        expires_at=expires_at,
                    )
                )
                await uow.session.flush()
        except IntegrityError as exc:
            raise SupportAdminReplayError('Launch assertion has already been consumed') from exc
        except WorkspaceError as exc:
            raise SupportAdminSessionError('Workspace is unavailable for support access') from exc

        return IssuedSupportAdminSession(
            token=self._encode_token(
                grant_jti_hash=grant_jti_hash,
                workspace_uuid=workspace_uuid,
                actor_account_uuid=actor_account_uuid,
                issued_at=issued_at,
                expires_at=expires_at,
            ),
            grant_jti_hash=grant_jti_hash,
            workspace_uuid=workspace_uuid,
            actor_account_uuid=actor_account_uuid,
            issued_at=issued_at,
            expires_at=expires_at,
        )

    def is_support_admin_token(self, token: str) -> bool:
        """Return True only for compact JWTs marked as support-admin tokens."""

        if not isinstance(token, str) or token.count('.') != 2:
            return False
        try:
            header = jwt.get_unverified_header(token)
        except jwt.PyJWTError:
            return False
        if header.get('typ') == SUPPORT_ADMIN_TOKEN_TYP:
            return True
        try:
            payload = jwt.decode(token, options={'verify_signature': False})
        except jwt.PyJWTError:
            return False
        return payload.get('kind') == SUPPORT_ADMIN_TOKEN_KIND

    async def authenticate_token(
        self,
        token: str,
        *,
        requested_workspace_uuid: str | None,
    ) -> SupportAdminSessionIdentity:
        if not self.is_support_admin_token(token):
            raise SupportAdminSessionError('Not a support admin token')
        workspace_uuid = (requested_workspace_uuid or '').strip()
        if not workspace_uuid:
            raise SupportAdminSessionError('Support admin token requires an explicit Workspace selector')

        jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
        try:
            payload = jwt.decode(
                token,
                jwt_secret,
                algorithms=['HS256'],
                issuer='langbot-core',
                audience=self._audience(workspace_uuid),
                options={'require': ['exp', 'iat', 'nbf', 'iss', 'aud']},
            )
        except jwt.PyJWTError as exc:
            raise SupportAdminSessionError('Invalid support admin token') from exc
        self._validate_payload(payload, workspace_uuid)
        grant_jti_hash = payload['grant_jti_hash']
        actor_account_uuid = payload['actor_account_uuid']

        tenant_uow = getattr(self.ap.persistence_mgr, 'tenant_uow', None)
        if not callable(tenant_uow):
            raise SupportAdminSessionError('Support admin sessions require tenant persistence')

        now = self._utcnow()
        async with tenant_uow(workspace_uuid) as uow:
            session = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
            if (
                session is None
                or session.workspace_uuid != workspace_uuid
                or session.actor_account_uuid != actor_account_uuid
                or session.revoked_at is not None
                or session.expires_at <= now
            ):
                raise SupportAdminSessionError('Support admin session is inactive')
            binding = await self.ap.workspace_service.get_execution_binding(workspace_uuid, session=uow.session)
            session.last_used_at = now
            await uow.session.flush()

        return SupportAdminSessionIdentity(
            grant_jti_hash=grant_jti_hash,
            workspace_uuid=workspace_uuid,
            actor_account_uuid=actor_account_uuid,
            instance_uuid=binding.instance_uuid,
            placement_generation=binding.placement_generation,
        )

    async def revoke_session(self, grant_jti_hash: str, workspace_uuid: str) -> None:
        self._validate_grant_hash(grant_jti_hash)
        now = self._utcnow()
        async with self.ap.persistence_mgr.tenant_uow(workspace_uuid) as uow:
            row = await uow.session.get(SupportAdminTemporarySession, grant_jti_hash)
            if row is not None and row.revoked_at is None:
                row.revoked_at = now

    def _encode_token(
        self,
        *,
        grant_jti_hash: str,
        workspace_uuid: str,
        actor_account_uuid: str,
        issued_at: datetime.datetime,
        expires_at: datetime.datetime,
    ) -> str:
        jwt_secret = self.ap.instance_config.data['system']['jwt']['secret']
        payload: dict[str, typing.Any] = {
            'kind': SUPPORT_ADMIN_TOKEN_KIND,
            'iss': 'langbot-core',
            'aud': self._audience(workspace_uuid),
            'sub': f'support-admin:{actor_account_uuid}',
            'iat': issued_at,
            'nbf': issued_at,
            'exp': expires_at,
            'actor_account_uuid': actor_account_uuid,
            'workspace_uuid': workspace_uuid,
            'effective_role': SUPPORT_ADMIN_EFFECTIVE_ROLE,
            'grant_jti_hash': grant_jti_hash,
        }
        return jwt.encode(payload, jwt_secret, algorithm='HS256', headers={'typ': SUPPORT_ADMIN_TOKEN_TYP})

    def _validate_payload(self, payload: dict[str, typing.Any], workspace_uuid: str) -> None:
        if payload.get('kind') != SUPPORT_ADMIN_TOKEN_KIND:
            raise SupportAdminSessionError('Invalid support admin token kind')
        if payload.get('workspace_uuid') != workspace_uuid:
            raise SupportAdminSessionError('Support admin session is scoped to another Workspace')
        if payload.get('effective_role') != SUPPORT_ADMIN_EFFECTIVE_ROLE:
            raise SupportAdminSessionError('Invalid support admin token role')
        actor_account_uuid = payload.get('actor_account_uuid')
        if not isinstance(actor_account_uuid, str) or not actor_account_uuid.strip():
            raise SupportAdminSessionError('Invalid support admin actor')
        grant_jti_hash = payload.get('grant_jti_hash')
        if not isinstance(grant_jti_hash, str) or not _SHA256_HEX.match(grant_jti_hash):
            raise SupportAdminSessionError('Invalid support admin grant')

    def _audience(self, workspace_uuid: str) -> str:
        return f'langbot-support-admin:{self.ap.workspace_service.instance_uuid}:{workspace_uuid}'

    @staticmethod
    def _validate_grant_hash(grant_jti_hash: str) -> None:
        if not _SHA256_HEX.match(grant_jti_hash):
            raise SupportAdminSessionError('Invalid support admin grant')

    def _utcnow(self) -> datetime.datetime:
        return datetime.datetime.fromtimestamp(self._wall_time(), datetime.UTC).replace(tzinfo=None)