Spaces:
Sleeping
Sleeping
File size: 4,107 Bytes
936f0bf c5638b0 936f0bf c5638b0 936f0bf c5638b0 936f0bf c5638b0 936f0bf c5638b0 936f0bf c5638b0 | 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 | 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)"
)
|