diff --git a/src/mcp/oauth/__init__.py b/src/mcp/oauth/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/mcp/oauth/store.py b/src/mcp/oauth/store.py new file mode 100644 index 0000000..2d63403 --- /dev/null +++ b/src/mcp/oauth/store.py @@ -0,0 +1,90 @@ +import hashlib +import json +import sqlite3 +from contextlib import contextmanager +from datetime import datetime, timezone +from pathlib import Path + +SCHEMA = """ +CREATE TABLE IF NOT EXISTS oauth_clients ( + client_id TEXT PRIMARY KEY, + client_metadata TEXT NOT NULL, + created_at TEXT NOT NULL +); +CREATE TABLE IF NOT EXISTS oauth_codes ( + code_hash TEXT PRIMARY KEY, + client_id TEXT NOT NULL, + user_id INTEGER NOT NULL, + redirect_uri TEXT, + code_challenge TEXT, + code_challenge_method TEXT, + scope TEXT, + resource TEXT, + expires_at TEXT NOT NULL, + used INTEGER NOT NULL DEFAULT 0 +); +CREATE TABLE IF NOT EXISTS oauth_tokens ( + id INTEGER PRIMARY KEY, + access_token_hash TEXT NOT NULL UNIQUE, + refresh_token_hash TEXT UNIQUE, + client_id TEXT NOT NULL, + user_id INTEGER NOT NULL, + scope TEXT, + resource TEXT NOT NULL, + issued_at TEXT NOT NULL, + access_expires_at TEXT NOT NULL, + refresh_expires_at TEXT, + revoked_at TEXT, + count_total INTEGER NOT NULL DEFAULT 0, + last_used_at TEXT +); +CREATE INDEX IF NOT EXISTS idx_oauth_tokens_access ON oauth_tokens(access_token_hash); +CREATE INDEX IF NOT EXISTS idx_oauth_tokens_refresh ON oauth_tokens(refresh_token_hash); +""" + + +def _utcnow_iso() -> str: + return datetime.now(timezone.utc).isoformat(timespec="seconds") + + +def _hash(value: str) -> str: + return hashlib.sha256(value.encode()).hexdigest() + + +@contextmanager +def _connect(db_path): + conn = sqlite3.connect(str(db_path)) + conn.row_factory = sqlite3.Row + try: + yield conn + finally: + conn.close() + + +def init_schema(db_path) -> None: + Path(db_path).parent.mkdir(parents=True, exist_ok=True) + with _connect(db_path) as conn: + conn.executescript(SCHEMA) + conn.commit() + + +def create_client(db_path, client_id: str, metadata: dict) -> None: + with _connect(db_path) as conn: + conn.execute( + "INSERT INTO oauth_clients (client_id, client_metadata, created_at) " + "VALUES (?, ?, ?)", + (client_id, json.dumps(metadata), _utcnow_iso()), + ) + conn.commit() + + +def get_client(db_path, client_id: str) -> dict | None: + with _connect(db_path) as conn: + row = conn.execute( + "SELECT * FROM oauth_clients WHERE client_id = ?", (client_id,) + ).fetchone() + if row is None: + return None + d = dict(row) + d["client_metadata"] = json.loads(d["client_metadata"]) + return d diff --git a/tests/mcp/test_oauth_store.py b/tests/mcp/test_oauth_store.py new file mode 100644 index 0000000..bca7a65 --- /dev/null +++ b/tests/mcp/test_oauth_store.py @@ -0,0 +1,15 @@ +from src.mcp.oauth import store + + +def test_create_and_get_client(tmp_path): + db = tmp_path / "u.sqlite" + store.init_schema(db) + meta = { + "redirect_uris": ["https://claude.ai/api/mcp/auth_callback"], + "scope": "mcp", + } + store.create_client(db, "abc123", meta) + row = store.get_client(db, "abc123") + assert row["client_id"] == "abc123" + assert row["client_metadata"] == meta + assert store.get_client(db, "nope") is None