Spaces:
Sleeping
Sleeping
File size: 5,528 Bytes
7d7841f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 | import asyncio
import time
from dataclasses import dataclass
from urllib.parse import urljoin, urlparse
import httpx
from app.crawler_registry import bot_user_agents
from app.pinned_dns import PinnedHTTPTransport, pinned_dns
from app.ssrf import resolve_public_address
DEFAULT_USER_AGENT = "MachineRead/1.0"
BROWSER_USER_AGENT = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
"(KHTML, like Gecko) Chrome/126.0.0.0 Safari/537.36"
)
BOT_USER_AGENTS: dict[str, str] = bot_user_agents()
@dataclass(frozen=True)
class FetchResult:
requested_url: str
final_url: str
status_code: int | None
headers: dict[str, str]
text: str
elapsed_ms: int
redirect_chain: list[str]
error: str | None = None
@property
def ok(self) -> bool:
# 3xx is treated as OK because the HTTP client follows redirects by
# default; arriving at a 3xx here means we hit the redirect cap
# (FetchLimit.max_redirects) rather than the redirect itself being
# successful. Callers that care about redirect-cap awareness should
# inspect ``status_code`` and ``redirect_chain`` separately rather
# than relying on ``ok`` alone.
return self.error is None and self.status_code is not None and 200 <= self.status_code < 400
@property
def blocked(self) -> bool:
return self.status_code in {401, 403, 429, 503}
def make_root_url(base_url: str, path: str) -> str:
return urljoin(base_url, path)
def _response_headers(response: httpx.Response) -> dict[str, str]:
return {key.lower(): value for key, value in response.headers.items()}
async def fetch_url(
url: str,
user_agent: str = DEFAULT_USER_AGENT,
accept: str = "text/html,application/xhtml+xml,application/xml;q=0.9,text/plain;q=0.8,*/*;q=0.7",
timeout: float = 15.0,
max_redirects: int = 5,
method: str = "GET",
) -> FetchResult:
"""Fetch a public URL while validating every redirect target.
``method`` is HTTP verb (``GET`` by default). ``HEAD`` is supported as a
lightweight probe for endpoints such as NLWeb's ``/ask``. All existing
callers default to GET so behaviour is unchanged.
"""
start = time.monotonic()
method_upper = (method or "GET").upper()
try:
current_url, current_ip, _ = await asyncio.to_thread(resolve_public_address, url)
except ValueError as exc:
return FetchResult(url, url, None, {}, "", 0, [], str(exc))
headers = {
"User-Agent": user_agent,
"Accept": accept,
}
redirect_chain: list[str] = []
try:
async with httpx.AsyncClient(
timeout=timeout,
follow_redirects=False,
transport=PinnedHTTPTransport(),
) as client:
for _ in range(max_redirects + 1):
pin_token = pinned_dns.set((urlparse(current_url).hostname or "", current_ip))
try:
response = await client.request(method_upper, current_url, headers=headers)
finally:
pinned_dns.reset(pin_token)
elapsed_ms = round((time.monotonic() - start) * 1000)
final_url = str(response.url)
if response.status_code not in {301, 302, 303, 307, 308}:
return FetchResult(
requested_url=url,
final_url=final_url,
status_code=response.status_code,
headers=_response_headers(response),
text=response.text,
elapsed_ms=elapsed_ms,
redirect_chain=redirect_chain,
)
location = response.headers.get("location")
if not location:
return FetchResult(
requested_url=url,
final_url=final_url,
status_code=response.status_code,
headers=_response_headers(response),
text=response.text,
elapsed_ms=elapsed_ms,
redirect_chain=redirect_chain,
)
next_url = urljoin(final_url, location)
try:
next_url, next_ip, _ = await asyncio.to_thread(resolve_public_address, next_url)
except ValueError as exc:
return FetchResult(
requested_url=url,
final_url=final_url,
status_code=response.status_code,
headers=_response_headers(response),
text="",
elapsed_ms=elapsed_ms,
redirect_chain=redirect_chain,
error=f"Unsafe redirect target: {exc}",
)
redirect_chain.append(next_url)
current_url, current_ip = next_url, next_ip
except httpx.RequestError as exc:
elapsed_ms = round((time.monotonic() - start) * 1000)
return FetchResult(url, current_url, None, {}, "", elapsed_ms, redirect_chain, str(exc))
elapsed_ms = round((time.monotonic() - start) * 1000)
return FetchResult(
requested_url=url,
final_url=current_url,
status_code=None,
headers={},
text="",
elapsed_ms=elapsed_ms,
redirect_chain=redirect_chain,
error="Too many redirects",
)
|