86 lines
3.7 KiB
Python
86 lines
3.7 KiB
Python
import asyncio
|
|
import json
|
|
import logging
|
|
import socket
|
|
|
|
import pytest
|
|
|
|
from app.checker import EndpointChecker, HopResponse, PinnedResolver
|
|
from app.config import Settings
|
|
from app.logging import JsonFormatter
|
|
from app.models import HealthState
|
|
from app.security import TargetPolicyError, ValidatedTarget, redact_url, validate_target
|
|
|
|
|
|
async def resolver_for(*addresses: str):
|
|
async def resolve(host: str, port: int) -> list[tuple[object, ...]]:
|
|
return [(socket.AF_INET6 if ":" in address else socket.AF_INET, 1, 6, "", (address, port)) for address in addresses]
|
|
return resolve
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize("address", ["127.0.0.1", "10.1.2.3", "169.254.169.254", "::1", "fc00::1", "0.0.0.0", "224.0.0.1"])
|
|
async def test_ip_classification_blocks_non_public_answers(address: str) -> None:
|
|
with pytest.raises(TargetPolicyError):
|
|
await validate_target("http://example.test/", await resolver_for(address))
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mixed_dns_answer_is_blocked_and_public_is_pinned() -> None:
|
|
with pytest.raises(TargetPolicyError):
|
|
await validate_target("https://example.test", await resolver_for("93.184.216.34", "127.0.0.1"))
|
|
target = await validate_target("https://example.test", await resolver_for("93.184.216.34"))
|
|
assert target.addresses == ("93.184.216.34",)
|
|
pinned = PinnedResolver(target.hostname, target.addresses)
|
|
assert (await pinned.resolve("example.test", 443))[0]["host"] == "93.184.216.34"
|
|
with pytest.raises(OSError):
|
|
await pinned.resolve("attacker.test", 443)
|
|
|
|
|
|
class RedirectFetcher:
|
|
def __init__(self) -> None:
|
|
self.calls: list[ValidatedTarget] = []
|
|
|
|
async def fetch(self, target: ValidatedTarget, settings: Settings, remaining: float) -> HopResponse:
|
|
self.calls.append(target)
|
|
return HopResponse(302, "http://internal.test/admin") if len(self.calls) == 1 else HopResponse(200, None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_redirect_hop_is_resolved_and_blocked_before_second_fetch() -> None:
|
|
seen: list[str] = []
|
|
async def resolver(host: str, port: int) -> list[tuple[object, ...]]:
|
|
seen.append(host)
|
|
address = "93.184.216.34" if host == "public.test" else "127.0.0.1"
|
|
return [(socket.AF_INET, 1, 6, "", (address, port))]
|
|
fetcher = RedirectFetcher()
|
|
result = await EndpointChecker(Settings(), fetcher=fetcher, resolver=resolver).check("http://public.test/start")
|
|
assert result.error_code == "blocked_target"
|
|
assert seen == ["public.test", "internal.test"]
|
|
assert len(fetcher.calls) == 1
|
|
|
|
|
|
class SlowFetcher:
|
|
async def fetch(self, target: ValidatedTarget, settings: Settings, remaining: float) -> HopResponse:
|
|
await asyncio.sleep(0.1)
|
|
return HopResponse(200, None)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_total_timeout_is_mapped_without_leaking_exception() -> None:
|
|
checker = EndpointChecker(Settings(total_timeout_seconds=0.01, connect_timeout_seconds=0.01, read_timeout_seconds=0.01), fetcher=SlowFetcher(), resolver=await resolver_for("93.184.216.34"))
|
|
result = await checker.check("https://example.test/?token=secret")
|
|
assert result.state == HealthState.down
|
|
assert result.error_code == "timeout"
|
|
assert result.error_message == "endpoint check timed out"
|
|
assert result.latency_ms >= 0
|
|
|
|
|
|
def test_query_redaction_and_json_log_output() -> None:
|
|
assert redact_url("https://user:pass@example.com:8443/p?q=secret#frag") == "https://example.com:8443/p"
|
|
record = logging.LogRecord("x", logging.INFO, "", 1, "done", (), None)
|
|
record.url = redact_url("https://example.com/path?api_key=hunter2")
|
|
payload = JsonFormatter().format(record)
|
|
assert json.loads(payload)["url"] == "https://example.com/path"
|
|
assert "hunter2" not in payload
|