dav-mcp/oauth.py

"""Single-user OAuth authorization server for ChatGPT's remote MCP connection."""

from __future__ import annotations

import hashlib
import hmac
import html
import json
import secrets
import time
from pathlib import Path
from typing import Any
from urllib.parse import urlencode

import aiosqlite
from mcp.server.auth.provider import (
    AccessToken,
    AuthorizationCode,
    AuthorizationParams,
    AuthorizeError,
    RegistrationError,
    RefreshToken,
    TokenError,
    construct_redirect_uri,
)
from mcp.shared.auth import OAuthClientInformationFull, OAuthToken
from starlette.exceptions import HTTPException
from starlette.requests import Request
from starlette.responses import HTMLResponse, RedirectResponse, Response


ACCESS_TOKEN_SECONDS = 60 * 60
REFRESH_TOKEN_SECONDS = 60 * 60 * 24 * 30
FLOW_SECONDS = 10 * 60
LOGIN_WINDOW_SECONDS = 15 * 60
LOGIN_FAILURE_LIMIT = 8
MAX_REGISTERED_CLIENTS = 32


class LocalOAuthProvider:
    """OAuth provider with PKCE, rotating refresh tokens, and SQLite persistence."""

    def __init__(
        self,
        *,
        issuer_url: str,
        resource_url: str,
        database_path: str,
        username: str,
        password: str,
        writable: bool,
    ) -> None:
        self.issuer_url = issuer_url.rstrip("/")
        self.resource_url = resource_url
        self.database_path = Path(database_path)
        self.username = username
        self.password = password
        self.allowed_scopes = ["notes:read"] + (["notes:write"] if writable else [])
        self.default_scopes = list(self.allowed_scopes)

    async def initialize(self) -> None:
        self.database_path.parent.mkdir(parents=True, exist_ok=True)
        async with aiosqlite.connect(self.database_path) as db:
            await db.execute("PRAGMA journal_mode=WAL")
            await db.executescript(
                """
                CREATE TABLE IF NOT EXISTS oauth_clients (
                    client_id TEXT PRIMARY KEY,
                    metadata_json TEXT NOT NULL
                );
                CREATE TABLE IF NOT EXISTS oauth_flows (
                    flow_id TEXT PRIMARY KEY,
                    client_id TEXT NOT NULL,
                    client_state TEXT,
                    redirect_uri TEXT NOT NULL,
                    redirect_uri_provided_explicitly INTEGER NOT NULL,
                    code_challenge TEXT NOT NULL,
                    resource TEXT NOT NULL,
                    scopes_json TEXT NOT NULL,
                    expires_at INTEGER NOT NULL
                );
                CREATE TABLE IF NOT EXISTS oauth_codes (
                    code TEXT PRIMARY KEY,
                    client_id TEXT NOT NULL,
                    redirect_uri TEXT NOT NULL,
                    redirect_uri_provided_explicitly INTEGER NOT NULL,
                    code_challenge TEXT NOT NULL,
                    resource TEXT NOT NULL,
                    subject TEXT NOT NULL,
                    scopes_json TEXT NOT NULL,
                    expires_at INTEGER NOT NULL
                );
                CREATE TABLE IF NOT EXISTS oauth_tokens (
                    access_token TEXT PRIMARY KEY,
                    refresh_token TEXT NOT NULL UNIQUE,
                    client_id TEXT NOT NULL,
                    resource TEXT NOT NULL,
                    subject TEXT NOT NULL,
                    scopes_json TEXT NOT NULL,
                    access_expires_at INTEGER NOT NULL,
                    refresh_expires_at INTEGER NOT NULL
                );
                CREATE TABLE IF NOT EXISTS oauth_login_failures (
                    client_ip TEXT NOT NULL,
                    failed_at INTEGER NOT NULL
                );
                CREATE INDEX IF NOT EXISTS oauth_login_failures_by_ip_time
                    ON oauth_login_failures(client_ip, failed_at);
                """
            )
            await db.commit()

    async def _connect(self) -> aiosqlite.Connection:
        db = await aiosqlite.connect(self.database_path)
        db.row_factory = aiosqlite.Row
        return db

    async def get_client(self, client_id: str) -> OAuthClientInformationFull | None:
        async with await self._connect() as db:
            async with db.execute(
                "SELECT metadata_json FROM oauth_clients WHERE client_id = ?", (client_id,)
            ) as cursor:
                row = await cursor.fetchone()
        if row is None:
            return None
        return OAuthClientInformationFull.model_validate_json(row["metadata_json"])

    async def register_client(self, client_info: OAuthClientInformationFull) -> None:
        if not client_info.client_id:
            raise ValueError("OAuth client registration did not include a client ID")
        data = client_info.model_dump_json()
        async with await self._connect() as db:
            async with db.execute(
                "SELECT COUNT(*) AS n FROM oauth_clients WHERE client_id != ?",
                (client_info.client_id,),
            ) as cursor:
                row = await cursor.fetchone()
            if row is not None and row["n"] >= MAX_REGISTERED_CLIENTS:
                raise RegistrationError(
                    "invalid_client_metadata",
                    "This single-user MCP server has reached its OAuth client registration limit",
                )
            await db.execute(
                "INSERT OR REPLACE INTO oauth_clients (client_id, metadata_json) VALUES (?, ?)",
                (client_info.client_id, data),
            )
            await db.commit()

    async def authorize(self, client: OAuthClientInformationFull, params: AuthorizationParams) -> str:
        requested_scopes = list(params.scopes or self.default_scopes)
        requested_scopes = list(dict.fromkeys(requested_scopes))
        if not set(requested_scopes).issubset(self.allowed_scopes):
            raise AuthorizeError("invalid_scope", "The requested notes permission is not enabled")
        if "notes:read" not in requested_scopes:
            requested_scopes.insert(0, "notes:read")

        if params.resource is not None and params.resource != self.resource_url:
            raise AuthorizeError("invalid_target", "This authorization server only issues tokens for this MCP endpoint")

        flow_id = secrets.token_urlsafe(32)
        async with await self._connect() as db:
            now = int(time.time())
            await db.execute("DELETE FROM oauth_flows WHERE expires_at < ?", (now,))
            await db.execute("DELETE FROM oauth_codes WHERE expires_at < ?", (now,))
            await db.execute("DELETE FROM oauth_tokens WHERE refresh_expires_at < ?", (now,))
            await db.execute(
                """INSERT INTO oauth_flows
                   (flow_id, client_id, client_state, redirect_uri,
                    redirect_uri_provided_explicitly, code_challenge, resource,
                    scopes_json, expires_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    flow_id,
                    client.client_id,
                    params.state,
                    str(params.redirect_uri),
                    int(params.redirect_uri_provided_explicitly),
                    params.code_challenge,
                    self.resource_url,
                    json.dumps(requested_scopes),
                    now + FLOW_SECONDS,
                ),
            )
            await db.commit()
        return f"{self.issuer_url}/login?{urlencode({'flow': flow_id})}"

    async def get_flow(self, flow_id: str) -> dict[str, Any] | None:
        async with await self._connect() as db:
            async with db.execute(
                "SELECT * FROM oauth_flows WHERE flow_id = ?", (flow_id,)
            ) as cursor:
                row = await cursor.fetchone()
            if row is not None and row["expires_at"] < int(time.time()):
                await db.execute("DELETE FROM oauth_flows WHERE flow_id = ?", (flow_id,))
                await db.commit()
                return None
        if row is None:
            return None
        result = dict(row)
        result["scopes"] = json.loads(result.pop("scopes_json"))
        return result

    async def login_is_limited(self, client_ip: str) -> bool:
        cutoff = int(time.time()) - LOGIN_WINDOW_SECONDS
        async with await self._connect() as db:
            async with db.execute(
                "SELECT COUNT(*) AS n FROM oauth_login_failures WHERE client_ip = ? AND failed_at >= ?",
                (client_ip, cutoff),
            ) as cursor:
                row = await cursor.fetchone()
        return bool(row and row["n"] >= LOGIN_FAILURE_LIMIT)

    async def record_login_failure(self, client_ip: str) -> None:
        now = int(time.time())
        cutoff = now - LOGIN_WINDOW_SECONDS
        async with await self._connect() as db:
            await db.execute("DELETE FROM oauth_login_failures WHERE failed_at < ?", (cutoff,))
            await db.execute(
                "INSERT INTO oauth_login_failures (client_ip, failed_at) VALUES (?, ?)",
                (client_ip, now),
            )
            await db.commit()

    async def clear_login_failures(self, client_ip: str) -> None:
        async with await self._connect() as db:
            await db.execute("DELETE FROM oauth_login_failures WHERE client_ip = ?", (client_ip,))
            await db.commit()

    def credentials_match(self, username: str, password: str) -> bool:
        name_ok = hmac.compare_digest(username.encode("utf-8"), self.username.encode("utf-8"))
        password_ok = hmac.compare_digest(password.encode("utf-8"), self.password.encode("utf-8"))
        return name_ok and password_ok

    async def finish_login(self, flow_id: str, username: str) -> str:
        flow = await self.get_flow(flow_id)
        if flow is None:
            raise HTTPException(status_code=400, detail="This sign-in request has expired. Start again from ChatGPT.")

        code_value = secrets.token_urlsafe(48)
        now = int(time.time())
        async with await self._connect() as db:
            await db.execute(
                """INSERT INTO oauth_codes
                   (code, client_id, redirect_uri, redirect_uri_provided_explicitly,
                    code_challenge, resource, subject, scopes_json, expires_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    code_value,
                    flow["client_id"],
                    flow["redirect_uri"],
                    flow["redirect_uri_provided_explicitly"],
                    flow["code_challenge"],
                    flow["resource"],
                    username,
                    json.dumps(flow["scopes"]),
                    now + 300,
                ),
            )
            await db.execute("DELETE FROM oauth_flows WHERE flow_id = ?", (flow_id,))
            await db.commit()
        return construct_redirect_uri(
            flow["redirect_uri"],
            code=code_value,
            state=flow["client_state"],
        )

    async def load_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: str
    ) -> AuthorizationCode | None:
        async with await self._connect() as db:
            async with db.execute(
                "SELECT * FROM oauth_codes WHERE code = ? AND client_id = ?",
                (authorization_code, client.client_id),
            ) as cursor:
                row = await cursor.fetchone()
        if row is None or row["expires_at"] < int(time.time()):
            return None
        return AuthorizationCode(
            code=row["code"],
            client_id=row["client_id"],
            redirect_uri=row["redirect_uri"],
            redirect_uri_provided_explicitly=bool(row["redirect_uri_provided_explicitly"]),
            code_challenge=row["code_challenge"],
            resource=row["resource"],
            subject=row["subject"],
            scopes=json.loads(row["scopes_json"]),
            expires_at=row["expires_at"],
        )

    async def exchange_authorization_code(
        self, client: OAuthClientInformationFull, authorization_code: AuthorizationCode
    ) -> OAuthToken:
        if client.client_id != authorization_code.client_id:
            raise TokenError("invalid_grant", "Authorization code was issued to a different client")
        now = int(time.time())
        access = secrets.token_urlsafe(48)
        refresh = secrets.token_urlsafe(48)
        async with await self._connect() as db:
            await db.execute("BEGIN IMMEDIATE")
            async with db.execute(
                "SELECT expires_at FROM oauth_codes WHERE code = ? AND client_id = ?",
                (authorization_code.code, client.client_id),
            ) as cursor:
                row = await cursor.fetchone()
            if row is None or row["expires_at"] < now:
                await db.rollback()
                raise TokenError("invalid_grant", "Authorization code is invalid or expired")
            await db.execute("DELETE FROM oauth_codes WHERE code = ?", (authorization_code.code,))
            await db.execute(
                """INSERT INTO oauth_tokens
                   (access_token, refresh_token, client_id, resource, subject,
                    scopes_json, access_expires_at, refresh_expires_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    access,
                    refresh,
                    client.client_id,
                    authorization_code.resource or self.resource_url,
                    authorization_code.subject or self.username,
                    json.dumps(authorization_code.scopes),
                    now + ACCESS_TOKEN_SECONDS,
                    now + REFRESH_TOKEN_SECONDS,
                ),
            )
            await db.commit()
        return OAuthToken(
            access_token=access,
            token_type="Bearer",
            expires_in=ACCESS_TOKEN_SECONDS,
            scope=" ".join(authorization_code.scopes),
            refresh_token=refresh,
        )

    async def load_refresh_token(
        self, client: OAuthClientInformationFull, refresh_token: str
    ) -> RefreshToken | None:
        now = int(time.time())
        async with await self._connect() as db:
            async with db.execute(
                """SELECT * FROM oauth_tokens
                   WHERE refresh_token = ? AND client_id = ? AND refresh_expires_at >= ?""",
                (refresh_token, client.client_id, now),
            ) as cursor:
                row = await cursor.fetchone()
        if row is None:
            return None
        return RefreshToken(
            token=row["refresh_token"],
            client_id=row["client_id"],
            scopes=json.loads(row["scopes_json"]),
            expires_at=row["refresh_expires_at"],
            resource=row["resource"],
            subject=row["subject"],
        )

    async def exchange_refresh_token(
        self,
        client: OAuthClientInformationFull,
        refresh_token: RefreshToken,
        scopes: list[str],
    ) -> OAuthToken:
        now = int(time.time())
        granted_scopes = list(dict.fromkeys(scopes or refresh_token.scopes))
        if not set(granted_scopes).issubset(refresh_token.scopes):
            raise TokenError("invalid_scope", "A refresh token cannot grant additional permissions")
        access = secrets.token_urlsafe(48)
        refresh = secrets.token_urlsafe(48)
        async with await self._connect() as db:
            await db.execute("BEGIN IMMEDIATE")
            async with db.execute(
                """SELECT * FROM oauth_tokens WHERE refresh_token = ? AND client_id = ?
                   AND refresh_expires_at >= ?""",
                (refresh_token.token, client.client_id, now),
            ) as cursor:
                row = await cursor.fetchone()
            if row is None:
                await db.rollback()
                raise TokenError("invalid_grant", "Refresh token is invalid or expired")
            await db.execute(
                "DELETE FROM oauth_tokens WHERE refresh_token = ?", (refresh_token.token,)
            )
            await db.execute(
                """INSERT INTO oauth_tokens
                   (access_token, refresh_token, client_id, resource, subject,
                    scopes_json, access_expires_at, refresh_expires_at)
                   VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
                (
                    access,
                    refresh,
                    client.client_id,
                    row["resource"],
                    row["subject"],
                    json.dumps(granted_scopes),
                    now + ACCESS_TOKEN_SECONDS,
                    now + REFRESH_TOKEN_SECONDS,
                ),
            )
            await db.commit()
        return OAuthToken(
            access_token=access,
            token_type="Bearer",
            expires_in=ACCESS_TOKEN_SECONDS,
            scope=" ".join(granted_scopes),
            refresh_token=refresh,
        )

    async def load_access_token(self, token: str) -> AccessToken | None:
        now = int(time.time())
        async with await self._connect() as db:
            async with db.execute(
                "SELECT * FROM oauth_tokens WHERE access_token = ? AND access_expires_at >= ?",
                (token, now),
            ) as cursor:
                row = await cursor.fetchone()
        if row is None:
            return None
        return AccessToken(
            token=row["access_token"],
            client_id=row["client_id"],
            scopes=json.loads(row["scopes_json"]),
            expires_at=row["access_expires_at"],
            resource=row["resource"],
            subject=row["subject"],
        )

    async def revoke_token(self, token: AccessToken | RefreshToken) -> None:
        column = "access_token" if isinstance(token, AccessToken) else "refresh_token"
        async with await self._connect() as db:
            await db.execute(f"DELETE FROM oauth_tokens WHERE {column} = ?", (token.token,))
            await db.commit()


def _csrf_cookie_name(flow_id: str) -> str:
    return "mcp_csrf_" + hashlib.sha256(flow_id.encode("utf-8")).hexdigest()[:20]


def _protect_login_response(response: Response) -> Response:
    response.headers["Content-Security-Policy"] = (
        "default-src 'none'; style-src 'unsafe-inline'; form-action 'self'; "
        "base-uri 'none'; frame-ancestors 'none'"
    )
    response.headers["X-Content-Type-Options"] = "nosniff"
    response.headers["X-Frame-Options"] = "DENY"
    response.headers["Referrer-Policy"] = "no-referrer"
    response.headers["Cache-Control"] = "no-store"
    response.headers["Pragma"] = "no-cache"
    return response


def _login_html(flow: dict[str, Any], flow_id: str, csrf: str, error: str | None = None) -> str:
    scopes = flow["scopes"]
    requested = []
    if "notes:read" in scopes:
        requested.append("Read and search notes in /Notes/Personal")
    if "notes:write" in scopes:
        requested.append("Create or replace notes and create folders in /Notes/Personal")
    scope_list = "".join(f"<li>{html.escape(item)}</li>" for item in requested)
    error_html = f'<p class="error">{html.escape(error)}</p>' if error else ""
    return f"""<!doctype html>
<html lang="en">
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Connect WebDAV Notes</title>
<style>
body {{ font: 16px system-ui, sans-serif; max-width: 34rem; margin: 4rem auto; padding: 0 1rem; color: #17202a; }}
input {{ display: block; box-sizing: border-box; width: 100%; padding: .7rem; margin: .35rem 0 1rem; }}
button {{ padding: .7rem 1rem; }} .error {{ color: #a00; }} li {{ margin: .4rem 0; }}
</style>
<h1>Connect WebDAV Notes</h1>
<p>ChatGPT is requesting access to your Obsidian vault:</p>
<ul>{scope_list}</ul>
{error_html}
<form method="post" action="{html.escape(flow['issuer_url'])}/login/callback" autocomplete="on">
  <input type="hidden" name="flow" value="{html.escape(flow_id)}">
  <input type="hidden" name="csrf" value="{html.escape(csrf)}">
  <label>Username<input name="username" autocomplete="username" required></label>
  <label>Password<input type="password" name="password" autocomplete="current-password" required></label>
  <label><input type="checkbox" name="approve" value="yes" required style="display:inline;width:auto"> Allow the permissions above</label>
  <button type="submit">Authorize</button>
</form>
</html>"""


def make_login_routes(provider: LocalOAuthProvider) -> list[Any]:
    async def login_page(request: Request) -> Response:
        flow_id = request.query_params.get("flow", "")
        flow = await provider.get_flow(flow_id)
        if flow is None:
            return HTMLResponse("This sign-in request has expired. Start again from ChatGPT.", status_code=400)
        flow["issuer_url"] = provider.issuer_url
        csrf = secrets.token_urlsafe(24)
        response = _protect_login_response(HTMLResponse(_login_html(flow, flow_id, csrf)))
        response.set_cookie(
            _csrf_cookie_name(flow_id), csrf, max_age=FLOW_SECONDS, secure=True,
            httponly=True, samesite="strict", path="/"
        )
        return response

    async def login_callback(request: Request) -> Response:
        form = await request.form()
        flow_id = form.get("flow")
        username = form.get("username")
        password = form.get("password")
        csrf = form.get("csrf")
        if not all(isinstance(value, str) for value in (flow_id, username, password, csrf)):
            raise HTTPException(status_code=400, detail="Invalid sign-in request")
        assert isinstance(flow_id, str)
        assert isinstance(username, str)
        assert isinstance(password, str)
        assert isinstance(csrf, str)

        flow = await provider.get_flow(flow_id)
        if flow is None:
            raise HTTPException(status_code=400, detail="This sign-in request has expired. Start again from ChatGPT.")
        flow["issuer_url"] = provider.issuer_url
        client_ip = request.client.host if request.client else "unknown"
        cookie_csrf = request.cookies.get(_csrf_cookie_name(flow_id), "")
        if not cookie_csrf or not hmac.compare_digest(cookie_csrf, csrf):
            return HTMLResponse("Sign-in verification failed. Reload the authorization page and try again.", status_code=403)

        if await provider.login_is_limited(client_ip):
            return HTMLResponse("Too many failed sign-in attempts. Please wait 15 minutes before trying again.", status_code=429)

        approved = form.get("approve") == "yes"
        if not approved or not provider.credentials_match(username, password):
            await provider.record_login_failure(client_ip)
            response = _protect_login_response(
                HTMLResponse(
                    _login_html(flow, flow_id, csrf, "Sign-in failed. Check your credentials and permission selection."),
                    status_code=401,
                )
            )
            response.set_cookie(
                _csrf_cookie_name(flow_id), csrf, max_age=FLOW_SECONDS, secure=True,
                httponly=True, samesite="strict", path="/"
            )
            return response

        await provider.clear_login_failures(client_ip)
        redirect_url = await provider.finish_login(flow_id, username)
        response = _protect_login_response(RedirectResponse(redirect_url, status_code=303))
        response.delete_cookie(_csrf_cookie_name(flow_id), path="/", secure=True, httponly=True, samesite="strict")
        return response

    from starlette.routing import Route

    return [
        Route("/login", login_page, methods=["GET"]),
        Route("/login/callback", login_callback, methods=["POST"]),
    ]