520 lines
23 KiB
Python
520 lines
23 KiB
Python
"""Project-local OAuth provider for the MCP SDK: consent, PKCE grants and rotating tokens.
|
|
|
|
The SDK validates protocol requests, redirects, client credentials and S256 PKCE.
|
|
This module stores grants and authenticates the project owner on a separate consent page.
|
|
It implements a shared project access code, not a multi-user identity service.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
import hashlib
|
|
import hmac
|
|
from html import escape
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import secrets
|
|
import sqlite3
|
|
import time
|
|
from urllib.parse import parse_qs, urlsplit
|
|
|
|
from mcp.server.auth.provider import (
|
|
AccessToken,
|
|
AuthorizationCode,
|
|
AuthorizationParams,
|
|
AuthorizeError,
|
|
RefreshToken,
|
|
RegistrationError,
|
|
TokenError,
|
|
construct_redirect_uri,
|
|
)
|
|
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
|
|
from starlette.responses import HTMLResponse, JSONResponse, RedirectResponse
|
|
|
|
SCOPE = "zpwiki:read"
|
|
SCOPES = [SCOPE, "offline_access"]
|
|
CALLBACKS = {"https://claude.ai/api/mcp/auth_callback"}
|
|
COOKIE = "__Host-zp-oauth"
|
|
PENDING_TTL, CODE_TTL, ACCESS_TTL, REFRESH_TTL = 600, 120, 3600, 30 * 86400
|
|
|
|
|
|
def digest(value: str) -> str:
|
|
return hashlib.sha256(value.encode()).hexdigest()
|
|
|
|
|
|
class Store:
|
|
def __init__(self, path: Path, namespace: str):
|
|
self.path, self.namespace = path, namespace
|
|
path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
|
|
if not path.exists():
|
|
fd = os.open(path, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0o600)
|
|
os.close(fd)
|
|
path.chmod(0o600)
|
|
with self.transaction() as conn:
|
|
conn.execute("""CREATE TABLE IF NOT EXISTS records (
|
|
namespace TEXT NOT NULL, kind TEXT NOT NULL, key TEXT NOT NULL,
|
|
body TEXT NOT NULL, expires REAL NOT NULL, family TEXT NOT NULL DEFAULT '',
|
|
PRIMARY KEY(namespace,kind,key))""")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS oauth_family ON records(namespace,family)")
|
|
conn.execute("CREATE INDEX IF NOT EXISTS oauth_expiry ON records(expires)")
|
|
|
|
@contextmanager
|
|
def transaction(self):
|
|
conn = sqlite3.connect(self.path, timeout=5)
|
|
try:
|
|
conn.execute("BEGIN IMMEDIATE")
|
|
yield conn
|
|
conn.commit()
|
|
except BaseException:
|
|
conn.rollback()
|
|
raise
|
|
finally:
|
|
conn.close()
|
|
|
|
def get(self, conn, kind, key):
|
|
row = conn.execute(
|
|
"SELECT body FROM records WHERE namespace=? AND kind=? AND key=? AND expires>?",
|
|
(self.namespace, kind, key, time.time()),
|
|
).fetchone()
|
|
return json.loads(row[0]) if row else None
|
|
|
|
def put(self, conn, kind, key, body, expires, family=""):
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO records VALUES (?,?,?,?,?,?)",
|
|
(self.namespace, kind, key, json.dumps(body), expires, family),
|
|
)
|
|
|
|
def delete(self, conn, kind, key):
|
|
conn.execute(
|
|
"DELETE FROM records WHERE namespace=? AND kind=? AND key=?",
|
|
(self.namespace, kind, key),
|
|
)
|
|
|
|
def capacity(self, conn, kind, maximum):
|
|
conn.execute("DELETE FROM records WHERE expires<=?", (time.time(),))
|
|
return (
|
|
conn.execute(
|
|
"SELECT count(*) FROM records WHERE namespace=? AND kind=?", (self.namespace, kind)
|
|
).fetchone()[0]
|
|
< maximum
|
|
)
|
|
|
|
def revoke_family(self, conn, family):
|
|
conn.execute("DELETE FROM records WHERE namespace=? AND family=?", (self.namespace, family))
|
|
|
|
|
|
class OAuthProvider:
|
|
def __init__(self, resource: str, static_token: str, password: str, database: Path):
|
|
if (
|
|
len(password) < 32
|
|
or not password.isascii()
|
|
or any(c.isspace() for c in password)
|
|
or hmac.compare_digest(password, static_token)
|
|
):
|
|
raise ValueError(
|
|
"Configure a separate ZP_MCP_LOGIN_PASSWORD with at least 32 characters"
|
|
)
|
|
self.resource = resource
|
|
url = urlsplit(resource)
|
|
self.origin = url.scheme + "://" + url.netloc
|
|
self.static_token, self.password_hash = static_token, digest(password)
|
|
# A changed public URL or rotated owner credential invalidates previous sessions.
|
|
self.store = Store(database, digest(resource + "\0" + static_token + "\0" + password))
|
|
|
|
async def get_client(self, client_id):
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "client", client_id)
|
|
return OAuthClientInformationFull.model_validate(value) if value else None
|
|
|
|
async def register_client(self, client_info):
|
|
redirects = [str(uri) for uri in client_info.redirect_uris or []]
|
|
if (
|
|
not redirects
|
|
or len(redirects) > 4
|
|
or any(uri not in CALLBACKS for uri in redirects)
|
|
or client_info.token_endpoint_auth_method
|
|
not in {"none", "client_secret_post", "client_secret_basic"}
|
|
or set(client_info.grant_types) != {"authorization_code", "refresh_token"}
|
|
or client_info.response_types != ["code"]
|
|
or not set((client_info.scope or "").split()) <= set(SCOPES)
|
|
):
|
|
raise RegistrationError(
|
|
"invalid_client_metadata",
|
|
"Only the documented Claude web callback and code flow are supported",
|
|
)
|
|
with self.store.transaction() as conn:
|
|
if not self.store.capacity(conn, "client", 1024):
|
|
raise RegistrationError(
|
|
"invalid_client_metadata", "Client registration capacity reached"
|
|
)
|
|
self.store.put(
|
|
conn,
|
|
"client",
|
|
client_info.client_id,
|
|
client_info.model_dump(mode="json"),
|
|
time.time() + 90 * 86400,
|
|
)
|
|
|
|
async def authorize(self, client, params: AuthorizationParams):
|
|
if params.resource not in {None, self.resource}:
|
|
raise AuthorizeError(
|
|
"invalid_request", "The requested resource does not match this MCP server"
|
|
)
|
|
scopes = params.scopes or (client.scope or SCOPE).split()
|
|
if SCOPE not in scopes or not set(scopes) <= set(SCOPES):
|
|
raise AuthorizeError("invalid_scope", "The read scope is required")
|
|
if not re.fullmatch(r"[A-Za-z0-9_-]{43}", params.code_challenge):
|
|
raise AuthorizeError("invalid_request", "S256 PKCE is required")
|
|
request_id = secrets.token_urlsafe(32)
|
|
value = {
|
|
"client_id": client.client_id,
|
|
"client_name": (client.client_name or "Claude")[:100],
|
|
"params": params.model_copy(
|
|
update={"scopes": scopes, "resource": self.resource}
|
|
).model_dump(mode="json"),
|
|
"attempts": 0,
|
|
"expires": time.time() + PENDING_TTL,
|
|
}
|
|
with self.store.transaction() as conn:
|
|
if not self.store.capacity(conn, "pending", 1024):
|
|
raise AuthorizeError("temporarily_unavailable", "Please retry later")
|
|
self.store.put(conn, "pending", digest(request_id), value, value["expires"])
|
|
return self.origin + "/oauth/login?request=" + request_id
|
|
|
|
async def load_authorization_code(self, client, authorization_code):
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "code", digest(authorization_code))
|
|
if not value or value["client_id"] != client.client_id:
|
|
return None
|
|
return AuthorizationCode(code=authorization_code, **value)
|
|
|
|
def _issue(self, conn, client_id, scopes, family, refresh_expiry=None):
|
|
now = int(time.time())
|
|
access = secrets.token_urlsafe(48)
|
|
refresh_expiry = refresh_expiry or now + REFRESH_TTL
|
|
expiry = min(now + ACCESS_TTL, refresh_expiry)
|
|
common = {
|
|
"client_id": client_id,
|
|
"scopes": scopes,
|
|
"resource": self.resource,
|
|
"subject": "project-user",
|
|
}
|
|
self.store.put(
|
|
conn, "access", digest(access), {**common, "expires_at": expiry}, expiry, family
|
|
)
|
|
refresh = None
|
|
if "offline_access" in scopes:
|
|
refresh = secrets.token_urlsafe(48)
|
|
self.store.put(
|
|
conn,
|
|
"refresh",
|
|
digest(refresh),
|
|
{**common, "expires_at": refresh_expiry, "family": family, "used": False},
|
|
refresh_expiry,
|
|
family,
|
|
)
|
|
return OAuthToken(
|
|
access_token=access,
|
|
expires_in=expiry - now,
|
|
refresh_token=refresh,
|
|
scope=" ".join(scopes),
|
|
)
|
|
|
|
async def exchange_authorization_code(self, client, authorization_code):
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "code", digest(authorization_code.code))
|
|
if not value or value["client_id"] != client.client_id:
|
|
raise TokenError("invalid_grant", "The code has expired or was already used")
|
|
self.store.delete(conn, "code", digest(authorization_code.code))
|
|
return self._issue(conn, client.client_id, value["scopes"], secrets.token_hex(24))
|
|
|
|
async def load_refresh_token(self, client, refresh_token):
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "refresh", digest(refresh_token))
|
|
if not value or value["client_id"] != client.client_id:
|
|
return None
|
|
return RefreshToken(
|
|
token=refresh_token, **{k: v for k, v in value.items() if k not in {"family", "used"}}
|
|
)
|
|
|
|
async def exchange_refresh_token(self, client, refresh_token, scopes):
|
|
replay = False
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "refresh", digest(refresh_token.token))
|
|
if not value or value["client_id"] != client.client_id:
|
|
raise TokenError("invalid_grant", "Invalid refresh token")
|
|
if value["used"]:
|
|
self.store.revoke_family(conn, value["family"])
|
|
replay = True
|
|
else:
|
|
if SCOPE not in scopes or not set(scopes) <= set(value["scopes"]):
|
|
raise TokenError("invalid_scope", "Requested scopes exceed the grant")
|
|
value["used"] = True
|
|
self.store.put(
|
|
conn,
|
|
"refresh",
|
|
digest(refresh_token.token),
|
|
value,
|
|
value["expires_at"],
|
|
value["family"],
|
|
)
|
|
# Rotate both credentials; retire old access tokens in the same family.
|
|
conn.execute(
|
|
"DELETE FROM records WHERE namespace=? AND kind=? AND family=?",
|
|
(self.store.namespace, "access", value["family"]),
|
|
)
|
|
issued = self._issue(
|
|
conn, client.client_id, scopes, value["family"], value["expires_at"]
|
|
)
|
|
if replay:
|
|
raise TokenError("invalid_grant", "Refresh token reuse revoked this grant")
|
|
return issued
|
|
|
|
async def load_access_token(self, token):
|
|
if hmac.compare_digest(token.encode(), self.static_token.encode()):
|
|
return AccessToken(
|
|
token=token, client_id="local-static", scopes=[SCOPE], resource=self.resource
|
|
)
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "access", digest(token))
|
|
return AccessToken(token=token, **value) if value else None
|
|
|
|
async def revoke_token(self, token):
|
|
with self.store.transaction() as conn:
|
|
value = self.store.get(conn, "access", digest(token.token)) or self.store.get(
|
|
conn, "refresh", digest(token.token)
|
|
)
|
|
row = conn.execute(
|
|
"SELECT family FROM records WHERE namespace=? AND key=? AND kind IN (?,?)",
|
|
(self.store.namespace, digest(token.token), "access", "refresh"),
|
|
).fetchone()
|
|
if value and row:
|
|
self.store.revoke_family(conn, row[0])
|
|
|
|
def _form(self, request_id, pending, csrf, message="", status=200):
|
|
target = escape(urlsplit(pending["params"]["redirect_uri"]).hostname or "")
|
|
name = escape(pending["client_name"])
|
|
document = f"""<!doctype html><html lang="sk"><meta charset="utf-8">
|
|
<meta name="viewport" content="width=device-width,initial-scale=1"><title>Prihlásenie — ZP Agent</title>
|
|
<style>body{{font:17px system-ui;max-width:520px;margin:8vh auto;padding:24px;line-height:1.5;background:#f5f5f5;color:#222}}
|
|
form{{background:white;padding:24px;border-radius:12px}}input,button{{box-sizing:border-box;width:100%;padding:12px;margin:8px 0;font:inherit}}
|
|
.error{{color:#a00}}small{{display:block}}</style>
|
|
<h1>Pripojiť ZP Agenta</h1><p>Aplikácia <strong>{name}</strong> žiada prístup k vyhľadávaniu RAG a GraphRAG.</p>
|
|
<p>Po potvrdení sa vrátiš na <strong>{target}</strong>.</p><form method="post" action="/oauth/login">
|
|
<input type="hidden" name="request" value="{escape(request_id)}"><input type="hidden" name="csrf" value="{escape(csrf)}">
|
|
<label for="password">Prístupový kód ZP Agenta</label><input id="password" name="password" type="password" autocomplete="current-password" required autofocus>
|
|
<small>Kód ti poskytne správca ZP Agenta. Nezadávaj sem heslo svojho účtu Claude.</small>
|
|
<p class="error">{escape(message)}</p><button name="action" value="allow">Prihlásiť a povoliť prístup</button>
|
|
<button name="action" value="deny" formnovalidate>Zrušiť</button></form></html>"""
|
|
return HTMLResponse(document, status_code=status)
|
|
|
|
async def login(self, request):
|
|
if request.method == "GET":
|
|
request_id = request.query_params.get("request", "")
|
|
with self.store.transaction() as conn:
|
|
pending = self.store.get(conn, "pending", digest(request_id))
|
|
if not pending:
|
|
return HTMLResponse(
|
|
"Prihlásenie vypršalo. Začni pripojenie znova v Claude.", status_code=400
|
|
)
|
|
csrf = secrets.token_urlsafe(32)
|
|
pending["csrf"] = digest(csrf)
|
|
self.store.put(conn, "pending", digest(request_id), pending, pending["expires"])
|
|
response = self._form(request_id, pending, csrf)
|
|
response.set_cookie(
|
|
COOKIE,
|
|
csrf,
|
|
max_age=PENDING_TTL,
|
|
secure=True,
|
|
httponly=True,
|
|
samesite="lax",
|
|
path="/",
|
|
)
|
|
return response
|
|
if request.headers.get("origin") not in {None, self.origin}:
|
|
return JSONResponse({"error": "invalid_origin"}, status_code=403)
|
|
if (
|
|
request.headers.get("content-type", "").split(";")[0]
|
|
!= "application/x-www-form-urlencoded"
|
|
):
|
|
return JSONResponse({"error": "invalid_request"}, status_code=415)
|
|
form = parse_qs(
|
|
(await request.body()).decode("utf-8", errors="replace"), keep_blank_values=True
|
|
)
|
|
if any(len(values) != 1 for values in form.values()):
|
|
return JSONResponse({"error": "invalid_request"}, status_code=400)
|
|
get = lambda name: form.get(name, [""])[0]
|
|
request_id, csrf = get("request"), get("csrf")
|
|
with self.store.transaction() as conn:
|
|
pending = self.store.get(conn, "pending", digest(request_id))
|
|
if (
|
|
not pending
|
|
or not csrf
|
|
or not request.cookies.get(COOKIE)
|
|
or not hmac.compare_digest(csrf.encode(), request.cookies[COOKIE].encode())
|
|
or not hmac.compare_digest(digest(csrf), pending.get("csrf", ""))
|
|
):
|
|
return JSONResponse({"error": "invalid_csrf_or_expired_login"}, status_code=403)
|
|
if get("action") == "deny":
|
|
self.store.delete(conn, "pending", digest(request_id))
|
|
response = RedirectResponse(
|
|
construct_redirect_uri(
|
|
pending["params"]["redirect_uri"],
|
|
error="access_denied",
|
|
state=pending["params"].get("state"),
|
|
),
|
|
status_code=303,
|
|
)
|
|
response.delete_cookie(COOKIE, secure=True, httponly=True, samesite="lax")
|
|
return response
|
|
bucket = str(int(time.time()) // 60)
|
|
rate = self.store.get(conn, "login_rate", bucket) or {"count": 0}
|
|
rate["count"] += 1
|
|
self.store.put(conn, "login_rate", bucket, rate, time.time() + 120)
|
|
if rate["count"] > 30 or pending["attempts"] >= 5:
|
|
return self._form(
|
|
request_id, pending, csrf, "Príliš veľa pokusov. Skús to neskôr.", 429
|
|
)
|
|
if get("action") != "allow" or not hmac.compare_digest(
|
|
digest(get("password")), self.password_hash
|
|
):
|
|
pending["attempts"] += 1
|
|
self.store.put(conn, "pending", digest(request_id), pending, pending["expires"])
|
|
return self._form(request_id, pending, csrf, "Nesprávny prístupový kód.", 401)
|
|
code = secrets.token_urlsafe(32)
|
|
params = pending["params"]
|
|
value = {
|
|
k: params[k]
|
|
for k in [
|
|
"scopes",
|
|
"code_challenge",
|
|
"redirect_uri",
|
|
"redirect_uri_provided_explicitly",
|
|
"resource",
|
|
]
|
|
}
|
|
value.update(
|
|
client_id=pending["client_id"],
|
|
expires_at=time.time() + CODE_TTL,
|
|
subject="project-user",
|
|
)
|
|
self.store.put(conn, "code", digest(code), value, value["expires_at"])
|
|
self.store.delete(conn, "pending", digest(request_id))
|
|
response = RedirectResponse(
|
|
construct_redirect_uri(params["redirect_uri"], code=code, state=params.get("state")),
|
|
status_code=303,
|
|
)
|
|
response.delete_cookie(COOKIE, secure=True, httponly=True, samesite="lax")
|
|
return response
|
|
|
|
|
|
class OAuthBoundary:
|
|
"""Bound auth request bodies, enforce resource binding and harden the consent page."""
|
|
|
|
def __init__(self, app, provider):
|
|
self.app, self.provider = app, provider
|
|
|
|
async def __call__(self, scope, receive, send):
|
|
if scope["type"] != "http":
|
|
return await self.app(scope, receive, send)
|
|
headers = dict(scope["headers"])
|
|
public_host = urlsplit(self.provider.origin).netloc
|
|
host = headers.get(b"host", b"").decode("latin1")
|
|
if host != public_host and not re.fullmatch(r"(?:127\.0\.0\.1|localhost)(?::\d+)?", host):
|
|
return await JSONResponse({"error": "invalid_host"}, status_code=421)(
|
|
scope, receive, send
|
|
)
|
|
path = scope["path"]
|
|
if scope["method"] == "POST" and path in {
|
|
"/token",
|
|
"/register",
|
|
"/authorize",
|
|
"/revoke",
|
|
"/oauth/login",
|
|
}:
|
|
body = bytearray()
|
|
while True:
|
|
message = await receive()
|
|
if message["type"] == "http.disconnect":
|
|
return
|
|
body.extend(message.get("body", b""))
|
|
if len(body) > 16384:
|
|
return await JSONResponse({"error": "request_too_large"}, status_code=413)(
|
|
scope, receive, send
|
|
)
|
|
if not message.get("more_body"):
|
|
break
|
|
if path in {"/token", "/revoke"}:
|
|
if (
|
|
headers.get(b"content-type", b"").split(b";")[0]
|
|
!= b"application/x-www-form-urlencoded"
|
|
):
|
|
return await JSONResponse({"error": "invalid_request"}, status_code=415)(
|
|
scope, receive, send
|
|
)
|
|
fields = parse_qs(body.decode("utf-8", errors="replace"), keep_blank_values=True)
|
|
if any(len(v) != 1 for v in fields.values()) or fields.get(
|
|
"resource", [self.provider.resource]
|
|
) != [self.provider.resource]:
|
|
return await JSONResponse({"error": "invalid_request"}, status_code=400)(
|
|
scope, receive, send
|
|
)
|
|
# SDK 1.x models an optional revocation secret as a required nullable field.
|
|
# Public clients have no secret; normalize that omission before its parser.
|
|
if path == "/revoke" and "client_secret" not in fields:
|
|
body.extend(b"&client_secret=")
|
|
if path == "/register":
|
|
try:
|
|
value = json.loads(body)
|
|
if not isinstance(value, dict):
|
|
raise ValueError()
|
|
except (ValueError, UnicodeError):
|
|
return await JSONResponse(
|
|
{"error": "invalid_client_metadata"}, status_code=400
|
|
)(scope, receive, send)
|
|
original_receive = receive
|
|
sent = False
|
|
|
|
async def replay():
|
|
nonlocal sent
|
|
if not sent:
|
|
sent = True
|
|
return {"type": "http.request", "body": bytes(body), "more_body": False}
|
|
return await original_receive()
|
|
|
|
receive = replay
|
|
|
|
async def secured(message):
|
|
if message["type"] == "http.response.start":
|
|
existing = list(message.get("headers", []))
|
|
names = {k.lower() for k, v in existing}
|
|
defaults = {
|
|
b"cache-control": b"no-store",
|
|
b"pragma": b"no-cache",
|
|
b"referrer-policy": b"no-referrer",
|
|
b"x-content-type-options": b"nosniff",
|
|
b"x-frame-options": b"DENY",
|
|
}
|
|
if path == "/oauth/login":
|
|
# no-referrer makes native form POSTs send Origin: null (Fetch).
|
|
# Keep the real origin for this same-origin form while suppressing
|
|
# the referrer on the redirect to Claude. CSRF checks stay strict.
|
|
defaults[b"referrer-policy"] = b"same-origin"
|
|
# Chromium applies form-action to POST redirects too. Permit only
|
|
# the same fixed callbacks accepted during OAuth registration.
|
|
callbacks = " ".join(sorted(CALLBACKS))
|
|
defaults[b"content-security-policy"] = (
|
|
"default-src 'none'; style-src 'unsafe-inline'; form-action 'self' "
|
|
+ callbacks
|
|
+ "; frame-ancestors 'none'; base-uri 'none'"
|
|
).encode()
|
|
message = {
|
|
**message,
|
|
"headers": existing + [(k, v) for k, v in defaults.items() if k not in names],
|
|
}
|
|
await send(message)
|
|
|
|
return await self.app(scope, receive, secured)
|