GEWU/tests/adapters/test_url_fetch.py

from __future__ import annotations

import socket
import time
from collections import deque
from typing import Any

import pytest

import gewu.adapters.url_fetch as url_fetch
from gewu.adapters.url_fetch import (
    MAX_CAPTURE_BYTES,
    PublicUrlFetcher,
    UrlCaptureError,
)


class FakeResponse:
    def __init__(self, status: int, headers: dict[str, str], body: bytes = b"") -> None:
        self.status = status
        self.headers = headers
        self.body = body
        self.read_calls = 0

    def getheader(self, name: str) -> str | None:
        return self.headers.get(name)

    def read(self, amount: int) -> bytes:
        self.read_calls += 1
        block, self.body = self.body[:amount], self.body[amount:]
        return block


class FakeConnection:
    def __init__(self, response: FakeResponse) -> None:
        self.response = response
        self.request_data: tuple[str, str, dict[str, str]] | None = None
        self.closed = False

    def request(self, method: str, path: str, headers: dict[str, str]) -> None:
        self.request_data = (method, path, headers)

    def getresponse(self) -> FakeResponse:
        return self.response

    def close(self) -> None:
        self.closed = True


def public_resolver(host: str, port: int, **_kwargs: Any) -> list[tuple[Any, ...]]:
    return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", port))]


def test_capture_requires_exact_host_allowlist_and_pins_validated_dns_address() -> None:
    response = FakeResponse(
        200,
        {"Content-Type": "text/plain; charset=utf-8", "X-Source-Title": "Fixture title"},
        b"fixture source",
    )
    connections: list[FakeConnection] = []
    pins: list[tuple[str, str, int, str, int]] = []

    def factory(scheme: str, host: str, port: int, address: str, timeout: int) -> FakeConnection:
        pins.append((scheme, host, port, address, timeout))
        connection = FakeConnection(response)
        connections.append(connection)
        return connection

    fetcher = PublicUrlFetcher(
        {"research.example"}, resolver=public_resolver, connection_factory=factory
    )
    captured = fetcher.fetch("https://research.example/reference?q=1#ignored")

    assert captured.original_url == "https://research.example/reference?q=1#ignored"
    assert captured.final_url == "https://research.example/reference?q=1#ignored"
    assert captured.media_type == "text/plain"
    assert captured.title == "Fixture title"
    assert captured.content == b"fixture source"
    assert pins[0][1:4] == ("research.example", 443, "93.184.216.34")
    assert connections[0].request_data == (
        "GET",
        "/reference?q=1",
        {
            "Host": "research.example",
            "User-Agent": "GEWU-source-capture/1.0",
            "Accept": "application/pdf, text/html, text/plain, text/markdown",
            "Accept-Encoding": "identity",
            "Connection": "close",
        },
    )
    assert connections[0].closed


def test_private_dns_result_is_rejected_before_opening_a_connection() -> None:
    factory_called = False

    def resolver(host: str, port: int, **_kwargs: Any) -> list[tuple[Any, ...]]:
        return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("127.0.0.1", port))]

    def factory(*_args: Any) -> FakeConnection:
        nonlocal factory_called
        factory_called = True
        return FakeConnection(FakeResponse(200, {"Content-Type": "text/plain"}))

    fetcher = PublicUrlFetcher({"research.example"}, resolver=resolver, connection_factory=factory)
    with pytest.raises(UrlCaptureError) as error:
        fetcher.fetch("http://research.example/source")
    assert error.value.code == "FETCH_BLOCKED"
    assert not factory_called


def test_dns_resolution_has_a_bounded_wait(monkeypatch: pytest.MonkeyPatch) -> None:
    def slow_resolver(host: str, port: int, **_kwargs: Any) -> list[tuple[Any, ...]]:
        time.sleep(0.03)
        return public_resolver(host, port)

    monkeypatch.setattr(url_fetch, "DNS_TIMEOUT_SECONDS", 0.001)
    fetcher = PublicUrlFetcher({"research.example"}, resolver=slow_resolver)
    with pytest.raises(UrlCaptureError) as error:
        fetcher.fetch("https://research.example/source")
    assert error.value.code == "FETCH_TIMEOUT"


def test_redirect_is_revalidated_against_allowlist_before_next_connection() -> None:
    first = FakeResponse(302, {"Location": "http://127.0.0.1/metadata"})
    connection_count = 0

    def factory(*_args: Any) -> FakeConnection:
        nonlocal connection_count
        connection_count += 1
        return FakeConnection(first)

    fetcher = PublicUrlFetcher(
        {"research.example"}, resolver=public_resolver, connection_factory=factory
    )
    with pytest.raises(UrlCaptureError) as error:
        fetcher.fetch("http://research.example/start")
    assert error.value.code == "FETCH_BLOCKED"
    assert connection_count == 1


def test_capture_rejects_unsupported_types_encoding_and_oversize() -> None:
    responses = deque(
        [
            FakeResponse(200, {"Content-Type": "application/zip"}),
            FakeResponse(
                200,
                {"Content-Type": "text/plain", "Content-Encoding": "gzip"},
            ),
            FakeResponse(
                200,
                {
                    "Content-Type": "text/plain",
                    "Content-Length": str(MAX_CAPTURE_BYTES + 1),
                },
            ),
        ]
    )

    def factory(*_args: Any) -> FakeConnection:
        return FakeConnection(responses.popleft())

    fetcher = PublicUrlFetcher(
        {"research.example"}, resolver=public_resolver, connection_factory=factory
    )
    for expected_code in (
        "UNSUPPORTED_MEDIA_TYPE",
        "UNSUPPORTED_ENCODING",
        "CAPTURE_TOO_LARGE",
    ):
        with pytest.raises(UrlCaptureError) as error:
            fetcher.fetch("http://research.example/source")
        assert error.value.code == expected_code


@pytest.mark.parametrize(
    "url",
    [
        "file:///etc/passwd",
        "http://user:password@research.example/source",
        "http://research.example:8080/source",
        "http://research.example/line\nInjected: header",
    ],
)
def test_capture_rejects_unsafe_url_forms(url: str) -> None:
    fetcher = PublicUrlFetcher({"research.example"}, resolver=public_resolver)
    with pytest.raises(UrlCaptureError) as error:
        fetcher.fetch(url)
    assert error.value.code == "URL_INVALID"