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"