#!/usr/bin/env python3
"""bolt HTTP 模式端到端测试"""
import hashlib, http.client, json, os, shutil, subprocess, sys, tempfile, threading, time
from urllib.parse import urlsplit

# 二进制路径:环境变量 BOLT_BIN 优先,默认为仓库内的 release 产物
BOLT = os.environ.get("BOLT_BIN") or os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "target", "release", "bolt")
PORT = 18080
PORT2 = 18081  # 带 token + --file-root 的第二个实例
TOKEN = "e2e-secret"
passed = failed = 0
def check(name, cond, detail=""):
    global passed, failed
    if cond:
        passed += 1
        print(f"  PASS {name}")
    else:
        failed += 1
        print(f"  FAIL {name} {detail}")

tmpdir = tempfile.mkdtemp(prefix="bolt-http-e2e-")
proc = subprocess.Popen([BOLT, "--http", f"127.0.0.1:{PORT}"], stderr=subprocess.PIPE)
proc2 = subprocess.Popen([BOLT, "--http", f"127.0.0.1:{PORT2}", "--token", TOKEN, "--file-root", tmpdir], stderr=subprocess.PIPE)
time.sleep(0.5)

def request(method, path, body=None, headers=None, port=PORT, timeout=30):
    c = http.client.HTTPConnection("127.0.0.1", port, timeout=timeout)
    c.request(method, path, body=body, headers=headers or {})
    r = c.getresponse()
    data = r.read()
    hdrs = {k.lower(): v for k, v in r.getheaders()}
    c.close()
    return r.status, hdrs, data

def post(obj, auth=None, origin=None, accept="application/json", path="/mcp", sid=None, port=PORT):
    headers = {"Content-Type": "application/json", "Accept": accept}
    if auth: headers["Authorization"] = auth
    if origin: headers["Origin"] = origin
    if sid: headers["Mcp-Session-Id"] = sid
    st, hdrs, data = request("POST", path, json.dumps(obj), headers, port=port)
    return st, hdrs.get("content-type", ""), data

def post_h(obj, sid=None, accept="application/json", port=PORT, auth=None):
    """返回 (status, headers, data)"""
    headers = {"Content-Type": "application/json", "Accept": accept}
    if sid: headers["Mcp-Session-Id"] = sid
    if auth: headers["Authorization"] = auth
    return request("POST", "/mcp", json.dumps(obj), headers, port=port)

def call(name, args, rid=100, sid=None, port=PORT, auth=None):
    st, _, data = post_h({"jsonrpc": "2.0", "id": rid, "method": "tools/call", "params": {"name": name, "arguments": args}}, sid=sid, port=port, auth=auth)
    return json.loads(data)

def sse_events(data):
    return [json.loads(line[len(b"data: "):]) for line in data.split(b"\r\n") if line.startswith(b"data: ")]

INIT = {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "t", "version": "1"}}}

# 1. initialize(JSON 响应 + 签发 Mcp-Session-Id)
st, hdrs, data = post_h(INIT)
m = json.loads(data)
sid_a = hdrs.get("mcp-session-id")
check("init json", st == 200 and "application/json" in hdrs.get("content-type", "") and m["result"]["serverInfo"]["name"] == "bolt", f"{st} {data[:100]}")
check("init issues session id", bool(sid_a) and len(sid_a) == 32, str(hdrs))

# 2. SSE 响应(accept 含 text/event-stream)
st, ct, data = post({"jsonrpc": "2.0", "id": 2, "method": "ping"}, accept="application/json, text/event-stream")
sseline = data.split(b"data: ")[1].split(b"\r\n")[0].strip()
m = json.loads(sseline)
check("ping sse", st == 200 and "text/event-stream" in ct and m.get("result") == {}, f"{st} {ct} {data[:150]}")

# 3. 通知 → 202
st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/initialized"})
check("notif 202", st == 202 and len(data) == 0, f"{st} {data[:50]}")

# 4. tools/call
st, ct, data = post({"jsonrpc": "2.0", "id": 3, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "echo HTTP_OK && exit 5"}}})
m = json.loads(data)
sc = m["result"]["structuredContent"]
check("run over http", sc["exit_code"] == 5 and sc["stdout"].strip() == "HTTP_OK", str(sc))

# 5. GET SSE 流:开流收心跳,同作用域命令的进度通知推送到流上
got_heartbeat = threading.Event()
msgs = []
def sse_get():
    c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=35)
    c.request("GET", "/mcp", headers={"Accept": "text/event-stream"})
    r = c.getresponse()
    assert r.status == 200 and "text/event-stream" in r.getheader("Content-Type")
    buf = b""
    end = time.time() + 30
    while time.time() < end:
        chunk = r.read(1)
        if not chunk: break
        buf += chunk
        if buf.endswith(b"\r\n\r\n"):
            msgs.append(buf.decode(errors="replace"))
            buf = b""
            if msgs and (": ping" in msgs[-1]):
                got_heartbeat.set()
                break
    c.close()
t = threading.Thread(target=sse_get)
t.start()
time.sleep(0.7)
# 发一个带进度 token 的命令(约 1.2s)
st, ct, data = post({"jsonrpc": "2.0", "id": 4, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "for x in 1 2 3 4; do echo h-$x; sleep 0.3; done"}, "_meta": {"progressToken": "h1"}}})
m = json.loads(data)
check("long cmd done", m["result"]["structuredContent"]["stdout"].count("h-") == 4, str(m)[:100])
saw_progress = any("notifications/progress" in x and "h1" in x for x in msgs)
check("sse got progress on GET stream", saw_progress, str(msgs)[:200])
t.join(timeout=25)
check("sse heartbeat", got_heartbeat.is_set(), str(msgs)[:200])

# 5b. POST + SSE:进度通知与结果在同一个响应流里
st, hdrs, data = post_h({"jsonrpc": "2.0", "id": 41, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "for x in 1 2 3; do echo p-$x; sleep 0.3; done"}, "_meta": {"progressToken": "p1"}}}, accept="application/json, text/event-stream", sid=sid_a)
evs = sse_events(data)
prog = [e for e in evs if e.get("method") == "notifications/progress" and e["params"]["progressToken"] == "p1"]
final = [e for e in evs if e.get("id") == 41]
check("post sse progress + result", st == 200 and prog and final and final[0]["result"]["structuredContent"]["stdout"].count("p-") == 3, str(evs)[:300])
check("post sse result is last", evs and evs[-1].get("id") == 41, str(evs)[-200:])

# 6. 取消:连接 A 长命令,连接 B 发 cancelled(同一作用域)
def long_cmd():
    return post({"jsonrpc": "2.0", "id": 6, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": "sleep 30"}}})
res = {}
t2 = threading.Thread(target=lambda: res.update(r=long_cmd()))
t2.start()
time.sleep(0.5)
# 其他会话发同 id 的取消不能误杀
st, _, _ = post({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": 6}}, sid=sid_a)
time.sleep(0.3)
check("cancel from other scope ignored", t2.is_alive())
st, ct, data = post({"jsonrpc": "2.0", "method": "notifications/cancelled", "params": {"requestId": 6}})
check("cancel notif 202", st == 202)
t2.join(timeout=10)
m = json.loads(res["r"][2])
check("http cancel works", m["result"]["structuredContent"]["cancelled"] is True, str(m)[:150])

# 7. 并发两个命令
def run(cmd, rid):
    return post({"jsonrpc": "2.0", "id": rid, "method": "tools/call", "params": {"name": "run_command", "arguments": {"command": cmd}}})
results = {}
ts = [threading.Thread(target=lambda c=cmd, i=i: results.update({i: run(c, i)})) for i, cmd in [(10, "sleep 0.5; echo A"), (11, "sleep 0.1; echo B")]]
for t in ts: t.start()
for t in ts: t.join(timeout=10)
check("http concurrent", json.loads(results[10][2])["result"]["structuredContent"]["stdout"].strip() == "A" and json.loads(results[11][2])["result"]["structuredContent"]["stdout"].strip() == "B")

# 8. 错误路径
st, _, _ = post({"jsonrpc": "2.0", "id": 99, "method": "nope"})
check("unknown method", st == 200)
st, _, _ = request("POST", "/mcp", b"this is {not json", {"Content-Type": "application/json", "Accept": "application/json"})
check("bad json -> 400", st == 400, str(st))
st, _, data = post({"jsonrpc": "2.0", "id": 99, "method": "tools/call", "params": {"name": "nope", "arguments": {}}})
check("unknown tool -32602", json.loads(data)["error"]["code"] == -32602)
st, _, _ = request("GET", "/other", headers={"Accept": "text/event-stream"})
check("404 other path", st == 404)
st, _, _ = request("PUT", "/mcp")
check("405 method", st == 405)
st, hdrs, _ = request("OPTIONS", "/mcp")
check("options cors", st == 204 and hdrs.get("access-control-allow-origin") == "*" and "Mcp-Session-Id" in hdrs.get("access-control-allow-headers", ""))

# 9. origin 校验
st, ct, data = post({"jsonrpc": "2.0", "id": 50, "method": "ping"}, origin="https://evil.example.com")
check("evil origin 403", st == 403, str(st))
st, ct, data = post({"jsonrpc": "2.0", "id": 51, "method": "ping"}, origin="http://localhost:4173")
check("localhost origin ok", st == 200, str(st))

# 10. GET 无 Accept 头也返回 SSE(宽松)
c = http.client.HTTPConnection("127.0.0.1", PORT, timeout=5)
c.request("GET", "/mcp")
r = c.getresponse()
c.close()
check("get sse default", r.status == 200 and "text/event-stream" in r.getheader("Content-Type"))

# 11. MCP 会话:隔离、未知 ID → 404、DELETE 终止
st, hdrs, _ = post_h(INIT)
sid_b = hdrs.get("mcp-session-id")
r = call("shell_open", {}, sid=sid_a)
shell_id = r["result"]["structuredContent"]["session_id"]
r = call("shell_read", {"session_id": shell_id, "timeout_ms": 10}, sid=sid_b)
check("shell session isolated", r["result"].get("isError") is True, str(r))
r = call("shell_read", {"session_id": shell_id, "timeout_ms": 10}, sid=sid_a)
check("shell session owner ok", r["result"].get("isError") is None, str(r))
st, _, _ = post_h({"jsonrpc": "2.0", "id": 60, "method": "ping"}, sid="0" * 32)
check("unknown session 404", st == 404, str(st))
st, _, _ = request("DELETE", "/mcp", headers={"Mcp-Session-Id": sid_a})
check("delete session", st == 200, str(st))
st, _, _ = post_h({"jsonrpc": "2.0", "id": 61, "method": "ping"}, sid=sid_a)
check("deleted session 404", st == 404, str(st))
r = call("shell_read", {"session_id": shell_id, "timeout_ms": 10}, sid=sid_b)
check("deleted session shells closed", r["result"].get("isError") is True)

# 12. HTTP 客户端不能用 notifications/exit 关停服务
post({"jsonrpc": "2.0", "method": "notifications/exit"})
time.sleep(0.4)
st, _, _ = post({"jsonrpc": "2.0", "id": 70, "method": "ping"})
check("exit notif ignored over http", st == 200 and proc.poll() is None)

# 13. tools/list 含 file_link
st, _, data = post({"jsonrpc": "2.0", "id": 71, "method": "tools/list"})
names = [t["name"] for t in json.loads(data)["result"]["tools"]]
check("tools/list http", names[-4:] == ["file_read", "file_write", "file_stat", "file_link"], str(names))

# 14. /files:Bearer 认证、签名链接、Range、chunked 上传、目录限制
AUTH = {"Authorization": f"Bearer {TOKEN}"}
fpath = os.path.join(tmpdir, "data.bin")
payload = os.urandom(300_000)
# 未认证请求在读取 body 前即返回 401 并关闭连接,这里用小 body 避免客户端写入时 broken pipe
st, _, _ = request("PUT", f"/files?path={fpath}", b"x", port=PORT2)
check("files put needs token", st == 401 and not os.path.exists(fpath), str(st))
st, _, data = request("PUT", f"/files?path={fpath}", payload, AUTH, port=PORT2)
info = json.loads(data)
check("files put", st == 201 and info["sha256"] == hashlib.sha256(payload).hexdigest(), f"{st} {data[:200]}")
st, hdrs, data = request("GET", f"/files?path={fpath}", headers=AUTH, port=PORT2)
check("files get", st == 200 and data == payload and hdrs.get("accept-ranges") == "bytes")
st, hdrs, data = request("GET", f"/files?path={fpath}", headers={**AUTH, "Range": "bytes=100-199"}, port=PORT2)
check("files range", st == 206 and data == payload[100:200] and hdrs.get("content-range") == f"bytes 100-199/{len(payload)}", f"{st} {hdrs}")
st, _, _ = request("GET", "/files?path=/etc/hosts", headers=AUTH, port=PORT2)
check("files outside root 403", st == 403, str(st))
# chunked 上传(http.client 对迭代器 body 使用 chunked 编码)
cpath = os.path.join(tmpdir, "chunked.txt")
c = http.client.HTTPConnection("127.0.0.1", PORT2, timeout=10)
c.request("PUT", f"/files?path={cpath}&mkdirs=1", body=iter([b"hello ", b"chunked ", b"world"]), headers=AUTH, encode_chunked=True)
r = c.getresponse(); r.read(); c.close()
check("files chunked put", r.status == 201 and open(cpath, "rb").read() == b"hello chunked world", str(r.status))
# 签名链接:不带 token 也能访问,且仅限签名的路径/方法
st, _, data = post_h(INIT, port=PORT2, auth=AUTH["Authorization"])
r = call("file_link", {"path": fpath}, port=PORT2, auth=AUTH["Authorization"])
url = r["result"]["structuredContent"]["url"]
u = urlsplit(url)
check("file_link url", u.netloc == f"127.0.0.1:{PORT2}" and u.path == "/files" and "curl" in r["result"]["structuredContent"]["curl"], url)
st, _, data = request("GET", f"{u.path}?{u.query}", port=PORT2)
check("signed get", st == 200 and data == payload, str(st))
st, _, _ = request("PUT", f"{u.path}?{u.query}", b"x", port=PORT2)
check("signed get cannot put", st == 403, str(st))
r = call("file_link", {"path": os.path.join(tmpdir, "up", "new.txt"), "method": "PUT", "mkdirs": True}, port=PORT2, auth=AUTH["Authorization"])
u = urlsplit(r["result"]["structuredContent"]["url"])
st, _, data = request("PUT", f"{u.path}?{u.query}", b"uploaded via link", {"Expect": "100-continue"}, port=PORT2)
check("signed put", st == 201 and open(os.path.join(tmpdir, "up", "new.txt"), "rb").read() == b"uploaded via link", f"{st} {data[:200]}")
st, _, _ = request("GET", f"{u.path}?{u.query.replace('sig=', 'sig=0')}", port=PORT2)
check("bad signature 403", st == 403, str(st))
r = call("file_link", {"path": "/etc/hosts"}, port=PORT2, auth=AUTH["Authorization"])
check("file_link outside root", r["result"].get("isError") is True)

proc.kill()
proc2.kill()
err = proc.stderr.read().decode()
shutil.rmtree(tmpdir, ignore_errors=True)
print(f"\n=== HTTP: {passed} passed, {failed} failed ===")
if failed and err.strip():
    print("bolt stderr:", err[:400])
sys.exit(1 if failed else 0)