"""SEC-17: the egress gate judges the address a guarded download actually dials.

Two loopback addresses stand in for the two sides: 127.0.0.2 plays a public host, 127.0.0.1 our own
network. The gate under test refuses exactly 127.0.0.1, so every assertion is about the judgement and
the wiring, with real sockets, a real yt-dlp and a real ffmpeg — no mocked transport.
"""

import contextlib
import http.server
import ipaddress
import socket
import socketserver
import struct
import threading
import time
from pathlib import Path

import pytest
import yt_dlp
from yt_dlp.utils import download_range_func

from app.downloads import adapters, egress, errors, formats

PUBLIC = "127.0.0.2"
INTERNAL = "127.0.0.1"


def _forbid_internal(ip: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
    return str(ip) == INTERNAL


@pytest.fixture
def gate():
    g = egress.start_gate(forbidden=_forbid_internal)
    yield g
    g.shutdown()
    g.server_close()


class _Listener:
    """A TCP listener that only counts who connected — the internal service under attack."""

    def __init__(self, host: str = INTERNAL, port: int = 0):
        self.hits = 0
        self.sock = socket.socket()
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        self.sock.bind((host, port))
        self.sock.listen()
        self.port = self.sock.getsockname()[1]
        threading.Thread(target=self._serve, daemon=True).start()

    def _serve(self):
        while True:
            try:
                conn, _ = self.sock.accept()
            except OSError:
                return
            self.hits += 1
            with contextlib.suppress(OSError):
                conn.sendall(b"HTTP/1.1 200 OK\r\nContent-Length: 2\r\n\r\nok")
            conn.close()

    def close(self):
        self.sock.close()


def _public_server(routes: dict, port: int = 0, directory: Path | None = None):
    """An HTTP server on the 'public' address: `routes` maps a path to (status, headers); anything
    else is served from `directory`."""

    class H(http.server.SimpleHTTPRequestHandler):
        def __init__(self, *a, **k):
            super().__init__(*a, directory=str(directory or "."), **k)

        def log_message(self, *a):
            pass

        def do_GET(self):
            if self.path in routes:
                status, headers, body = routes[self.path]
                self.send_response(status)
                for k, v in headers.items():
                    self.send_header(k, v)
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                return
            super().do_GET()

    srv = socketserver.ThreadingTCPServer((PUBLIC, port), H)
    srv.daemon_threads = True
    threading.Thread(target=srv.serve_forever, daemon=True).start()
    return srv


def _guarded_opts(tmp_path: Path, gate_url: str, **extra) -> dict:
    opts = formats.build_ydl_opts({}, str(tmp_path / "%(id)s.%(ext)s"), lambda d: None, egress=True)
    opts.update(retries=0, fragment_retries=0, extractor_retries=0, socket_timeout=5)
    opts["proxy"] = gate_url  # the test's gate (it lets 127.0.0.2 through), not the worker's
    opts.update(extra)
    return opts


def _download_error(opts: dict, url: str) -> str:
    with pytest.raises(yt_dlp.utils.DownloadError) as exc, egress.open_ydl(opts) as ydl:
        ydl.extract_info(url, download=True)
    return str(exc.value)


# --- the judgement ---------------------------------------------------------------------------------


def test_a_name_with_any_forbidden_answer_is_refused():
    def resolve(answers):
        return lambda host, port, *a: [(2, 1, 6, "", (ip, port)) for ip in answers]

    assert egress.vetted_address(
        "x", 80, resolve=resolve([PUBLIC]), forbidden=_forbid_internal
    ) == (
        PUBLIC,
        80,
    )
    assert (
        egress.vetted_address(
            "x", 80, resolve=resolve([PUBLIC, INTERNAL]), forbidden=_forbid_internal
        )
        is None
    )

    def no_answer(*a):
        raise socket.gaierror("no such name")

    assert egress.vetted_address("x", 80, resolve=no_answer, forbidden=_forbid_internal) is None


def test_the_real_judgement_refuses_our_own_networks():
    # The worker's gate uses urlguard's judgement unchanged — the same one the enqueue check uses.
    for addr in ("127.0.0.1", "10.1.2.3", "172.18.0.5", "192.168.1.1", "169.254.169.254"):
        assert egress.vetted_address(addr, 80, forbidden=egress.urlguard.is_forbidden) is None, addr
    assert egress.vetted_address("93.184.215.14", 80, forbidden=egress.urlguard.is_forbidden)


def test_the_gate_dials_the_address_it_vetted_not_a_second_answer(gate):
    # DNS rebinding: a name that answers public for the check and private for the fetch. The gate
    # resolves once per connection and dials that answer.
    public = _public_server({"/": (200, {}, b"public")})
    internal = _Listener(INTERNAL, public.server_address[1])
    calls = []

    def rebinding(host, port, *a):
        calls.append(host)
        ip = PUBLIC if len(calls) == 1 else INTERNAL
        return [(2, 1, 6, "", (ip, port))]

    gate.resolve = rebinding
    try:
        with socket.create_connection(("127.0.0.1", gate.server_address[1])) as s:
            s.sendall(
                f"GET http://rebind.test:{public.server_address[1]}/ HTTP/1.1\r\n"
                f"Host: rebind.test\r\n\r\n".encode()
            )
            answer = b"".join(iter(lambda: s.recv(4096), b""))
        assert answer.endswith(b"public")
        assert calls == ["rebind.test"]
        assert internal.hits == 0
    finally:
        public.shutdown()
        internal.close()


def test_a_tunnel_into_the_network_is_refused_before_it_is_dialled(gate):
    internal = _Listener()
    try:
        with socket.create_connection(("127.0.0.1", gate.server_address[1])) as s:
            s.sendall(f"CONNECT {INTERNAL}:{internal.port} HTTP/1.1\r\n\r\n".encode())
            answer = s.recv(4096)
        assert answer.startswith(b"HTTP/1.1 403")
        assert internal.hits == 0
    finally:
        internal.close()


def test_a_client_that_hangs_up_mid_transfer_is_not_an_error(gate, monkeypatch):
    # yt-dlp closes a connection once it has read what it needed; seen live as a traceback the gate
    # printed for every such download.
    errors_seen = []
    monkeypatch.setattr(gate, "handle_error", lambda *a: errors_seen.append(a))
    public = _public_server({"/big": (200, {}, b"x" * 4_000_000)})
    try:
        s = socket.create_connection(("127.0.0.1", gate.server_address[1]))
        s.sendall(f"GET http://{PUBLIC}:{public.server_address[1]}/big HTTP/1.1\r\n\r\n".encode())
        s.recv(1024)
        s.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0))
        s.close()  # a reset, not a polite close
        time.sleep(0.5)
        assert errors_seen == []
    finally:
        public.shutdown()


# --- a real yt-dlp through it ----------------------------------------------------------------------


def test_a_redirect_into_the_network_reaches_it_unguarded_and_not_through_the_gate(tmp_path, gate):
    internal = _Listener()
    public = _public_server(
        {"/v.mp4": (302, {"Location": f"http://{INTERNAL}:{internal.port}/v.mp4"}, b"")}
    )
    url = f"http://{PUBLIC}:{public.server_address[1]}/v.mp4"
    try:
        # The hole, as it was: yt-dlp follows the 3xx into our network.
        bare = {"outtmpl": str(tmp_path / "bare" / "%(id)s.%(ext)s"), "quiet": True, "retries": 0}
        with contextlib.suppress(yt_dlp.utils.DownloadError), yt_dlp.YoutubeDL(bare) as ydl:
            ydl.extract_info(url, download=True)
        assert internal.hits >= 1
        before = internal.hits

        _download_error(_guarded_opts(tmp_path, gate.url), url)
        assert internal.hits == before
    finally:
        public.shutdown()
        internal.close()


def test_a_live_and_a_dead_internal_port_leave_the_same_text(tmp_path, gate):
    live = _Listener()
    dead = _Listener()
    dead_port = dead.port
    dead.close()
    public = _public_server(
        {
            "/live.mp4": (302, {"Location": f"http://{INTERNAL}:{live.port}/x"}, b""),
            "/dead.mp4": (302, {"Location": f"http://{INTERNAL}:{dead_port}/x"}, b""),
        }
    )
    base = f"http://{PUBLIC}:{public.server_address[1]}"
    try:
        stored = []
        for name in ("live", "dead"):
            msg = _download_error(_guarded_opts(tmp_path / name, gate.url), f"{base}/{name}.mp4")
            stored.append(errors.user_error_text("url", msg, errors.classify(msg)))
        assert stored[0] == stored[1]
        assert live.hits == 0
    finally:
        public.shutdown()
        live.close()


def test_a_format_fetched_past_the_gate_is_never_chosen(tmp_path):
    # ffmpeg speaks rtmp, rtsp, tcp… directly, where no HTTP proxy is asked. A guarded download keeps
    # only the formats fetched over http(s).
    opts = _guarded_opts(tmp_path, "http://127.0.0.1:9")
    info = {
        "id": "x",
        "title": "x",
        "extractor": "generic",
        "extractor_key": "Generic",
        "webpage_url": f"http://{PUBLIC}/page",
        "formats": [
            {"format_id": "rtmp", "url": f"rtmp://{INTERNAL}/live", "ext": "flv", "height": 1080},
            {"format_id": "web", "url": f"http://{PUBLIC}/a.mp4", "ext": "mp4", "height": 360},
        ],
    }
    with egress.open_ydl(opts) as ydl:
        chosen = ydl.process_ie_result(dict(info), download=False)
    assert chosen["format_id"] == "web"

    only_rtmp = dict(info, formats=[info["formats"][0]])
    with (
        pytest.raises(yt_dlp.utils.YoutubeDLError, match="Requested format is not available"),
        egress.open_ydl(opts) as ydl,
    ):
        ydl.process_ie_result(only_rtmp, download=False)


def test_a_format_yt_dlp_fetches_itself_is_kept():
    # f4m, ism and the DASH generator fetch through yt-dlp's own urlopen; a bare m3u8 goes to the
    # native HLS fetcher under hls_prefer_native (set for every guarded download). All ask the gate.
    kept = [
        {"format_id": p, "url": f"https://{PUBLIC}/{p}", "protocol": p}
        for p in ("f4m", "ism", "m3u8", "m3u8_native", "http_dash_segments_generator", "https")
    ]
    raw = {"format_id": "rtmp", "url": f"rtmp://{PUBLIC}/live", "protocol": "rtmp"}
    _, info = egress.HttpOnlyFormats().run({"formats": [*kept, raw]})
    assert [f["format_id"] for f in info["formats"]] == [f["format_id"] for f in kept]


def test_a_single_url_result_past_the_gate_loses_its_url():
    _, info = egress.HttpOnlyFormats().run({"url": f"rtmp://{INTERNAL}/live"})
    assert info["url"] is None
    _, info = egress.HttpOnlyFormats().run({"url": f"https://{PUBLIC}/a.mp4"})
    assert info["url"] == f"https://{PUBLIC}/a.mp4"


def test_the_format_filter_says_no_phase_of_its_own(tmp_path):
    # It is a post-processor only to run before the format choice; the job's postprocessor hook
    # must not see it (it painted every guarded job "processing" before a byte was fetched).
    calls: list[dict] = []
    opts = _guarded_opts(tmp_path, "http://127.0.0.1:9", postprocessor_hooks=[calls.append])
    with egress.open_ydl(opts) as ydl:
        ydl.run_all_pps("pre_process", {"id": "x", "formats": []})
    assert calls == []


def test_an_allowed_tunnel_carries_bytes_both_ways(gate):
    # The CONNECT path a site served over https takes: a 200, then the bytes the client sent after
    # the head (in the same packet here) reach the far end, and its answer comes back.
    far = socket.socket()
    far.bind((PUBLIC, 0))
    far.listen()

    def echo():
        conn, _ = far.accept()
        with conn:
            conn.sendall(b"PONG:" + conn.recv(64))

    threading.Thread(target=echo, daemon=True).start()
    try:
        with socket.create_connection(("127.0.0.1", gate.server_address[1]), timeout=5) as s:
            s.sendall(f"CONNECT {PUBLIC}:{far.getsockname()[1]} HTTP/1.1\r\n\r\nPING".encode())
            got = b""
            while b"PONG:PING" not in got:
                chunk = s.recv(4096)
                if not chunk:
                    break
                got += chunk
        assert got.startswith(b"HTTP/1.1 200")
        assert got.endswith(b"PONG:PING")
    finally:
        far.close()


def test_the_workers_own_gate_refuses_our_network():
    # Every other gate here is built with a test rule; this is the one `proxy_url()` starts.
    internal = _Listener()
    try:
        port = int(egress.proxy_url().rsplit(":", 1)[1])
        with socket.create_connection(("127.0.0.1", port), timeout=5) as s:
            s.sendall(f"CONNECT {INTERNAL}:{internal.port} HTTP/1.1\r\n\r\n".encode())
            assert s.recv(4096).startswith(b"HTTP/1.1 403")
        assert internal.hits == 0
    finally:
        internal.close()


def _clip_hits(tmp_path: Path, segment: str, proxy: str) -> int:
    """How many connections the internal listener takes while a clip (`download_ranges`, which
    yt-dlp hands to ffmpeg even for an HLS format) is cut from a playlist whose one segment is
    `segment` — formatted with the listener's address."""
    internal = _Listener()
    playlist = (
        "#EXTM3U\n#EXT-X-VERSION:3\n#EXT-X-TARGETDURATION:2\n#EXTINF:2.0,\n"
        + segment.format(host=INTERNAL, port=internal.port)
        + "\n#EXT-X-ENDLIST\n"
    ).encode()
    public = _public_server(
        {"/p.m3u8": (200, {"Content-Type": "application/vnd.apple.mpegurl"}, playlist)}
    )
    try:
        opts = _guarded_opts(tmp_path, proxy, download_ranges=download_range_func([], [(0, 1)]))
        with contextlib.suppress(yt_dlp.utils.DownloadError), egress.open_ydl(opts) as ydl:
            ydl.extract_info(f"http://{PUBLIC}:{public.server_address[1]}/p.m3u8", download=True)
        return internal.hits
    finally:
        public.shutdown()
        internal.close()


def test_a_clip_cut_by_ffmpeg_reaches_a_segment_only_through_the_gate(tmp_path, gate):
    # Without the gate, ffmpeg fetches a hostile playlist's segment from inside our network — the
    # control that shows this path is real. Through the gate it is never dialled.
    segment = "http://{host}:{port}/s.ts"
    assert _clip_hits(tmp_path / "bare", segment, proxy="") >= 1
    assert _clip_hits(tmp_path / "gated", segment, proxy=gate.url) == 0


def test_ffmpeg_itself_refuses_a_raw_tcp_segment(tmp_path):
    # A boundary, not our code: the HTTP proxy cannot see a tcp:// segment, and measured, ffmpeg's
    # HLS reader does not open one from an http playlist. This fails if an ffmpeg upgrade starts to.
    assert _clip_hits(tmp_path, "tcp://{host}:{port}/", proxy="") == 0


# --- the wiring ------------------------------------------------------------------------------------


def test_only_youtube_downloads_skip_the_gate():
    assert adapters.SourceAdapter.egress_guarded is True
    assert adapters.resolve_adapter("url", "https://www.facebook.com/x").egress_guarded is True
    assert adapters.resolve_adapter("url", "https://vimeo.com/1").egress_guarded is True
    assert adapters.resolve_adapter("youtube", "dQw4w9WgXcQ").egress_guarded is False


def test_a_guarded_download_goes_through_the_gate_whatever_its_adapter_adds():
    opts = formats.build_ydl_opts(
        {}, "/tmp/x", lambda d: None, extra={"proxy": "", "hls_prefer_native": False}, egress=True
    )
    assert opts["proxy"] == egress.proxy_url()
    assert opts["hls_prefer_native"] is True
    assert opts[egress.GUARDED] is True
    assert "proxy" not in formats.build_ydl_opts({}, "/tmp/x", lambda d: None, egress=False)


def test_every_yt_dlp_this_app_opens_goes_through_open_ydl():
    app = Path(egress.__file__).resolve().parents[1]
    offenders = [
        f"{p.relative_to(app)}:{n}"
        for p in sorted(app.rglob("*.py"))
        if p.name != "egress.py"
        for n, line in enumerate(p.read_text(encoding="utf-8").splitlines(), 1)
        if "YoutubeDL(" in line
    ]
    assert offenders == []