feat(org): model grants, org memory, MCP gateway (v2 phases A1-3..5)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
mjalmoudy94
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
c93b25741e
commit
dee2f72c46
@@ -0,0 +1,337 @@
|
||||
"""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",
|
||||
]
|
||||
Reference in New Issue
Block a user