"""Unit tests for the spens DNS forwarder (no network dependencies beyond
loopback sockets; no mitmproxy/docker required)."""

from __future__ import annotations

import json
import socket
import struct
import threading

from spens.data.dns_forwarder import (
    DnsForwarder,
    DnsQueryLogger,
    _read_tcp_message,
    build_refused,
    build_servfail,
    load_domain_rules,
    log_entry,
    name_allowed,
    parse_questions,
    questions_allowed,
    upstream_from_resolv_conf,
    write_ready_marker,
)


def _query(name: str = "example.com", qtype: int = 1, qdcount: int = 1) -> bytes:
    """Build a DNS query message with ``qdcount`` questions (the last is
    ``name``/``qtype``; preceding ones are filler)."""
    questions = b""
    for i in range(qdcount - 1):
        qname = b"".join(bytes([len(l)]) + l.encode() for l in f"filler{i}.test".split(".")) + b"\x00"
        questions += qname + struct.pack("!HH", 1, 1)
    qname = b"".join(bytes([len(l)]) + l.encode() for l in name.split(".")) + b"\x00"
    questions += qname + struct.pack("!HH", qtype, 1)
    header = struct.pack("!HHHHHH", 0x1234, 0x0100, qdcount, 0, 0, 0)
    return header + questions


# -- pure parsing/formatting -------------------------------------------------


def test_parse_questions_single() -> None:
    assert parse_questions(_query()) == [("example.com", 1)]


def test_parse_questions_multiple() -> None:
    data = _query("api.anthropic.com", qtype=28, qdcount=2)
    assert parse_questions(data) == [
        ("filler0.test", 1),
        ("api.anthropic.com", 28),
    ]


def test_parse_questions_unknown_type_preserved() -> None:
    assert parse_questions(_query("example.com", qtype=999)) == [("example.com", 999)]


def test_parse_questions_truncated_is_tolerated() -> None:
    data = _query("example.com")
    # cut inside the question section: must not raise
    assert parse_questions(data[:-3]) == []


def test_parse_questions_garbage_is_tolerated() -> None:
    assert parse_questions(b"\x00" * 30) == []
    assert parse_questions(b"") == []


def test_build_servfail() -> None:
    query = _query()
    resp = build_servfail(query)
    assert struct.unpack_from("!H", resp, 0)[0] == 0x1234  # id preserved
    flags = struct.unpack_from("!H", resp, 2)[0]
    assert flags & 0x8000  # QR: response
    assert flags & 0x000F == 0x0002  # RCODE = SERVFAIL
    # question section is echoed back
    assert resp[12:] == query[12:]


def test_build_servfail_short_query_passthrough() -> None:
    assert build_servfail(b"ab") == b"ab"


def test_build_refused_sets_qr_and_rcode() -> None:
    query = _query("evil.example")
    resp = build_refused(query)
    assert struct.unpack_from("!H", resp, 0)[0] == 0x1234  # id preserved
    flags = struct.unpack_from("!H", resp, 2)[0]
    assert flags & 0x8000  # QR: response
    assert flags & 0x000F == 0x0005  # RCODE = REFUSED
    assert resp[12:] == query[12:]  # question echoed back


def test_build_refused_short_query_passthrough() -> None:
    assert build_refused(b"ab") == b"ab"


# -- domain allowlist policy -------------------------------------------------

RULES = [
    {"pattern": "*api.anthropic.com*", "allow": ["*"]},
    {"pattern": "pypi.org"},
]
ALWAYS = ["spens-interceptor", "localhost"]


def test_name_allowed_matches_domain_and_subdomains() -> None:
    assert name_allowed("api.anthropic.com", RULES, ALWAYS)
    assert name_allowed("eu.api.anthropic.com", RULES, ALWAYS)
    assert name_allowed("pypi.org", RULES, ALWAYS)
    assert name_allowed("files.pypi.org", RULES, ALWAYS)
    # trailing dot (FQDN root) is tolerated
    assert name_allowed("pypi.org.", RULES, ALWAYS)


def test_name_allowed_blocks_exfil_lookalikes() -> None:
    # the classic recursive-DNS exfil channel: data smuggled in the labels of
    # an attacker-controlled domain must NOT resolve
    assert not name_allowed("secret-payload.attacker.example", RULES, ALWAYS)
    # allowed domain as a *prefix* of an attacker domain (label-boundary check)
    assert not name_allowed("api.anthropic.com.attacker.example", RULES, ALWAYS)
    assert not name_allowed("data.pypi.org.evil.test", RULES, ALWAYS)
    # substring / lookalike, not a real suffix
    assert not name_allowed("evil-anthropic.com", RULES, ALWAYS)
    # a bare parent of the allowed pattern is not itself allowed
    assert not name_allowed("anthropic.com", RULES, ALWAYS)
    # reverse-DNS is not needed by the agent and would be an exfil vector too
    assert not name_allowed("1.2.3.4.in-addr.arpa", RULES, ALWAYS)


def test_name_allowed_always_allow_internal_names() -> None:
    # the interceptor alias must always resolve or the agent cannot reach the
    # proxy at all, even under a strict allowlist
    assert name_allowed("spens-interceptor", RULES, ALWAYS)
    assert name_allowed("localhost", RULES, ALWAYS)


def test_name_allowed_fail_open_without_rules() -> None:
    # no domain policy configured -> DNS is not filtered (mirrors HTTP egress)
    assert name_allowed("anything.example", [], ALWAYS)


def test_name_allowed_lone_wildcard_allows_all() -> None:
    assert name_allowed("whatever.test", [{"pattern": "*"}], [])


def test_questions_allowed_requires_every_question() -> None:
    good = [("api.anthropic.com", 1)]
    mixed = [("api.anthropic.com", 1), ("evil.example", 1)]
    assert questions_allowed(good, RULES, ALWAYS)
    assert not questions_allowed(mixed, RULES, ALWAYS)  # fail closed on any bad
    # empty question section is refused when a policy is active, allowed w/o
    assert not questions_allowed([], RULES, ALWAYS)
    assert questions_allowed([], [], ALWAYS)


def test_load_domain_rules(tmp_path) -> None:
    cfg = tmp_path / "cfg.json"
    cfg.write_text(
        json.dumps({"domain_rules": [{"pattern": "pypi.org"}, "bogus", {}]}),
        encoding="utf-8",
    )
    rules = load_domain_rules(str(cfg))
    assert rules == [{"pattern": "pypi.org"}, {}]


def test_load_domain_rules_missing_or_malformed(tmp_path) -> None:
    assert load_domain_rules(str(tmp_path / "nope.json")) == []
    bad = tmp_path / "bad.json"
    bad.write_text("{not json", encoding="utf-8")
    assert load_domain_rules(str(bad)) == []
    nolist = tmp_path / "nolist.json"
    nolist.write_text(json.dumps({"domain_rules": "x"}), encoding="utf-8")
    assert load_domain_rules(str(nolist)) == []


def test_log_entry_format() -> None:
    entry = log_entry([("example.com", 1), ("foo.test", 999)], "172.18.0.2", "udp")
    assert entry["type"] == "dns_query"
    assert entry["client"] == "172.18.0.2"
    assert entry["proto"] == "udp"
    assert entry["questions"] == [
        {"name": "example.com", "type": "A"},
        {"name": "foo.test", "type": "TYPE999"},
    ]
    json.dumps(entry)  # must be JSON-serializable


def test_upstream_from_resolv_conf(tmp_path) -> None:
    resolv = tmp_path / "resolv.conf"
    resolv.write_text(
        "# comment\nnameserver 10.0.0.2\nnameserver 10.0.0.3\n",
        encoding="utf-8",
    )
    assert upstream_from_resolv_conf(str(resolv)) == ("10.0.0.2", 53)


def test_upstream_from_resolv_conf_missing_file(tmp_path) -> None:
    assert upstream_from_resolv_conf(str(tmp_path / "nope")) == ("127.0.0.11", 53)


def test_write_ready_marker(tmp_path) -> None:
    marker = write_ready_marker(tmp_path)
    assert marker.name == "spens_dns_ready"
    assert marker.exists()
    assert marker.read_text(encoding="utf-8").strip()  # timestamp content


def test_dns_query_logger_appends_jsonl(tmp_path) -> None:
    log = DnsQueryLogger(tmp_path / "dns_log.jsonl")
    log.log(log_entry([("a.test", 1)], "172.18.0.2", "udp"))
    log.log(log_entry([("b.test", 28)], "172.18.0.2", "tcp"))
    lines = (tmp_path / "dns_log.jsonl").read_text(encoding="utf-8").splitlines()
    assert len(lines) == 2
    assert json.loads(lines[0])["questions"] == [{"name": "a.test", "type": "A"}]
    assert json.loads(lines[1])["questions"] == [{"name": "b.test", "type": "AAAA"}]


def test_dns_query_logger_swallows_write_errors(tmp_path) -> None:
    # path is a directory -> open() raises OSError -> must not propagate
    log = DnsQueryLogger(tmp_path)
    log.log(log_entry([("a.test", 1)], "c", "udp"))


# -- forwarding (loopback sockets) -------------------------------------------


def _echo_upstream_udp(sock: socket.socket) -> None:
    data, addr = sock.recvfrom(65535)
    flags = struct.unpack_from("!H", data, 2)[0] | 0x8000
    sock.sendto(data[:2] + struct.pack("!H", flags) + data[4:], addr)


def test_forwarder_udp_roundtrip_and_logging(tmp_path) -> None:
    upstream = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    upstream.bind(("127.0.0.1", 0))
    upstream_port = upstream.getsockname()[1]
    threading.Thread(target=_echo_upstream_udp, args=(upstream,), daemon=True).start()

    logger = DnsQueryLogger(tmp_path / "dns_log.jsonl")
    forwarder = DnsForwarder(("127.0.0.1", upstream_port), logger, timeout=5)

    server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server.bind(("127.0.0.1", 0))
    server_port = server.getsockname()[1]
    threading.Thread(target=forwarder.serve_udp, args=(server,), daemon=True).start()

    client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client.settimeout(5)
    client.sendto(_query("example.com"), ("127.0.0.1", server_port))
    resp, _ = client.recvfrom(65535)
    assert struct.unpack_from("!H", resp, 2)[0] & 0x8000  # a response came back

    # the query was logged before forwarding
    entry = json.loads(
        (tmp_path / "dns_log.jsonl").read_text(encoding="utf-8").splitlines()[0]
    )
    assert entry["questions"] == [{"name": "example.com", "type": "A"}]
    assert entry["proto"] == "udp"
    assert entry["client"] == "127.0.0.1"


def test_forwarder_udp_blocked_name_refused_and_not_forwarded(tmp_path) -> None:
    """A name outside the allowlist must be REFUSED without ever reaching the
    upstream resolver -- this is what closes the recursive-DNS exfil channel."""
    contacted = {"hit": False}
    upstream = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    upstream.bind(("127.0.0.1", 0))
    upstream_port = upstream.getsockname()[1]

    def upstream_server() -> None:
        while True:
            data, addr = upstream.recvfrom(65535)
            contacted["hit"] = True
            flags = struct.unpack_from("!H", data, 2)[0] | 0x8000
            upstream.sendto(data[:2] + struct.pack("!H", flags) + data[4:], addr)

    threading.Thread(target=upstream_server, daemon=True).start()

    logger = DnsQueryLogger(tmp_path / "dns_log.jsonl")
    forwarder = DnsForwarder(
        ("127.0.0.1", upstream_port),
        logger,
        timeout=5,
        domain_rules=[{"pattern": "pypi.org"}],
        always_allow=["spens-interceptor"],
    )

    server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server.bind(("127.0.0.1", 0))
    server_port = server.getsockname()[1]
    threading.Thread(target=forwarder.serve_udp, args=(server,), daemon=True).start()

    client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client.settimeout(5)
    client.sendto(_query("exfil-data.attacker.example"), ("127.0.0.1", server_port))
    resp, _ = client.recvfrom(65535)
    assert struct.unpack_from("!H", resp, 2)[0] & 0x000F == 0x0005  # REFUSED
    assert contacted["hit"] is False  # the blocked query never left the box

    entry = json.loads(
        (tmp_path / "dns_log.jsonl").read_text(encoding="utf-8").splitlines()[0]
    )
    assert entry["action"] == "blocked"
    assert entry["questions"] == [{"name": "exfil-data.attacker.example", "type": "A"}]


def test_forwarder_udp_allowed_name_is_forwarded(tmp_path) -> None:
    upstream = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    upstream.bind(("127.0.0.1", 0))
    upstream_port = upstream.getsockname()[1]
    threading.Thread(target=_echo_upstream_udp, args=(upstream,), daemon=True).start()

    logger = DnsQueryLogger(tmp_path / "dns_log.jsonl")
    forwarder = DnsForwarder(
        ("127.0.0.1", upstream_port),
        logger,
        timeout=5,
        domain_rules=[{"pattern": "pypi.org"}],
    )

    server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server.bind(("127.0.0.1", 0))
    server_port = server.getsockname()[1]
    threading.Thread(target=forwarder.serve_udp, args=(server,), daemon=True).start()

    client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client.settimeout(5)
    client.sendto(_query("files.pypi.org"), ("127.0.0.1", server_port))
    resp, _ = client.recvfrom(65535)
    assert struct.unpack_from("!H", resp, 2)[0] & 0x8000  # a response came back

    entry = json.loads(
        (tmp_path / "dns_log.jsonl").read_text(encoding="utf-8").splitlines()[0]
    )
    assert entry["action"] == "forward"


def test_forwarder_udp_returns_servfail_when_upstream_dead(tmp_path) -> None:
    # occupy then release a port so nothing is listening on it
    probe = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    probe.bind(("127.0.0.1", 0))
    dead_port = probe.getsockname()[1]
    probe.close()

    logger = DnsQueryLogger(tmp_path / "dns_log.jsonl")
    forwarder = DnsForwarder(("127.0.0.1", dead_port), logger, timeout=0.2)

    server = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    server.bind(("127.0.0.1", 0))
    server_port = server.getsockname()[1]
    threading.Thread(target=forwarder.serve_udp, args=(server,), daemon=True).start()

    query = _query("example.com")
    client = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    client.settimeout(5)
    client.sendto(query, ("127.0.0.1", server_port))
    resp, _ = client.recvfrom(65535)
    assert struct.unpack_from("!H", resp, 2)[0] & 0x000F == 0x0002  # SERVFAIL

    # the query was still logged (audit before forwarding)
    entry = json.loads(
        (tmp_path / "dns_log.jsonl").read_text(encoding="utf-8").splitlines()[0]
    )
    assert entry["questions"] == [{"name": "example.com", "type": "A"}]


def _free_tcp_port() -> int:
    probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    probe.bind(("127.0.0.1", 0))
    port = probe.getsockname()[1]
    probe.close()
    return port


def test_forwarder_tcp_roundtrip_and_logging(tmp_path) -> None:
    upstream_port = _free_tcp_port()

    def upstream_server() -> None:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as srv:
            srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
            srv.bind(("127.0.0.1", upstream_port))
            srv.listen(1)
            conn, _ = srv.accept()
            with conn:
                msg = _read_tcp_message(conn)
                assert msg is not None
                flags = struct.unpack_from("!H", msg, 2)[0] | 0x8000
                resp = msg[:2] + struct.pack("!H", flags) + msg[4:]
                conn.sendall(struct.pack("!H", len(resp)) + resp)

    threading.Thread(target=upstream_server, daemon=True).start()

    logger = DnsQueryLogger(tmp_path / "dns_log.jsonl")
    forwarder = DnsForwarder(("127.0.0.1", upstream_port), logger, timeout=5)

    server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server.bind(("127.0.0.1", 0))
    server.listen(1)
    server_port = server.getsockname()[1]
    threading.Thread(target=forwarder.serve_tcp, args=(server,), daemon=True).start()

    query = _query("example.com")
    client = socket.create_connection(("127.0.0.1", server_port), timeout=5)
    client.sendall(struct.pack("!H", len(query)) + query)
    header = client.recv(2)
    (length,) = struct.unpack("!H", header)
    resp = b""
    while len(resp) < length:
        resp += client.recv(length - len(resp))
    client.close()
    assert struct.unpack_from("!H", resp, 2)[0] & 0x8000

    entry = json.loads(
        (tmp_path / "dns_log.jsonl").read_text(encoding="utf-8").splitlines()[0]
    )
    assert entry["proto"] == "tcp"
    assert entry["questions"] == [{"name": "example.com", "type": "A"}]


def _free_tcp_port() -> int:
    probe = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    probe.bind(("127.0.0.1", 0))
    port = probe.getsockname()[1]
    probe.close()
    return port


def test_read_tcp_message_frame_and_eof() -> None:
    a, b = socket.socketpair()
    with a, b:
        a.sendall(struct.pack("!H", 4) + b"abcd")
        assert _read_tcp_message(b) == b"abcd"
        a.close()
        assert _read_tcp_message(b) is None


# -- DNS / HTTP policy parity ------------------------------------------------


def test_dns_and_http_enforce_the_same_domain_rule() -> None:
    """DNS resolution and HTTP egress share one matching rule.

    The forwarder used to carry its own copy of the suffix-matching logic;
    this pins that the two channels agree, so a name that cannot be resolved
    is exactly a host that cannot be reached.
    """
    from spens.data.mitmproxy_helpers import host_allowed

    rules = [{"pattern": "*api.anthropic.com*"}, {"pattern": "pypi.org"}]
    for host in [
        "api.anthropic.com",
        "eu.api.anthropic.com",
        "pypi.org",
        "files.pypi.org",
        "anthropic.com",
        "evil-anthropic.com",
        "api.anthropic.com.attacker.example",
        "secret-payload.attacker.example",
        "",
    ]:
        assert name_allowed(host, rules, []) == host_allowed(host, rules), host


def test_name_allowed_ignores_blank_always_allow_entries() -> None:
    """A blank entry must not strip to a bare wildcard and allow everything."""
    rules = [{"pattern": "pypi.org"}]
    assert not name_allowed("attacker.example", rules, ["", "  "])
