335 lines
14 KiB
Python
335 lines
14 KiB
Python
#!/usr/bin/env python3
|
|
"""Verify the live OAuth+MCP flow without opening Claude or calling a language model.
|
|
|
|
Reads the owner code from local .env/environment, never prints credentials and
|
|
never follows the final redirect to the Claude account. Test grants are revoked.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import argparse
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
from pathlib import Path
|
|
import re
|
|
import secrets
|
|
import sys
|
|
from urllib.parse import parse_qs, urlsplit
|
|
|
|
import httpx
|
|
from dotenv import dotenv_values
|
|
from mcp import ClientSession
|
|
from mcp.client.streamable_http import streamable_http_client
|
|
|
|
ROOT = Path(__file__).resolve().parents[2]
|
|
CALLBACK = "https://claude.ai/api/mcp/auth_callback"
|
|
|
|
|
|
def require(condition, message):
|
|
if not condition:
|
|
raise RuntimeError(message)
|
|
|
|
|
|
async def browser_login(login_url, origin, password, engine):
|
|
"""Use a fresh browser profile and an actual form POST, with no forged Origin."""
|
|
from playwright.async_api import async_playwright
|
|
|
|
async with async_playwright() as playwright:
|
|
# Only the project's server bypasses the deliberately unreachable proxy.
|
|
# The browser can attempt Claude's redirect (testing CSP) but cannot send it.
|
|
browser = await getattr(playwright, engine).launch(
|
|
headless=True,
|
|
proxy={"server": "http://127.0.0.1:9", "bypass": urlsplit(origin).hostname},
|
|
)
|
|
try:
|
|
context = await browser.new_context()
|
|
seen = {}
|
|
callback_ready = asyncio.Event()
|
|
|
|
def callback_request(request):
|
|
if request.url.startswith(CALLBACK + "?"):
|
|
seen["url"] = request.url
|
|
seen["referrer"] = request.headers.get("referer")
|
|
callback_ready.set()
|
|
|
|
page = await context.new_page()
|
|
page.on("request", callback_request)
|
|
await page.goto(login_url)
|
|
await page.locator("input[name=password]").fill(password)
|
|
is_form = (
|
|
lambda request: request.method == "POST" and request.url == origin + "/oauth/login"
|
|
)
|
|
async with (
|
|
page.expect_request(is_form) as request_info,
|
|
page.expect_response(
|
|
lambda response: response.request.method == "POST"
|
|
and response.url == origin + "/oauth/login"
|
|
) as response_info,
|
|
):
|
|
await page.locator("button[value=allow]").click(no_wait_after=True)
|
|
request, response = await request_info.value, await response_info.value
|
|
actual_origin = await request.header_value("origin")
|
|
print(
|
|
"Browser form: Origin=" + str(actual_origin) + ", HTTP " + str(response.status),
|
|
flush=True,
|
|
)
|
|
require(actual_origin == origin, "Browser form sent an unexpected Origin")
|
|
require(response.status == 303, "Browser login form failed")
|
|
await asyncio.wait_for(callback_ready.wait(), timeout=10)
|
|
require(
|
|
seen.get("url") and seen.get("referrer") is None,
|
|
"Callback or referrer protection failed",
|
|
)
|
|
return seen["url"]
|
|
finally:
|
|
await browser.close()
|
|
|
|
|
|
async def check(browser=None):
|
|
values = dotenv_values(ROOT / ".env")
|
|
resource = os.getenv("ZP_MCP_PUBLIC_URL") or values.get("ZP_MCP_PUBLIC_URL", "")
|
|
password = os.getenv("ZP_MCP_LOGIN_PASSWORD") or values.get("ZP_MCP_LOGIN_PASSWORD", "")
|
|
url = urlsplit(resource)
|
|
require(
|
|
url.scheme == "https"
|
|
and url.hostname
|
|
and url.path == "/mcp"
|
|
and not url.query
|
|
and not url.fragment
|
|
and url.username is None
|
|
and url.password is None,
|
|
"Configure the HTTPS MCP URL",
|
|
)
|
|
require(password and len(password) >= 32, "Configure the local OAuth login code")
|
|
origin = url.scheme + "://" + url.netloc
|
|
async with httpx.AsyncClient(timeout=30, follow_redirects=False) as http:
|
|
response = await http.post(resource, json={})
|
|
require(
|
|
response.status_code == 401
|
|
and "resource_metadata=" in response.headers.get("www-authenticate", ""),
|
|
"OAuth discovery challenge failed",
|
|
)
|
|
metadata_url = re.search(
|
|
r'resource_metadata="([^"]+)"', response.headers["www-authenticate"]
|
|
)[1]
|
|
require(
|
|
metadata_url == origin + "/.well-known/oauth-protected-resource/mcp",
|
|
"Unexpected metadata destination",
|
|
)
|
|
response = await http.get(metadata_url)
|
|
require(
|
|
response.status_code == 200 and response.json().get("resource") == resource,
|
|
"Resource metadata failed",
|
|
)
|
|
response = await http.get(origin + "/.well-known/oauth-authorization-server")
|
|
require(response.status_code == 200, "Authorization metadata failed")
|
|
metadata = response.json()
|
|
require(
|
|
all(
|
|
metadata.get(key) == origin + path
|
|
for key, path in {
|
|
"authorization_endpoint": "/authorize",
|
|
"registration_endpoint": "/register",
|
|
"token_endpoint": "/token",
|
|
}.items()
|
|
),
|
|
"Unexpected authorization endpoint destination",
|
|
)
|
|
require(
|
|
"S256" in metadata.get("code_challenge_methods_supported", [])
|
|
and "none" in metadata.get("token_endpoint_auth_methods_supported", []),
|
|
"Missing public-client PKCE support",
|
|
)
|
|
print("OAuth discovery and S256 PKCE metadata: OK", flush=True)
|
|
|
|
response = await http.post(
|
|
origin + "/register",
|
|
json={
|
|
"client_name": "ZP Agent connection verification",
|
|
"redirect_uris": [CALLBACK],
|
|
"token_endpoint_auth_method": "none",
|
|
"grant_types": ["authorization_code", "refresh_token"],
|
|
"response_types": ["code"],
|
|
"scope": "zpwiki:read offline_access",
|
|
},
|
|
)
|
|
require(response.status_code == 201, "Dynamic client registration failed")
|
|
client_id = response.json()["client_id"]
|
|
verifier, state = secrets.token_urlsafe(48), secrets.token_urlsafe(24)
|
|
challenge = (
|
|
base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest())
|
|
.decode()
|
|
.rstrip("=")
|
|
)
|
|
response = await http.get(
|
|
origin + "/authorize",
|
|
params={
|
|
"response_type": "code",
|
|
"client_id": client_id,
|
|
"redirect_uri": CALLBACK,
|
|
"code_challenge": challenge,
|
|
"code_challenge_method": "S256",
|
|
"scope": "zpwiki:read offline_access",
|
|
"state": state,
|
|
"resource": resource,
|
|
},
|
|
)
|
|
require(response.status_code == 302, "Authorization request failed")
|
|
login = urlsplit(response.headers["location"])
|
|
require(
|
|
login.scheme + "://" + login.netloc == origin and login.path == "/oauth/login",
|
|
"Unexpected login destination",
|
|
)
|
|
if browser:
|
|
callback_url = await browser_login(
|
|
response.headers["location"], origin, password, browser
|
|
)
|
|
else:
|
|
response = await http.get(response.headers["location"])
|
|
require(
|
|
response.status_code == 200
|
|
and "HttpOnly" in response.headers.get("set-cookie", "")
|
|
and "Secure" in response.headers.get("set-cookie", ""),
|
|
"Secure login page failed",
|
|
)
|
|
csrf = re.search(r'name="csrf" value="([^"]+)"', response.text)
|
|
require(csrf is not None, "Missing login CSRF field")
|
|
form = {
|
|
"request": parse_qs(login.query)["request"][0],
|
|
"csrf": csrf[1],
|
|
"action": "allow",
|
|
}
|
|
rejected = await http.post(
|
|
origin + "/oauth/login",
|
|
data={**form, "password": "invalid-test-code"},
|
|
headers={"Origin": origin},
|
|
)
|
|
require(rejected.status_code == 401, "Invalid owner code was not rejected")
|
|
response = await http.post(
|
|
origin + "/oauth/login",
|
|
data={**form, "password": password},
|
|
headers={"Origin": origin},
|
|
)
|
|
require(response.status_code == 303, "Owner login and consent failed")
|
|
callback_url = response.headers["location"]
|
|
callback = urlsplit(callback_url)
|
|
require(
|
|
callback.scheme + "://" + callback.netloc + callback.path == CALLBACK,
|
|
"Unexpected callback destination",
|
|
)
|
|
query = parse_qs(callback.query)
|
|
require(query.get("state") == [state] and query.get("code"), "OAuth state or code failed")
|
|
body = {
|
|
"grant_type": "authorization_code",
|
|
"client_id": client_id,
|
|
"code": query["code"][0],
|
|
"code_verifier": verifier,
|
|
"redirect_uri": CALLBACK,
|
|
"resource": resource,
|
|
}
|
|
wrong = await http.post(
|
|
origin + "/token", data={**body, "code_verifier": "incorrect-verifier"}
|
|
)
|
|
require(
|
|
wrong.status_code == 400 and wrong.json().get("error") == "invalid_grant",
|
|
"Invalid PKCE verifier accepted",
|
|
)
|
|
response = await http.post(origin + "/token", data=body)
|
|
require(response.status_code == 200, "Authorization code exchange failed")
|
|
token = response.json()
|
|
refresh = token.get("refresh_token")
|
|
require(
|
|
refresh and 0 < token.get("expires_in", 0) <= 3600,
|
|
"Missing limited token or refresh grant",
|
|
)
|
|
try:
|
|
replay = await http.post(origin + "/token", data=body)
|
|
require(replay.status_code == 400, "Authorization code replay accepted")
|
|
print(
|
|
"DCR, owner login, consent, state, code exchange and replay rejection: OK",
|
|
flush=True,
|
|
)
|
|
async with httpx.AsyncClient(
|
|
timeout=60, headers={"Authorization": "Bearer " + token["access_token"]}
|
|
) as authorized:
|
|
async with streamable_http_client(resource, http_client=authorized) as (
|
|
read,
|
|
write,
|
|
_,
|
|
):
|
|
async with ClientSession(read, write) as session:
|
|
await session.initialize()
|
|
names = {item.name for item in (await session.list_tools()).tools}
|
|
require(
|
|
names == {"zpwiki_rag", "zpwiki_graphrag", "zpwiki_health"},
|
|
"MCP tool discovery failed",
|
|
)
|
|
result = await session.call_tool("zpwiki_health")
|
|
health = result.structuredContent or {}
|
|
require(
|
|
not result.isError and health.get("ready"), "Backend readiness failed"
|
|
)
|
|
versions = [
|
|
health["tools"][engine].get("index_version")
|
|
for engine in ["rag", "graphrag"]
|
|
]
|
|
require(
|
|
versions[0] is not None
|
|
and versions[0] == versions[1] == health.get("index_version"),
|
|
"Different tool data versions",
|
|
)
|
|
for engine in ["rag", "graphrag"]:
|
|
invalid = await session.call_tool("zpwiki_" + engine, {"query": ""})
|
|
require(invalid.isError, "Invalid tool query accepted")
|
|
print("OAuth-authenticated MCP, both tools and common index version: OK", flush=True)
|
|
renewed = await http.post(
|
|
origin + "/token",
|
|
data={
|
|
"grant_type": "refresh_token",
|
|
"client_id": client_id,
|
|
"refresh_token": refresh,
|
|
"resource": resource,
|
|
},
|
|
)
|
|
require(renewed.status_code == 200, "Refresh failed")
|
|
new = renewed.json()
|
|
require(new["refresh_token"] != refresh, "Refresh token did not rotate")
|
|
retired = await http.get(
|
|
origin + "/health", headers={"Authorization": "Bearer " + token["access_token"]}
|
|
)
|
|
require(retired.status_code == 401, "Old access token still active after refresh")
|
|
refresh = new["refresh_token"]
|
|
revoked = await http.post(
|
|
origin + "/revoke", data={"client_id": client_id, "token": refresh}
|
|
)
|
|
require(revoked.status_code == 200, "Revocation failed")
|
|
denied = await http.get(
|
|
origin + "/health", headers={"Authorization": "Bearer " + new["access_token"]}
|
|
)
|
|
require(denied.status_code == 401, "Revoked grant still active")
|
|
print("Refresh rotation and grant revocation: OK; external LLM calls: 0", flush=True)
|
|
finally:
|
|
await http.post(origin + "/revoke", data={"client_id": client_id, "token": refresh})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--browser",
|
|
choices=["firefox", "chromium"],
|
|
help="Submit the login form in an isolated Playwright browser",
|
|
)
|
|
args = parser.parse_args()
|
|
try:
|
|
asyncio.run(check(args.browser))
|
|
except Exception as error:
|
|
# Never include HTTP bodies, credentials, token responses or callback URLs.
|
|
print(
|
|
"OAuth verification failed: "
|
|
+ (str(error) if type(error) is RuntimeError else type(error).__name__),
|
|
file=sys.stderr,
|
|
)
|
|
raise SystemExit(1)
|