"""MCP tokens: issue, list, revoke, verify. Owned by Agent 1. One table, one service, three rules that matter: 1. **A token is shown once.** Only ``sha256(token)`` and a six-character prefix are stored, so a database leak cannot authenticate anyone; the prefix is for the audit trail and the settings screen, and is never enough to log in. 2. **Verification is constant-time and total.** ``hmac.compare_digest`` on the hash, and every failure mode — unknown, revoked, expired — returns the same ``None``. A caller cannot tell "wrong token" from "revoked token", because a revoked token is not a fact about the token, it is a decision. 3. **Scopes are a floor, not a ceiling.** A token carries the scopes it was minted with and nothing else; an empty scope list is "no scopes", never "all". The per-organization switch lives in ``org_policies.mcp_enabled`` (Phase 3). """ from __future__ import annotations import hashlib import hmac import secrets from collections import deque from dataclasses import dataclass, field from datetime import timedelta from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from app.core.config import settings from app.core.errors import NotFoundError, ValidationError from app.domain.base import new_id, utcnow from app.infrastructure.db.models_org import McpTokenRow #: Every scope the gateway understands. A scope nobody granted a token is a #: scope the gateway refuses, whatever the asks say. MCP_SCOPES: tuple[str, ...] = ( "app:read", "tasks:read", "tasks:write", "memory:read", "memory:write", "delivery:read", ) #: Enough entropy that a token is unguessable, short enough to paste into a #: client config file. TOKEN_ENTROPY_BYTES = 24 #: How the raw token looks to a human: the prefix is what the UI shows back. TOKEN_PREFIX = "smo_" #: Displayed characters; ``token_hash`` is what authenticates. PREFIX_DISPLAY_CHARS = 6 @dataclass(slots=True, kw_only=True) class McpToken: """A token as the service hands it back — never the raw secret.""" id: str organization_id: str name: str prefix: str scopes: list[str] created_by_user_id: str | None = None expires_at: object | None = None last_used_at: object | None = None revoked_at: object | None = None created_at: object | None = None @dataclass(slots=True, kw_only=True) class IssuedToken: """The one and only response that carries the raw token.""" token: str record: McpToken @dataclass(slots=True, kw_only=True) class TokenIdentity: """What a verified token proves: which organization, and what it may do.""" organization_id: str token_id: str prefix: str scopes: frozenset[str] = field(default_factory=frozenset) def allows(self, scope: str) -> bool: return scope in self.scopes def as_utc(value): """Normalise a timestamp read from the database for comparison. SQLite (the development and test store) hands back naive datetimes; every expiry and revocation check compares against ``utcnow()``, so a naive value is assumed to be UTC rather than raising in the middle of a refusal. """ if value is None: return None from datetime import UTC return value.replace(tzinfo=UTC) if value.tzinfo is None else value def hash_token(raw: str) -> str: """The only form of a token that is ever stored.""" return hashlib.sha256(raw.encode("utf-8")).hexdigest() def _mint() -> str: return f"{TOKEN_PREFIX}{secrets.token_urlsafe(TOKEN_ENTROPY_BYTES)}" def _scopes_of(raw: str) -> list[str]: """Parse a CSV scope field, keeping only known scopes and their order.""" seen: list[str] = [] for piece in (raw or "").split(","): scope = piece.strip() if scope and scope in MCP_SCOPES and scope not in seen: seen.append(scope) return seen def validate_scopes(scopes: list[str] | None) -> list[str]: """Refuse an unknown scope rather than silently dropping it.""" if not scopes: return [] unknown = [scope for scope in scopes if scope not in MCP_SCOPES] if unknown: raise ValidationError(f"محدودهٔ ناشناخته: {', '.join(sorted(unknown))}") # De-duplicated, order preserved, so the UI's checkbox set round-trips. return list(dict.fromkeys(scopes)) class TokenRateLimiter: """In-memory, per-token call budget for the current minute. Deliberately process-local: the gateway is one process, and a distributed limiter would need the Redis the deployment does not have. A restart forgives a burst — which is the right trade for a spare process only the owner's own agents reach. """ def __init__(self, *, limit_per_minute: int | None = None) -> None: self.limit = int(limit_per_minute or settings.mcp_rate_limit_per_minute) self._calls: dict[str, deque[float]] = {} def allow(self, key: str) -> bool: import time now = time.monotonic() window_start = now - 60.0 bucket = self._calls.setdefault(key, deque()) while bucket and bucket[0] < window_start: bucket.popleft() if len(bucket) >= self.limit: return False bucket.append(now) return True def reset(self) -> None: self._calls.clear() class McpService: """Token lifecycle and verification for the MCP gateway.""" def __init__(self, session: AsyncSession) -> None: self.session = session # -- lifecycle ----------------------------------------------------------- async def issue( self, *, organization_id: str, name: str, scopes: list[str], created_by_user_id: str | None = None, ttl_days: int | None = None, ) -> IssuedToken: clean_name = name.strip() if not clean_name: raise ValidationError("نام توکن نمی‌تواند خالی باشد.") wanted = validate_scopes(scopes) if not wanted: raise ValidationError("حداقل یک محدودهٔ دسترسی انتخاب کنید.") days = int(ttl_days or settings.mcp_token_ttl_days) if days < 1: raise ValidationError("مدت اعتبار توکن نامعتبر است.") raw = _mint() row = McpTokenRow( id=new_id(), organization_id=organization_id, name=clean_name[:120], token_hash=hash_token(raw), prefix=raw[: len(TOKEN_PREFIX) + PREFIX_DISPLAY_CHARS], scopes=",".join(wanted), created_by_user_id=created_by_user_id, expires_at=utcnow() + timedelta(days=days), ) self.session.add(row) await self.session.flush() return IssuedToken(token=raw, record=_to_domain(row)) async def list(self, organization_id: str) -> list[McpToken]: rows = ( ( await self.session.execute( select(McpTokenRow) .where(McpTokenRow.organization_id == organization_id) .order_by(McpTokenRow.created_at.desc()) ) ) .scalars() .all() ) return [_to_domain(row) for row in rows] async def revoke(self, *, organization_id: str, token_id: str) -> None: """Revoke a token. Revoking twice is a 404, not a silent success — the same idempotence rule the grant endpoints follow, so a caller can tell "already gone" from "never existed" only by looking, not by guessing.""" row = await self._row(organization_id, token_id) if row is None or row.revoked_at is not None: raise NotFoundError("توکن یافت نشد.") row.revoked_at = utcnow() await self.session.flush() # -- verification -------------------------------------------------------- async def verify(self, raw_token: str) -> TokenIdentity | None: """Resolve a raw token, or ``None`` for every kind of refusal. The lookup is by hash (the token is never stored), and the comparison is constant-time so a token that shares a prefix with a real one cannot be discovered by timing. An expired or revoked token is simply not an identity any more. """ if not raw_token or not raw_token.startswith(TOKEN_PREFIX): return None digest = hash_token(raw_token) rows = ( ( await self.session.execute( select(McpTokenRow).where( McpTokenRow.token_hash == digest, McpTokenRow.revoked_at.is_(None), ) ) ) .scalars() .all() ) if not rows: return None row = rows[0] if not hmac.compare_digest(row.token_hash, digest): return None if as_utc(row.expires_at) is not None and as_utc(row.expires_at) <= utcnow(): return None return TokenIdentity( organization_id=row.organization_id, token_id=row.id, prefix=row.prefix, scopes=frozenset(_scopes_of(row.scopes)), ) async def creator_of(self, token_id: str) -> str | None: """The user who minted this token — the accountable actor for its writes. A task created through the gateway is attributed to the administrator who created the token, not to a ghost user: the token is a delegated key, and delegation never erases who holds it. """ row = await self.session.get(McpTokenRow, token_id) return row.created_by_user_id if row else None async def touch(self, token_id: str) -> None: """Record last use. Never allowed to fail a call.""" row = await self.session.get(McpTokenRow, token_id) if row is not None: row.last_used_at = utcnow() await self.session.flush() async def count_per_org(self) -> dict[str, int]: """Active tokens per organization, for the owner's /admin/mcp screen.""" rows = ( await self.session.execute( select(McpTokenRow.organization_id).where( McpTokenRow.revoked_at.is_(None) ) ) ).scalars() counts: dict[str, int] = {} for org_id in rows: counts[org_id] = counts.get(org_id, 0) + 1 return counts async def _row(self, organization_id: str, token_id: str) -> McpTokenRow | None: return await self.session.scalar( select(McpTokenRow).where( McpTokenRow.organization_id == organization_id, McpTokenRow.id == token_id, ) ) def _to_domain(row: McpTokenRow) -> McpToken: return McpToken( id=row.id, organization_id=row.organization_id, name=row.name, prefix=row.prefix, scopes=_scopes_of(row.scopes), created_by_user_id=row.created_by_user_id, expires_at=row.expires_at, last_used_at=row.last_used_at, revoked_at=row.revoked_at, created_at=row.created_at, ) __all__ = [ "MCP_SCOPES", "TOKEN_PREFIX", "IssuedToken", "McpService", "McpToken", "TokenIdentity", "TokenRateLimiter", "hash_token", "validate_scopes", ]