from __future__ import annotations import io import ipaddress import socket from urllib.parse import urlparse import httpx import structlog from fastapi import UploadFile from app.core.config import settings from app.core.exceptions import FileTooLargeError, InvalidFileTypeError, InvalidParameterError logger = structlog.get_logger(__name__) _PDF_MAGIC = b"%PDF" _CHUNK_SIZE = 64 * 1024 _MAX_URL_REDIRECTS = 5 _ALLOWED_SCHEMES = frozenset({"http", "https"}) # SSRF guard: pre-collapsed private/loopback/link-local networks. _PRIVATE_RANGES: tuple[ipaddress.IPv4Network | ipaddress.IPv6Network, ...] = ( ipaddress.ip_network("10.0.0.0/8"), ipaddress.ip_network("172.16.0.0/12"), ipaddress.ip_network("192.168.0.0/16"), ipaddress.ip_network("127.0.0.0/8"), ipaddress.ip_network("169.254.0.0/16"), ipaddress.ip_network("::1/128"), ipaddress.ip_network("fc00::/7"), ) def _is_private_ip(addr: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool: return any(addr in net for net in _PRIVATE_RANGES) async def read_and_validate_pdf(upload: UploadFile) -> bytes: buffer = io.BytesIO() total = 0 while chunk := await upload.read(_CHUNK_SIZE): total += len(chunk) if total > settings.MAX_FILE_SIZE_BYTES: raise FileTooLargeError(settings.MAX_FILE_SIZE_MB) buffer.write(chunk) if total == 0: raise InvalidFileTypeError("Uploaded file is empty") pdf_bytes = buffer.getvalue() _assert_pdf_magic(pdf_bytes) logger.debug("pdf_upload_validated", size_bytes=total) return pdf_bytes async def fetch_pdf_from_url(url: str) -> bytes: _validate_url_scheme(url) _validate_url_host(url) async with httpx.AsyncClient( follow_redirects=True, max_redirects=_MAX_URL_REDIRECTS, timeout=httpx.Timeout(30.0, connect=10.0), ) as client: try: response = await client.get(url) response.raise_for_status() except httpx.HTTPStatusError as exc: raise InvalidFileTypeError( f"Failed to fetch PDF from URL: HTTP {exc.response.status_code}" ) from exc except httpx.RequestError as exc: raise InvalidFileTypeError(f"Failed to fetch PDF from URL: {exc}") from exc _validate_url_host(str(response.url)) content_length = response.headers.get("content-length") if content_length and int(content_length) > settings.MAX_FILE_SIZE_BYTES: raise FileTooLargeError(settings.MAX_FILE_SIZE_MB) pdf_bytes = b"" async for chunk in response.aiter_bytes(chunk_size=_CHUNK_SIZE): pdf_bytes += chunk if len(pdf_bytes) > settings.MAX_FILE_SIZE_BYTES: raise FileTooLargeError(settings.MAX_FILE_SIZE_MB) if len(pdf_bytes) == 0: raise InvalidFileTypeError("URL returned an empty response") _assert_pdf_magic(pdf_bytes) logger.debug("pdf_url_validated", url=url, size_bytes=len(pdf_bytes)) return pdf_bytes def _assert_pdf_magic(data: bytes) -> None: if not data.startswith(_PDF_MAGIC): raise InvalidFileTypeError( "File does not appear to be a valid PDF (missing PDF magic bytes)" ) def _validate_url_scheme(url: str) -> None: parsed = urlparse(url) if parsed.scheme.lower() not in _ALLOWED_SCHEMES: raise InvalidParameterError( f"URL scheme '{parsed.scheme}' is not allowed. Only http and https are permitted." ) def _validate_url_host(url: str) -> None: parsed = urlparse(url) hostname = parsed.hostname if not hostname: raise InvalidParameterError("URL has no valid hostname") try: resolved_ip = socket.gethostbyname(hostname) addr = ipaddress.ip_address(resolved_ip) except (socket.gaierror, ValueError): raise InvalidParameterError(f"Could not resolve hostname: {hostname}") if _is_private_ip(addr): raise InvalidParameterError( "Requests to private or loopback IP addresses are not permitted (SSRF protection)" )