verae-mcp-timestamping-server/src/verae_mcp_timestamping_server/server.py

"""MCP resource server: stamp / verify / hash_report. Hex only, never content."""

from __future__ import annotations

import asyncio
import os
import sys
from contextlib import contextmanager

from mcp.server import MCPServer
from mcp.server.auth.middleware.auth_context import get_access_token
from mcp.server.auth.settings import AuthSettings, ClientRegistrationOptions
from mcp.server.mcpserver.exceptions import ToolError
from pydantic import AnyHttpUrl
from starlette.requests import Request
from starlette.responses import Response

from verae_mcp_timestamping_server.global_client import (
    CreateUserResult,
    GlobalClient,
    GlobalLedger,
    LoginResult,
    SessionResult,
    request_token,
)
from verae_mcp_timestamping_server.oauth_login import VeraeOAuthProvider
from verae_mcp_timestamping_server.setup_cmd import print_login_hint, run_session, run_setup
from verae_mcp_timestamping_server.guard import reject_content_keys, strict_arg_model
from verae_mcp_timestamping_server.ledger import (
    FileLedger,
    Ledger,
    MemoryLedger,
    Receipts,
    RoutingLedger,
    VerifyResult,
)

mcp = MCPServer(
    "verae-timestamping",
    instructions=(
        "Global stamps need a VeraeTime login on the machine that runs this server, "
        "or a bearer token from the browser signup on the public server. "
        "If session is not valid, ask the user to run "
        "`uvx verae-mcp-timestamping-server setup`. "
        "Do not ask them to paste a token into the chat or into mcp.json. "
        "The token file is ~/.verae/auth.properties, mode 600."
    ),
)
_ledger: Ledger = MemoryLedger()
_global_client: GlobalClient = GlobalClient()
_oauth: VeraeOAuthProvider | None = None


@contextmanager
def _caller_token():
    access = get_access_token()
    if access is None or not access.token:
        yield
        return
    reset = request_token.set(access.token)
    try:
        yield
    finally:
        request_token.reset(reset)


@mcp.custom_route("/account", methods=["GET", "POST"])
async def account_page(request: Request) -> Response:
    """Browser signup and login. The JWT is handed to the assistant, not printed."""
    if _oauth is None:
        return Response("OAuth is not enabled on this server.", status_code=404)
    return await _oauth.handle_account(request)


def default_ledger(client: GlobalClient | None = None) -> Ledger:
    held = client or GlobalClient.from_env()
    return RoutingLedger(
        local=FileLedger(),
        global_ledger=GlobalLedger(held),
    )


def _client_of(ledger: Ledger) -> GlobalClient | None:
    inner = getattr(ledger, "global_ledger", None)
    client = getattr(inner, "client", None)
    if isinstance(client, GlobalClient):
        return client
    return None


_HARDENED_TOOLS: set[int] = set()


def _harden_hash_only_tools(server: MCPServer) -> None:
    """Forbid unknown args (content/bytes included); advertise additionalProperties=false."""
    for tool in server._tool_manager.list_tools():
        tid = id(tool)
        if tid in _HARDENED_TOOLS:
            continue
        meta = tool.fn_metadata
        object.__setattr__(meta, "arg_model", strict_arg_model(meta.arg_model))

        def _make_validate(m):
            def validate_arguments(arguments_to_validate):
                reject_content_keys(arguments_to_validate)
                arguments_pre_parsed = m.pre_parse_json(arguments_to_validate)
                arguments_parsed_model = m.arg_model.model_validate(arguments_pre_parsed)
                return arguments_parsed_model.model_dump_one_level()

            return validate_arguments

        object.__setattr__(meta, "validate_arguments", _make_validate(meta))
        params = meta.arg_model.model_json_schema(by_alias=True)
        params["additionalProperties"] = False
        object.__setattr__(tool, "parameters", params)
        _HARDENED_TOOLS.add(tid)


def build_mcp(
    ledger: Ledger | None = None,
    global_client: GlobalClient | None = None,
) -> MCPServer:
    global _ledger, _global_client
    if ledger is None:
        client = global_client or GlobalClient.from_env()
        _ledger = default_ledger(client)
        _global_client = client
    else:
        _ledger = ledger
        _global_client = global_client or _client_of(ledger) or GlobalClient()
    _harden_hash_only_tools(mcp)
    return mcp


def _tool_error(exc: ValueError) -> ToolError:
    return ToolError(str(exc))


@mcp.tool()
async def stamp(
    sha256: str,
    scope: str = "local",
    cube_id: str = "",
    org_id: str = "",
    metadata: dict | None = None,
    restamp: bool = False,
    message: str = "",
) -> Receipts:
    """Commit a SHA-256. scope=local is the leaf file ledger. scope=global posts the digest to api.veraetime.net and reuses the first CID unless restamp is true. Optional metadata is sent as a verae-meta-git/1 history; blank values do not delete fields."""
    try:
        with _caller_token():
            receipts = await _ledger.stamp(
                sha256,
                scope=scope,
                cube_id=cube_id,
                org_id=org_id,
                metadata=metadata,
                restamp=restamp,
                message=message,
            )
    except ValueError as exc:
        raise _tool_error(exc) from exc
    return Receipts(receipts=receipts)


@mcp.tool()
async def verify(receipt_id: str, want_scope: str = "") -> VerifyResult:
    """Check a receipt. want_scope=global asks api.veraetime.net (central source of truth)."""
    with _caller_token():
        return await _ledger.verify(receipt_id, want_scope=want_scope)


@mcp.tool()
async def login(username: str, password: str) -> LoginResult:
    """Log in to api.veraetime.net. Stores the JWT on this host and does not return it."""
    try:
        return await asyncio.to_thread(_global_client.login, username, password)
    except ValueError as exc:
        raise _tool_error(exc) from exc


@mcp.tool()
async def session() -> SessionResult:
    """Validate the stored VeraeTime JWT. Returns valid, username, role, and expires_at."""
    return await asyncio.to_thread(_global_client.session)


@mcp.tool()
async def signup(username: str, password: str) -> CreateUserResult:
    """Create a VeraeTime user account. Role is user. Does not replace the stored JWT."""
    try:
        return await asyncio.to_thread(_global_client.signup, username, password)
    except ValueError as exc:
        raise _tool_error(exc) from exc


@mcp.tool()
async def create_user(username: str, password: str, role: str) -> CreateUserResult:
    """Create a VeraeTime user. The stored JWT must belong to an admin."""
    try:
        return await asyncio.to_thread(
            _global_client.create_user, username, password, role
        )
    except ValueError as exc:
        raise _tool_error(exc) from exc


@mcp.tool()
async def hash_report(
    sha256: str, org_id: str = "", want_scope: str = ""
) -> Receipts:
    """JSON-array history of receipts for one SHA-256 (Carfax)."""
    try:
        receipts = await _ledger.hash_report(
            sha256, org_id=org_id, want_scope=want_scope
        )
    except ValueError as exc:
        raise _tool_error(exc) from exc
    return Receipts(receipts=receipts)


def enable_oauth(public_url: str) -> VeraeOAuthProvider:
    """Turn on browser signup for the HTTP server. Stdio stays open unless this runs."""
    global _oauth
    issuer = public_url.rstrip("/")
    provider = VeraeOAuthProvider(api=_global_client, issuer=issuer)
    mcp.settings.auth = AuthSettings(
        issuer_url=AnyHttpUrl(issuer),
        resource_server_url=AnyHttpUrl(issuer + "/mcp"),
        client_registration_options=ClientRegistrationOptions(enabled=True, valid_scopes=["stamp"]),
        required_scopes=["stamp"],
        validate_token_resource=False,
    )
    mcp._auth_server_provider = provider
    from mcp.server.auth.provider import ProviderTokenVerifier

    mcp._token_verifier = ProviderTokenVerifier(provider)
    _oauth = provider
    return provider


def main() -> None:
    if len(sys.argv) > 1 and sys.argv[1] == "setup":
        raise SystemExit(run_setup(sys.argv[2:]))
    if len(sys.argv) > 1 and sys.argv[1] == "session":
        raise SystemExit(run_session())
    client = GlobalClient.from_env()
    build_mcp(default_ledger(client), global_client=client)
    print_login_hint(client)
    if os.environ.get("VERAE_MCP_OAUTH") == "1":
        public = os.environ.get("VERAE_MCP_PUBLIC_URL", "http://127.0.0.1:8000")
        enable_oauth(public)
    transport = os.environ.get("VERAE_MCP_TRANSPORT", "stdio")
    if transport == "streamable-http":
        host = os.environ.get("VERAE_MCP_HOST", "127.0.0.1")
        port = int(os.environ.get("VERAE_MCP_PORT", "8000"))
        from mcp.server.transport_security import TransportSecuritySettings

        # Behind Caddy on ns1; allow public hostname + local probes.
        allowed = [
            h.strip()
            for h in os.environ.get(
                "VERAE_MCP_ALLOWED_HOSTS",
                "mcp.veraetime.net,mcp.georgelambert.org,localhost,localhost:8000,127.0.0.1,127.0.0.1:8000,10.1.10.128,10.1.10.128:8000",
            ).split(",")
            if h.strip()
        ]
        transport_security = TransportSecuritySettings(
            enable_dns_rebinding_protection=True,
            allowed_hosts=allowed,
            allowed_origins=os.environ.get("VERAE_MCP_ALLOWED_ORIGINS", "*").split(","),
        )
        if os.environ.get("VERAE_MCP_CORS") == "1":
            import uvicorn
            from starlette.middleware.cors import CORSMiddleware

            app = mcp.streamable_http_app(
                json_response=True,
                transport_security=transport_security,
                host=host,
            )
            app.add_middleware(
                CORSMiddleware,
                allow_origins=["*"],
                allow_methods=["GET", "POST", "DELETE", "OPTIONS"],
                allow_headers=["*"],
                expose_headers=["mcp-session-id", "Mcp-Session-Id"],
            )
            uvicorn.run(app, host=host, port=port)
            return
        mcp.run(
            transport="streamable-http",
            host=host,
            port=port,
            transport_security=transport_security,
        )
        return
    mcp.run()


if __name__ == "__main__":
    main()