app/utils/validators.py

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)"
        )