126 lines
4.6 KiB
Python
126 lines
4.6 KiB
Python
import asyncio
|
|
import ipaddress
|
|
import socket
|
|
from dataclasses import dataclass
|
|
from time import perf_counter
|
|
from urllib.parse import urljoin, urlsplit
|
|
|
|
import aiohttp
|
|
from aiohttp.abc import AbstractResolver
|
|
|
|
|
|
class UnsafeTargetError(Exception):
|
|
pass
|
|
|
|
|
|
class CheckTimeoutError(Exception):
|
|
def __init__(self, latency_ms: float) -> None:
|
|
self.latency_ms = latency_ms
|
|
|
|
|
|
class CheckNetworkError(Exception):
|
|
def __init__(self, latency_ms: float) -> None:
|
|
self.latency_ms = latency_ms
|
|
|
|
|
|
class TooManyRedirectsError(Exception):
|
|
def __init__(self, latency_ms: float) -> None:
|
|
self.latency_ms = latency_ms
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class HttpOutcome:
|
|
status_code: int
|
|
latency_ms: float
|
|
final_url: str
|
|
|
|
|
|
class SafeResolver(AbstractResolver):
|
|
async def resolve(
|
|
self, host: str, port: int = 0, family: int = socket.AF_INET
|
|
) -> list[dict[str, object]]:
|
|
try:
|
|
infos = await asyncio.to_thread(
|
|
socket.getaddrinfo, host, port, family, socket.SOCK_STREAM
|
|
)
|
|
except socket.gaierror as exc:
|
|
raise OSError("DNS resolution failed") from exc
|
|
results: list[dict[str, object]] = []
|
|
seen: set[str] = set()
|
|
for resolved_family, _, proto, _, sockaddr in infos:
|
|
address = sockaddr[0]
|
|
ip = ipaddress.ip_address(address)
|
|
if not ip.is_global:
|
|
raise UnsafeTargetError("target resolves to a non-public address")
|
|
if address not in seen:
|
|
seen.add(address)
|
|
results.append({
|
|
"hostname": host, "host": address, "port": port,
|
|
"family": resolved_family, "proto": proto, "flags": 0,
|
|
})
|
|
if not results:
|
|
raise OSError("DNS returned no addresses")
|
|
return results
|
|
|
|
async def close(self) -> None:
|
|
return None
|
|
|
|
async def validate_url(self, url: str) -> None:
|
|
parsed = urlsplit(url)
|
|
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
|
|
raise UnsafeTargetError("only absolute HTTP(S) URLs are allowed")
|
|
if parsed.username is not None or parsed.password is not None:
|
|
raise UnsafeTargetError("URL userinfo is not allowed")
|
|
try:
|
|
port = parsed.port or (443 if parsed.scheme == "https" else 80)
|
|
except ValueError as exc:
|
|
raise UnsafeTargetError("invalid target port") from exc
|
|
await self.resolve(parsed.hostname, port, socket.AF_UNSPEC)
|
|
|
|
|
|
class EndpointChecker:
|
|
def __init__(self, timeout_seconds: float, max_redirects: int) -> None:
|
|
self.timeout_seconds = timeout_seconds
|
|
self.max_redirects = max_redirects
|
|
self.resolver = SafeResolver()
|
|
|
|
async def check(self, url: str) -> HttpOutcome:
|
|
started = perf_counter()
|
|
timeout = aiohttp.ClientTimeout(total=self.timeout_seconds)
|
|
connector = aiohttp.TCPConnector(
|
|
resolver=self.resolver, ttl_dns_cache=0, force_close=True
|
|
)
|
|
try:
|
|
async with aiohttp.ClientSession(
|
|
timeout=timeout, connector=connector, trust_env=False
|
|
) as session:
|
|
current = url
|
|
for hop in range(self.max_redirects + 1):
|
|
await self.resolver.validate_url(current)
|
|
async with session.get(current, allow_redirects=False) as response:
|
|
if response.status in {301, 302, 303, 307, 308}:
|
|
location = response.headers.get("Location")
|
|
if location is None:
|
|
return HttpOutcome(
|
|
response.status, self._elapsed(started), current
|
|
)
|
|
if hop == self.max_redirects:
|
|
raise TooManyRedirectsError(self._elapsed(started))
|
|
current = urljoin(current, location)
|
|
continue
|
|
return HttpOutcome(
|
|
response.status, self._elapsed(started), current
|
|
)
|
|
except UnsafeTargetError:
|
|
raise
|
|
except TooManyRedirectsError:
|
|
raise
|
|
except (asyncio.TimeoutError, TimeoutError) as exc:
|
|
raise CheckTimeoutError(self._elapsed(started)) from exc
|
|
except aiohttp.ClientError as exc:
|
|
raise CheckNetworkError(self._elapsed(started)) from exc
|
|
|
|
@staticmethod
|
|
def _elapsed(started: float) -> float:
|
|
return round((perf_counter() - started) * 1000, 3)
|