colibre/src/mcp/oauth/store.py

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


def save_code(
    db_path,
    code: str,
    *,
    client_id: str,
    user_id: int,
    redirect_uri: str | None,
    code_challenge: str | None,
    code_challenge_method: str | None,
    scope: str | None,
    resource: str | None,
    expires_at: str,
) -> None:
    with _connect(db_path) as conn:
        conn.execute(
            "INSERT INTO oauth_codes (code_hash, client_id, user_id, redirect_uri, "
            "code_challenge, code_challenge_method, scope, resource, expires_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                _hash(code),
                client_id,
                user_id,
                redirect_uri,
                code_challenge,
                code_challenge_method,
                scope,
                resource,
                expires_at,
            ),
        )
        conn.commit()


def get_code(db_path, code: str) -> dict | None:
    with _connect(db_path) as conn:
        row = conn.execute(
            "SELECT * FROM oauth_codes WHERE code_hash = ?", (_hash(code),)
        ).fetchone()
    return dict(row) if row else None


def delete_code(db_path, code: str) -> None:
    with _connect(db_path) as conn:
        conn.execute("DELETE FROM oauth_codes WHERE code_hash = ?", (_hash(code),))
        conn.commit()


def save_token(
    db_path,
    *,
    access_token: str,
    refresh_token: str | None,
    client_id: str,
    user_id: int,
    scope: str | None,
    resource: str,
    access_expires_at: str,
    refresh_expires_at: str | None,
) -> int:
    with _connect(db_path) as conn:
        cur = conn.execute(
            "INSERT INTO oauth_tokens (access_token_hash, refresh_token_hash, "
            "client_id, user_id, scope, resource, issued_at, access_expires_at, "
            "refresh_expires_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
            (
                _hash(access_token),
                _hash(refresh_token) if refresh_token else None,
                client_id,
                user_id,
                scope,
                resource,
                _utcnow_iso(),
                access_expires_at,
                refresh_expires_at,
            ),
        )
        conn.commit()
        return cur.lastrowid


def get_token_by_access(db_path, access_token: str) -> dict | None:
    with _connect(db_path) as conn:
        row = conn.execute(
            "SELECT * FROM oauth_tokens WHERE access_token_hash = ?",
            (_hash(access_token),),
        ).fetchone()
    return dict(row) if row else None


def get_token_by_refresh(db_path, refresh_token: str) -> dict | None:
    with _connect(db_path) as conn:
        row = conn.execute(
            "SELECT * FROM oauth_tokens WHERE refresh_token_hash = ?",
            (_hash(refresh_token),),
        ).fetchone()
    return dict(row) if row else None


def revoke_token(db_path, token_id: int) -> None:
    with _connect(db_path) as conn:
        conn.execute(
            "UPDATE oauth_tokens SET revoked_at = ? WHERE id = ?",
            (_utcnow_iso(), token_id),
        )
        conn.commit()


def increment_usage(db_path, token_id: int) -> None:
    with _connect(db_path) as conn:
        conn.execute(
            "UPDATE oauth_tokens SET count_total = count_total + 1, last_used_at = ? "
            "WHERE id = ?",
            (_utcnow_iso(), token_id),
        )
        conn.commit()