#!/usr/bin/env python3
"""Probe a Burrow proxy subscription node without enabling system proxying.

This script intentionally bypasses the Burrow daemon, Network Extension,
routes, DNS settings, and browsers. It loads the selected proxy node from the
Burrow SQLite database and speaks the provider protocol directly.
"""

from __future__ import annotations

import argparse
import hashlib
import ipaddress
import json
import os
import socket
import sqlite3
import ssl
import struct
import sys
import time
import urllib.parse
import urllib.request
from dataclasses import dataclass
from pathlib import Path
from typing import Any


DEFAULT_APP_GROUP = "2KPBQHRJ2D.com.tianruichen.burrow"
DEFAULT_DB = (
    Path.home()
    / "Library"
    / "Group Containers"
    / DEFAULT_APP_GROUP
    / "burrow.db"
)
DEFAULT_TARGET = "1.1.1.1:80"
DEFAULT_HTTP_HOST = "one.one.one.one"
DEFAULT_DNS_SERVER = "1.1.1.1:53"
DEFAULT_DNS_NAME = "google.com"
DEFAULT_TIMEOUT = 5.0
DOH_PROVIDERS = {
    "doh-google": "https://dns.google/resolve",
    "doh-cloudflare": "https://cloudflare-dns.com/dns-query",
}

TROJAN_CMD_CONNECT = 0x01
TROJAN_CMD_UDP_ASSOCIATE = 0x03


@dataclass(frozen=True)
class Endpoint:
    host: str
    port: int


@dataclass(frozen=True)
class ProbeResult:
    ok: bool
    detail: str


def parse_endpoint(value: str, default_port: int | None = None) -> Endpoint:
    text = value.strip()
    if not text:
        raise ValueError("endpoint must not be empty")

    if text.startswith("["):
        end = text.find("]")
        if end == -1:
            raise ValueError(f"invalid bracketed IPv6 endpoint: {value}")
        host = text[1:end]
        rest = text[end + 1 :]
        if rest.startswith(":"):
            port = int(rest[1:])
        elif default_port is not None:
            port = default_port
        else:
            raise ValueError(f"endpoint is missing a port: {value}")
        return Endpoint(host, port)

    if text.count(":") == 1:
        host, port_text = text.rsplit(":", 1)
        return Endpoint(host, int(port_text))

    try:
        ipaddress.ip_address(text)
    except ValueError:
        if ":" in text:
            if default_port is None:
                raise ValueError(f"IPv6 endpoint is missing a port: {value}")
            return Endpoint(text, default_port)
        if default_port is None:
            raise ValueError(f"endpoint is missing a port: {value}")
        return Endpoint(text, default_port)

    if default_port is None:
        raise ValueError(f"endpoint is missing a port: {value}")
    return Endpoint(text, default_port)


def encode_socks_addr(endpoint: Endpoint) -> bytes:
    try:
        ip = ipaddress.ip_address(endpoint.host)
    except ValueError:
        host = endpoint.host.encode("idna")
        if len(host) > 255:
            raise ValueError(f"SOCKS domain is too long: {endpoint.host}")
        return bytes([3, len(host)]) + host + struct.pack("!H", endpoint.port)

    if ip.version == 4:
        return bytes([1]) + ip.packed + struct.pack("!H", endpoint.port)
    return bytes([4]) + ip.packed + struct.pack("!H", endpoint.port)


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) -> Endpoint:
    atyp = read_exact(sock, 1)[0]
    if atyp == 1:
        host = str(ipaddress.IPv4Address(read_exact(sock, 4)))
    elif atyp == 4:
        host = str(ipaddress.IPv6Address(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 in response: {atyp}")
    port = struct.unpack("!H", read_exact(sock, 2))[0]
    return Endpoint(host, port)


def trojan_password_hash(password: str) -> bytes:
    return hashlib.sha224(password.encode()).hexdigest().encode()


def trojan_header(password: str, command: int, target: Endpoint) -> bytes:
    return (
        trojan_password_hash(password)
        + b"\r\n"
        + bytes([command])
        + encode_socks_addr(target)
        + b"\r\n"
    )


def trojan_udp_frame(target: Endpoint, payload: bytes) -> bytes:
    return encode_socks_addr(target) + struct.pack("!H", len(payload)) + b"\r\n" + payload


def build_dns_query(name: str) -> tuple[int, bytes]:
    query_id = int(time.time() * 1000) & 0xFFFF
    labels = [label for label in name.rstrip(".").split(".") if label]
    qname = b"".join(bytes([len(label)]) + label.encode("ascii") for label in labels) + b"\x00"
    header = struct.pack("!HHHHHH", query_id, 0x0100, 1, 0, 0, 0)
    question = qname + struct.pack("!HH", 1, 1)
    return query_id, header + question


def parse_dns_response(query_id: int, payload: bytes) -> str:
    if len(payload) < 12:
        return f"truncated DNS response: {len(payload)} bytes"
    response_id, flags, questions, answers, authority, additional = struct.unpack(
        "!HHHHHH", payload[:12]
    )
    rcode = flags & 0x000F
    qr = bool(flags & 0x8000)
    id_note = "matching id" if response_id == query_id else f"id mismatch {response_id:#06x}"
    return (
        f"{len(payload)} bytes, qr={qr}, rcode={rcode}, answers={answers}, "
        f"authority={authority}, additional={additional}, questions={questions}, {id_note}"
    )


def load_payload(db_path: Path, network_id: int | None) -> tuple[int, dict[str, Any]]:
    if not db_path.exists():
        raise FileNotFoundError(f"Burrow database not found: {db_path}")
    conn = sqlite3.connect(str(db_path))
    conn.row_factory = sqlite3.Row
    try:
        if network_id is None:
            row = conn.execute(
                "select id,payload from network where type = 'ProxySubscription' order by idx limit 1"
            ).fetchone()
        else:
            row = conn.execute(
                "select id,payload from network where id = ? and type = 'ProxySubscription'",
                (network_id,),
            ).fetchone()
    finally:
        conn.close()
    if row is None:
        selector = "first ProxySubscription" if network_id is None else f"ProxySubscription id={network_id}"
        raise LookupError(f"could not find {selector} in {db_path}")
    payload_blob = row["payload"]
    if isinstance(payload_blob, bytes):
        payload_text = payload_blob.decode("utf-8", errors="replace")
    else:
        payload_text = payload_blob or "{}"
    return int(row["id"]), json.loads(payload_text)


def choose_node(
    payload: dict[str, Any],
    node_name: str | None = None,
    node_ordinal: int | None = None,
) -> dict[str, Any]:
    nodes = payload.get("nodes") or []
    if node_name is not None:
        for node in nodes:
            if node.get("name") == node_name:
                return node
        raise LookupError(f"ProxySubscription payload has no node named {node_name!r}")
    if node_ordinal is not None:
        for node in nodes:
            if node.get("ordinal") == node_ordinal:
                return node
        raise LookupError(f"ProxySubscription payload has no node with ordinal {node_ordinal}")

    selected_name = payload.get("selected_name")
    selected_ordinal = payload.get("selected_ordinal")
    for node in nodes:
        if selected_name is not None and node.get("name") == selected_name:
            return node
        if selected_ordinal is not None and node.get("ordinal") == selected_ordinal:
            return node
    if nodes:
        return nodes[0]
    raise LookupError("ProxySubscription payload has no compatible nodes")


def resolved_stream_addresses(endpoint: Endpoint) -> list[tuple[int, tuple[Any, ...]]]:
    addresses: list[tuple[int, tuple[Any, ...]]] = []
    seen: set[tuple[Any, ...]] = set()
    for family, _, _, _, sockaddr in socket.getaddrinfo(
        endpoint.host, endpoint.port, type=socket.SOCK_STREAM
    ):
        key = (family, sockaddr)
        if key not in seen:
            seen.add(key)
            addresses.append((family, sockaddr))
    return addresses


def resolve_doh(host: str, port: int, provider: str, timeout: float) -> list[Endpoint]:
    base = DOH_PROVIDERS[provider]
    url = base + "?" + urllib.parse.urlencode({"name": host, "type": "A"})
    request = urllib.request.Request(
        url,
        headers={"Accept": "application/dns-json", "User-Agent": "burrow-proxy-ladder"},
    )
    with urllib.request.urlopen(request, timeout=timeout) as response:
        body = json.load(response)
    endpoints: list[Endpoint] = []
    for answer in body.get("Answer") or []:
        if answer.get("type") != 1:
            continue
        address = str(answer.get("data") or "")
        try:
            ipaddress.IPv4Address(address)
        except ValueError:
            continue
        endpoint = Endpoint(address, port)
        if endpoint not in endpoints:
            endpoints.append(endpoint)
    return endpoints


def resolve_doh_all(host: str, port: int, timeout: float) -> list[Endpoint]:
    endpoints: list[Endpoint] = []
    errors: list[str] = []
    for provider in DOH_PROVIDERS:
        try:
            provider_endpoints = resolve_doh(host, port, provider, timeout)
        except BaseException as exc:
            errors.append(f"{provider}={type(exc).__name__}: {exc}")
            continue
        print(
            f"{provider}_answers="
            + ",".join(f"{endpoint.host}:{endpoint.port}" for endpoint in provider_endpoints)
        )
        for endpoint in provider_endpoints:
            if endpoint not in endpoints:
                endpoints.append(endpoint)
    if errors:
        print("doh_warnings=" + "; ".join(errors))
    return endpoints


def is_fake_ip(host: str) -> bool:
    try:
        ip = ipaddress.ip_address(host)
    except ValueError:
        return False
    return ip in ipaddress.ip_network("198.18.0.0/15")


def server_hostname(server: str, sni: str | None) -> str | None:
    value = sni or server
    if not value:
        return None
    return value


def tls_context(config: dict[str, Any]) -> ssl.SSLContext:
    skip_verify = bool(config.get("skip_cert_verify"))
    if skip_verify:
        context = ssl._create_unverified_context()
    else:
        context = ssl.create_default_context()

    alpn_present = bool(config.get("alpn_present"))
    alpn = config.get("alpn") or []
    if not alpn_present:
        alpn = ["h2", "http/1.1"]
    if alpn:
        context.set_alpn_protocols([str(value) for value in alpn])
    return context


def connect_tls(
    connect_endpoint: Endpoint,
    tls_server_name: str,
    node: dict[str, Any],
    timeout: float,
) -> ssl.SSLSocket:
    config = node["config"]
    context = tls_context(config)
    sni = server_hostname(tls_server_name, config.get("sni"))
    last_error: BaseException | None = None
    for _family, sockaddr in resolved_stream_addresses(connect_endpoint):
        started = time.monotonic()
        raw_sock: socket.socket | None = None
        try:
            raw_sock = socket.create_connection(sockaddr, timeout=timeout)
            raw_sock.settimeout(timeout)
            tls_sock = context.wrap_socket(raw_sock, server_hostname=sni)
            tls_sock.settimeout(timeout)
            elapsed_ms = (time.monotonic() - started) * 1000
            peer = sockaddr[0] if isinstance(sockaddr, tuple) else str(sockaddr)
            print(
                f"  tls_connect_ok address={peer}:{connect_endpoint.port} "
                f"elapsed_ms={elapsed_ms:.1f} alpn={tls_sock.selected_alpn_protocol()!r}"
            )
            return tls_sock
        except BaseException as exc:
            last_error = exc
            elapsed_ms = (time.monotonic() - started) * 1000
            peer = sockaddr[0] if isinstance(sockaddr, tuple) else str(sockaddr)
            print(
                f"  tls_connect_error address={peer}:{connect_endpoint.port} "
                f"elapsed_ms={elapsed_ms:.1f} error={type(exc).__name__}: {exc}"
            )
            if raw_sock is not None:
                raw_sock.close()
    assert last_error is not None
    raise last_error


def probe_trojan_tcp(
    connect_endpoint: Endpoint,
    tls_server_name: str,
    node: dict[str, Any],
    target: Endpoint,
    http_host: str,
    timeout: float,
) -> ProbeResult:
    request = (
        f"GET / HTTP/1.1\r\n"
        f"Host: {http_host}\r\n"
        f"Connection: close\r\n"
        f"User-Agent: burrow-proxy-ladder\r\n\r\n"
    ).encode()
    try:
        with connect_tls(connect_endpoint, tls_server_name, node, timeout) as sock:
            sock.sendall(trojan_header(node["config"]["password"], TROJAN_CMD_CONNECT, target))
            sock.sendall(request)
            data = sock.recv(2048)
    except BaseException as exc:
        return ProbeResult(False, f"{type(exc).__name__}: {exc}")
    if not data:
        return ProbeResult(False, "proxy returned EOF before any HTTP response bytes")
    first_line = data.splitlines()[0].decode("utf-8", errors="replace")
    return ProbeResult(True, f"received {len(data)} bytes; first_line={first_line!r}")


def probe_trojan_udp_dns(
    connect_endpoint: Endpoint,
    tls_server_name: str,
    node: dict[str, Any],
    dns_server: Endpoint,
    dns_name: str,
    timeout: float,
) -> ProbeResult:
    if node["config"].get("udp") is False:
        return ProbeResult(False, "selected Trojan node has udp=false")
    query_id, query = build_dns_query(dns_name)
    try:
        with connect_tls(connect_endpoint, tls_server_name, node, timeout) as sock:
            sock.sendall(
                trojan_header(node["config"]["password"], TROJAN_CMD_UDP_ASSOCIATE, dns_server)
            )
            sock.sendall(trojan_udp_frame(dns_server, query))
            source = read_socks_addr(sock)
            payload_len = struct.unpack("!H", read_exact(sock, 2))[0]
            delimiter = read_exact(sock, 2)
            if delimiter != b"\r\n":
                return ProbeResult(False, f"invalid Trojan UDP delimiter: {delimiter!r}")
            payload = read_exact(sock, payload_len)
    except BaseException as exc:
        return ProbeResult(False, f"{type(exc).__name__}: {exc}")
    dns_summary = parse_dns_response(query_id, payload)
    return ProbeResult(
        True,
        f"received UDP frame from {source.host}:{source.port}; DNS response {dns_summary}",
    )


def print_node_summary(network_id: int, payload: dict[str, Any], node: dict[str, Any]) -> None:
    config = node.get("config") or {}
    print(f"network_id={network_id}")
    print(f"subscription={payload.get('name')!r}")
    print(f"selected_node={node.get('name')!r}")
    print(f"protocol={node.get('protocol')!r}")
    print(f"server={node.get('server')}:{node.get('port')}")
    print(f"sni={config.get('sni')!r}")
    print(f"skip_cert_verify={bool(config.get('skip_cert_verify'))}")
    print(f"alpn_present={bool(config.get('alpn_present'))} alpn={config.get('alpn') or []!r}")
    print(f"udp={config.get('udp', True)!r}")


def main() -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Directly probe the selected Burrow proxy subscription node without "
            "using the daemon, Network Extension, system proxy, routes, or DNS settings."
        )
    )
    parser.add_argument("--db", type=Path, default=Path(os.environ.get("BURROW_DB", DEFAULT_DB)))
    parser.add_argument("--network-id", type=int)
    parser.add_argument("--node-name", help="Probe this subscription node instead of the selected node.")
    parser.add_argument("--node-ordinal", type=int, help="Probe this subscription ordinal instead of the selected node.")
    parser.add_argument(
        "--server-address",
        help=(
            "Override the proxy connect address, as host or host:port. "
            "This bypasses system DNS for the proxy hostname while preserving SNI."
        ),
    )
    parser.add_argument(
        "--resolve-server",
        choices=["system", "doh-google", "doh-cloudflare", "doh-all"],
        default="system",
        help="How to resolve the proxy server when --server-address is not provided.",
    )
    parser.add_argument(
        "--all-server-addresses",
        action="store_true",
        help="Probe every resolved proxy server address and pass if any address works.",
    )
    parser.add_argument("--target", default=DEFAULT_TARGET, help="TCP HTTP target, host:port")
    parser.add_argument("--http-host", default=DEFAULT_HTTP_HOST)
    parser.add_argument("--dns-server", default=DEFAULT_DNS_SERVER, help="UDP DNS target, host:port")
    parser.add_argument("--dns-name", default=DEFAULT_DNS_NAME)
    parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
    parser.add_argument("--skip-udp", action="store_true")
    parser.add_argument("--require-udp", action="store_true")
    args = parser.parse_args()

    print("Burrow proxy ladder: direct provider/protocol probe")
    print("bypasses=daemon,NetworkExtension,system_proxy,routes,system_dns,browser")
    print(f"db={args.db}")

    try:
        if args.node_name is not None and args.node_ordinal is not None:
            raise ValueError("--node-name and --node-ordinal are mutually exclusive")
        network_id, payload = load_payload(args.db, args.network_id)
        node = choose_node(payload, args.node_name, args.node_ordinal)
        print_node_summary(network_id, payload, node)
        if node.get("protocol") != "trojan":
            print(f"unsupported protocol for this direct probe: {node.get('protocol')!r}")
            return 2
        server = Endpoint(str(node["server"]), int(node["port"]))
        if args.server_address:
            connect_endpoints = [parse_endpoint(args.server_address, default_port=server.port)]
        elif args.resolve_server == "system":
            system_addresses = resolved_stream_addresses(server)
            connect_endpoints = [
                Endpoint(str(sockaddr[0]), int(sockaddr[1])) for _family, sockaddr in system_addresses
            ]
            if not args.all_server_addresses and connect_endpoints:
                connect_endpoints = connect_endpoints[:1]
            if not connect_endpoints:
                connect_endpoints = [server]
        elif args.resolve_server == "doh-all":
            connect_endpoints = resolve_doh_all(server.host, server.port, args.timeout)
            if not args.all_server_addresses and connect_endpoints:
                connect_endpoints = connect_endpoints[:1]
        else:
            connect_endpoints = resolve_doh(server.host, server.port, args.resolve_server, args.timeout)
            print(
                f"{args.resolve_server}_answers="
                + ",".join(f"{endpoint.host}:{endpoint.port}" for endpoint in connect_endpoints)
            )
            if not args.all_server_addresses and connect_endpoints:
                connect_endpoints = connect_endpoints[:1]
        if not connect_endpoints:
            raise LookupError("proxy server resolution returned no usable A records")
        target = parse_endpoint(args.target)
        dns_server = parse_endpoint(args.dns_server)
    except BaseException as exc:
        print(f"setup_error={type(exc).__name__}: {exc}", file=sys.stderr)
        return 2

    print(
        "connect_endpoints="
        + ",".join(f"{endpoint.host}:{endpoint.port}" for endpoint in connect_endpoints)
    )
    if any(endpoint != server for endpoint in connect_endpoints):
        print(f"tls_server_name_source={server.host!r} sni_preserved=True")
    if any(is_fake_ip(endpoint.host) for endpoint in connect_endpoints):
        print(
            "warning=proxy hostname resolved into 198.18.0.0/15; "
            "rerun with --resolve-server doh-all or --server-address REAL_IP to avoid fake-IP DNS"
        )

    any_tcp_ok = False
    any_required_udp_ok = False
    for idx, connect_endpoint in enumerate(connect_endpoints, start=1):
        if len(connect_endpoints) > 1:
            print(f"\n=== proxy server candidate {idx}/{len(connect_endpoints)} ===")
            print(f"connect_endpoint={connect_endpoint.host}:{connect_endpoint.port}")

        print("\n[1/2] Trojan TCP CONNECT probe")
        print(f"target={target.host}:{target.port} http_host={args.http_host!r}")
        tcp_result = probe_trojan_tcp(
            connect_endpoint, server.host, node, target, args.http_host, args.timeout
        )
        print(("ok: " if tcp_result.ok else "fail: ") + tcp_result.detail)
        any_tcp_ok = any_tcp_ok or tcp_result.ok

        udp_result = ProbeResult(True, "skipped")
        if not args.skip_udp:
            print("\n[2/2] Trojan UDP DNS probe")
            print(f"dns_server={dns_server.host}:{dns_server.port} dns_name={args.dns_name!r}")
            udp_result = probe_trojan_udp_dns(
                connect_endpoint, server.host, node, dns_server, args.dns_name, args.timeout
            )
            print(("ok: " if udp_result.ok else "fail: ") + udp_result.detail)
        else:
            print("\n[2/2] Trojan UDP DNS probe")
            print("skipped by --skip-udp")
        any_required_udp_ok = any_required_udp_ok or udp_result.ok

        if tcp_result.ok and (not args.require_udp or udp_result.ok):
            break

    if not any_tcp_ok:
        return 1
    if args.require_udp and not any_required_udp_ok:
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())