60 lines
3.0 KiB
Python
60 lines
3.0 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""E2E: OAuth write token -> /write/sse and tools/list via SSE."""
|
||
|
|
import json, urllib.request, urllib.parse, hashlib, base64, secrets, http.client, time
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
BASE = "http://127.0.0.1:9091"
|
||
|
|
env_path = Path("/mnt/ai-storage/projects/ctxd/app/.env")
|
||
|
|
approval_key = ""
|
||
|
|
for line in env_path.read_text().splitlines():
|
||
|
|
if line.startswith("OAUTH_APPROVAL_KEY="):
|
||
|
|
approval_key = line.split("=", 1)[1]
|
||
|
|
break
|
||
|
|
if not approval_key:
|
||
|
|
raise SystemExit("no OAUTH_APPROVAL_KEY")
|
||
|
|
|
||
|
|
class NR(urllib.request.HTTPRedirectHandler):
|
||
|
|
def redirect_request(self, *a): return None
|
||
|
|
opener = urllib.request.build_opener(NR)
|
||
|
|
|
||
|
|
r = opener.open(urllib.request.Request(
|
||
|
|
f"{BASE}/oauth/register",
|
||
|
|
data=json.dumps({"redirect_uris": ["http://localhost:9999/cb"], "client_name": "write-test", "scope": "ctxd.read ctxd.write"}).encode(),
|
||
|
|
headers={"Content-Type": "application/json"}, method="POST"), timeout=15)
|
||
|
|
client = json.loads(r.read())
|
||
|
|
cid, secret = client["client_id"], client["client_secret"]
|
||
|
|
|
||
|
|
cv = secrets.token_urlsafe(32)
|
||
|
|
cc = base64.urlsafe_b64encode(hashlib.sha256(cv.encode()).digest()).rstrip(b"=").decode()
|
||
|
|
params = {"response_type": "code", "client_id": cid, "redirect_uri": "http://localhost:9999/cb",
|
||
|
|
"code_challenge": cc, "code_challenge_method": "S256", "scope": "ctxd.read ctxd.write", "state": "t"}
|
||
|
|
try:
|
||
|
|
opener.open(urllib.request.Request(f"{BASE}/oauth/authorize?" + urllib.parse.urlencode(params),
|
||
|
|
data=urllib.parse.urlencode({"approval_key": approval_key}).encode(),
|
||
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"}, method="POST"), timeout=15)
|
||
|
|
except urllib.error.HTTPError as e:
|
||
|
|
loc = e.headers.get("Location", "")
|
||
|
|
code = dict(urllib.parse.parse_qsl(urllib.parse.urlparse(loc).query)).get("code", "")
|
||
|
|
|
||
|
|
r3 = opener.open(urllib.request.Request(f"{BASE}/oauth/token",
|
||
|
|
data=urllib.parse.urlencode({"grant_type": "authorization_code", "code": code, "redirect_uri": "http://localhost:9999/cb",
|
||
|
|
"client_id": cid, "client_secret": secret, "code_verifier": cv}).encode(),
|
||
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"}, method="POST"), timeout=15)
|
||
|
|
token = json.loads(r3.read())["access_token"]
|
||
|
|
|
||
|
|
conn = http.client.HTTPConnection("127.0.0.1", 9091, timeout=5)
|
||
|
|
conn.request("GET", "/write/sse", headers={"Authorization": f"Bearer {token}"})
|
||
|
|
resp = conn.getresponse()
|
||
|
|
print("write/sse status:", resp.status, "ctype:", resp.getheader("content-type"))
|
||
|
|
chunk = resp.read(500).decode(errors="replace")
|
||
|
|
print("body[:200]:", chunk[:200])
|
||
|
|
conn.close()
|
||
|
|
|
||
|
|
# POST initialize to /write/messages (needs session from SSE - simplified check)
|
||
|
|
conn2 = http.client.HTTPConnection("127.0.0.1", 9091, timeout=10)
|
||
|
|
body = json.dumps({"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}})
|
||
|
|
conn2.request("POST", "/write/messages", body=body,
|
||
|
|
headers={"Authorization": f"Bearer {token}", "Content-Type": "application/json"})
|
||
|
|
r2 = conn2.getresponse()
|
||
|
|
print("write/messages POST:", r2.status, r2.read(300).decode(errors="replace")[:200])
|
||
|
|
conn2.close()
|