smart-org/apps/api/app/application/services/mcp_service.py

"""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",
]