#!/usr/bin/env bash
set -euo pipefail

usage() {
  cat <<'EOF'
Usage: Scripts/burrow-proxy-selftest

Run a controlled Burrow proxy packet-runtime test without changing system proxy
settings or using the user's Burrow database. The script starts a local Trojan
TLS server, starts a Burrow daemon against a temporary database/socket, imports a
temporary ProxySubscription pointed at that server, then runs the packet-level
proxy-tcp-probe through daemon DNS and TCP packet streaming.
EOF
}

if [[ "${1:-}" == "-h" || "${1:-}" == "--help" ]]; then
  usage
  exit 0
fi

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
burrow_bin="${BURROW_BIN:-$repo_root/target/debug/burrow}"

if [[ ! -x "$burrow_bin" ]]; then
  (cd "$repo_root" && cargo build -p burrow)
fi

if ! command -v openssl >/dev/null 2>&1; then
  echo "openssl is required for the local TLS server certificate" >&2
  exit 2
fi

tmpdir="$(mktemp -d "${TMPDIR:-/tmp}/burrow-proxy-selftest.XXXXXX")"
daemon_pid=""
server_pid=""
status=0

cleanup() {
  if [[ -n "$server_pid" ]]; then
    kill "$server_pid" >/dev/null 2>&1 || true
  fi
  if [[ -n "$daemon_pid" ]]; then
    kill "$daemon_pid" >/dev/null 2>&1 || true
  fi
  if [[ "${BURROW_SELFTEST_KEEP:-}" == "1" || "$status" -ne 0 ]]; then
    echo "preserving self-test artifacts: $tmpdir" >&2
  else
    rm -rf "$tmpdir"
  fi
}
trap 'status=$?; cleanup' EXIT

cert="$tmpdir/cert.pem"
key="$tmpdir/key.pem"
socket="$tmpdir/burrow.sock"
port_file="$tmpdir/server.port"
server_result="$tmpdir/server-result.json"
payload="$tmpdir/proxy-payload.json"
daemon_log="$tmpdir/daemon.log"
server_log="$tmpdir/server.log"

openssl req \
  -x509 \
  -newkey rsa:2048 \
  -nodes \
  -subj /CN=localhost \
  -keyout "$key" \
  -out "$cert" \
  -days 1 >/dev/null 2>&1

python3 - "$port_file" "$server_result" "$cert" "$key" >"$server_log" 2>&1 <<'PY' &
from __future__ import annotations

import hashlib
import json
import socket
import ssl
import struct
import sys

port_file, result_file, cert_file, key_file = sys.argv[1:]
password = "secret"


def read_exact(sock: ssl.SSLSocket, length: int) -> bytes:
    chunks: list[bytes] = []
    remaining = length
    while remaining:
        chunk = sock.recv(remaining)
        if not chunk:
            raise EOFError(f"expected {remaining} more bytes")
        chunks.append(chunk)
        remaining -= len(chunk)
    return b"".join(chunks)


def read_socks_addr(sock: ssl.SSLSocket) -> dict[str, object]:
    atyp = read_exact(sock, 1)[0]
    if atyp == 1:
        host = socket.inet_ntop(socket.AF_INET, read_exact(sock, 4))
    elif atyp == 4:
        host = socket.inet_ntop(socket.AF_INET6, read_exact(sock, 16))
    elif atyp == 3:
        length = read_exact(sock, 1)[0]
        host = read_exact(sock, length).decode("idna")
    else:
        raise ValueError(f"unsupported SOCKS address type {atyp}")
    port = struct.unpack("!H", read_exact(sock, 2))[0]
    return {"atyp": atyp, "host": host, "port": port}


def read_http_headers(sock: ssl.SSLSocket) -> bytes:
    data = bytearray()
    while b"\r\n\r\n" not in data:
        chunk = sock.recv(4096)
        if not chunk:
            break
        data.extend(chunk)
        if len(data) > 65535:
            raise ValueError("HTTP request headers exceeded 64 KiB")
    return bytes(data)


ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_SERVER)
ctx.load_cert_chain(cert_file, key_file)
ctx.set_alpn_protocols(["h2", "http/1.1"])

listener = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
listener.bind(("127.0.0.1", 0))
listener.listen(1)
port = listener.getsockname()[1]
with open(port_file, "w", encoding="utf-8") as f:
    f.write(str(port))

raw, peer = listener.accept()
raw.settimeout(10)
result: dict[str, object] = {"peer": str(peer)}
try:
    with ctx.wrap_socket(raw, server_side=True) as tls:
        tls.settimeout(10)
        result["alpn"] = tls.selected_alpn_protocol()
        got_hash = read_exact(tls, 56)
        expected_hash = hashlib.sha224(password.encode()).hexdigest().encode()
        result["password_hash_ok"] = got_hash == expected_hash
        if not result["password_hash_ok"]:
            raise ValueError("unexpected Trojan password hash")
        if read_exact(tls, 2) != b"\r\n":
            raise ValueError("missing Trojan password delimiter")
        command = read_exact(tls, 1)[0]
        result["command"] = command
        target = read_socks_addr(tls)
        result["target"] = target
        if read_exact(tls, 2) != b"\r\n":
            raise ValueError("missing Trojan header delimiter")
        request = read_http_headers(tls)
        result["request_first_line"] = request.splitlines()[0].decode("utf-8", "replace")
        body = json.dumps(
            {
                "ok": True,
                "target": target,
                "request_first_line": result["request_first_line"],
            },
            sort_keys=True,
        ).encode()
        tls.sendall(
            b"HTTP/1.1 200 OK\r\n"
            + f"Content-Length: {len(body)}\r\n".encode()
            + b"Connection: close\r\n"
            + b"Content-Type: application/json\r\n\r\n"
            + body
        )
finally:
    with open(result_file, "w", encoding="utf-8") as f:
        json.dump(result, f, ensure_ascii=False, indent=2, sort_keys=True)
    listener.close()
PY
server_pid=$!

for _ in {1..100}; do
  if [[ -s "$port_file" ]]; then
    break
  fi
  if ! kill -0 "$server_pid" >/dev/null 2>&1; then
    echo "local Trojan server exited before listening" >&2
    cat "$server_log" >&2 || true
    exit 1
  fi
  sleep 0.05
done

if [[ ! -s "$port_file" ]]; then
  echo "timed out waiting for local Trojan server" >&2
  cat "$server_log" >&2 || true
  exit 1
fi
server_port="$(cat "$port_file")"

cat >"$payload" <<EOF
{
  "name": "burrow proxy selftest",
  "source": { "url": "selftest://local-trojan" },
  "detected_format": "uri-list",
  "selected_ordinal": 1,
  "selected_name": "local trojan",
  "nodes": [
    {
      "ordinal": 1,
      "name": "local trojan",
      "protocol": "trojan",
      "server": "127.0.0.1",
      "port": $server_port,
      "warnings": [],
      "config": {
        "type": "trojan",
        "password": "secret",
        "sni": "localhost",
        "alpn": [],
        "alpn_present": true,
        "skip_cert_verify": true,
        "network": "tcp",
        "udp": true,
        "client_fingerprint": null,
        "fingerprint": null
      }
    }
  ],
  "warnings": []
}
EOF

(
  cd "$tmpdir"
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" daemon >"$daemon_log" 2>&1
) &
daemon_pid=$!

for _ in {1..100}; do
  if [[ -S "$socket" ]]; then
    break
  fi
  if ! kill -0 "$daemon_pid" >/dev/null 2>&1; then
    echo "Burrow daemon exited before creating socket" >&2
    cat "$daemon_log" >&2 || true
    exit 1
  fi
  sleep 0.05
done

if [[ ! -S "$socket" ]]; then
  echo "timed out waiting for Burrow daemon socket" >&2
  cat "$daemon_log" >&2 || true
  exit 1
fi

(
  cd "$tmpdir"
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" network-add 1 3 "$payload"
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" tunnel-config
  BURROW_SOCKET_PATH="$socket" "$burrow_bin" proxy-tcp-probe \
    --dns-name selftest.invalid \
    --remote 1.1.1.1:80 \
    --timeout-ms 8000
)

wait "$server_pid"
server_pid=""

echo
echo "Local Trojan server observed:"
cat "$server_result"
echo