from __future__ import annotations import asyncio import hashlib import ipaddress import os import re import shutil import socket import ssl import stat import struct import tempfile import threading import zipfile from collections.abc import ( AsyncIterator, Awaitable, Callable, Iterable, Iterator, Mapping, Sequence, ) from contextlib import contextmanager, suppress from dataclasses import dataclass, field from io import BytesIO from pathlib import Path from typing import Any, NotRequired, Protocol, TypedDict from urllib.parse import urljoin, urlsplit import httpcore import httpx import pypdf from defusedxml import ElementTree as ET from pypdf import PdfReader _pypdf_apply_configuration: Callable[..., Any] | None = getattr(pypdf, "apply_configuration", None) Resolver = Callable[[str], Awaitable[Sequence[str]]] ActorContext = object SocketOption = ( tuple[int, int, int] | tuple[int, int, bytes | bytearray] | tuple[int, int, None, int] ) PDF_MIME = "application/pdf" DOCX_MIME = "application/vnd.openxmlformats-officedocument.wordprocessingml.document" TEXT_MIME = "text/plain" _TEXT_MIME_TYPES = frozenset( { "text/plain", "text/markdown", "text/x-markdown", "text/csv", "application/json", } ) _DOCX_MAIN_TYPE = "application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml" _ALLOWED_REDIRECT_CODES = frozenset({301, 302, 303, 307, 308}) # OpenAI defines file_id as a required string without a character-set format. # Keep it opaque: this value is an identity only and never a path or URL part. _FILE_ID_MAX_CHARS = 128 _SAFE_INPUT_REASONS = frozenset( { "actor_context_missing", "attachment_not_object", "attachment_unknown_fields", "attachment_required_fields_missing", "attachment_file_id_invalid", "attachment_file_id_empty", "attachment_file_id_too_long", "attachment_file_id_non_printable", "attachment_mime_type_invalid", "attachment_file_name_invalid", "files_container_invalid", "duplicate_file_ids", "acquired_file_id_mismatch", } ) _DNS_LABEL = re.compile(r"^[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?$") _URL_MAX_CHARS = 8192 _TLS_SHUTDOWN_GRACE_SECONDS = 0.5 _MIB = 1024 * 1024 _HARD_MAX_FILES = 20 _HARD_MAX_FILE_BYTES = 20 * _MIB _HARD_MAX_TOTAL_BYTES = 50 * _MIB _HARD_MAX_REDIRECTS = 3 _HARD_MAX_TEXT_CHARS = 1_000_000 _HARD_MAX_DOCX_UNCOMPRESSED = 25 * _MIB _LEGACY_PDF_CONFIGURATION_LOCK = threading.RLock() class FileInput(TypedDict): download_url: str file_id: str mime_type: NotRequired[str] file_name: NotRequired[str] class MalwareScanner(Protocol): async def require_clean(self, path: Path, *, deadline: float) -> None: ... class FileIngestionError(Exception): """Safe, stable failure without request URLs, local paths, or source data.""" _MESSAGES: Mapping[str, str] = { "invalid_input": "File input is invalid.", "url_not_allowed": "File URL is not allowed.", "source_limit": "File set exceeds its allowed limit.", "download_failed": "File download failed.", "download_limit": "Downloaded file exceeds its allowed limit.", "download_timeout": "File download exceeded its time limit.", "scan_unavailable": "Malware scanning is unavailable.", "scan_failed": "File did not pass malware scanning.", "malware_detected": "File did not pass malware scanning.", "unsupported_format": "File format is not supported.", "mime_mismatch": "File MIME type does not match its contents.", "extraction_failed": "File text could not be extracted safely.", "extraction_limit": "File text exceeds its allowed limit.", "extraction_timeout": "File extraction exceeded its time limit.", "configuration": "File ingestion configuration is invalid.", } def __init__(self, code: str, *, input_reason: str | None = None) -> None: self.code = code if code in self._MESSAGES else "download_failed" self.input_reason = ( input_reason if self.code == "invalid_input" and isinstance(input_reason, str) and input_reason in _SAFE_INPUT_REASONS else None ) super().__init__(self._MESSAGES[self.code]) @dataclass(frozen=True, slots=True) class FilePolicy: max_files: int = 5 max_file_bytes: int = 10 * _MIB max_total_bytes: int = 25 * _MIB max_redirects: int = 2 download_timeout_seconds: float = 15.0 scan_timeout_seconds: float = 20.0 extraction_timeout_seconds: float = 15.0 max_text_chars_per_file: int = 100_000 max_total_text_chars: int = 250_000 max_pdf_pages: int = 50 max_pdf_decompressed_bytes: int = 10 * _MIB max_docx_entries: int = 256 max_docx_relationships: int = 256 max_docx_xml_elements: int = 100_000 max_docx_uncompressed_bytes: int = 20 * _MIB max_docx_compression_ratio: int = 100 scanner: MalwareScanner | None = field(default=None, repr=False, compare=False) temp_directory: Path | None = field(default=None, repr=False, compare=False) def __post_init__(self) -> None: bounds = { "max_files": (self.max_files, 1, _HARD_MAX_FILES), "max_file_bytes": (self.max_file_bytes, 1, _HARD_MAX_FILE_BYTES), "max_total_bytes": (self.max_total_bytes, 1, _HARD_MAX_TOTAL_BYTES), "max_redirects": (self.max_redirects, 0, _HARD_MAX_REDIRECTS), "max_text_chars_per_file": (self.max_text_chars_per_file, 1, _HARD_MAX_TEXT_CHARS), "max_total_text_chars": (self.max_total_text_chars, 1, _HARD_MAX_TEXT_CHARS), "max_pdf_pages": (self.max_pdf_pages, 1, 200), "max_pdf_decompressed_bytes": (self.max_pdf_decompressed_bytes, 1, 20 * _MIB), "max_docx_entries": (self.max_docx_entries, 1, 512), "max_docx_relationships": (self.max_docx_relationships, 1, 1024), "max_docx_xml_elements": (self.max_docx_xml_elements, 1, 200_000), "max_docx_uncompressed_bytes": ( self.max_docx_uncompressed_bytes, 1, _HARD_MAX_DOCX_UNCOMPRESSED, ), "max_docx_compression_ratio": (self.max_docx_compression_ratio, 1, 500), } for name, (value, minimum, int_upper_bound) in bounds.items(): if ( isinstance(value, bool) or not isinstance(value, int) or not minimum <= value <= int_upper_bound ): raise ValueError(f"{name} is outside its safe range") for name, timeout_value, time_upper_bound in ( ("download_timeout_seconds", self.download_timeout_seconds, 60.0), ("scan_timeout_seconds", self.scan_timeout_seconds, 60.0), ("extraction_timeout_seconds", self.extraction_timeout_seconds, 30.0), ): if ( isinstance(timeout_value, bool) or not isinstance(timeout_value, int | float) or not 0.1 <= timeout_value <= time_upper_bound ): raise ValueError(f"{name} is outside its safe range") if self.temp_directory is not None and not isinstance(self.temp_directory, Path): raise ValueError("temp_directory must be a Path") @classmethod def from_environment(cls, *, scanner: MalwareScanner | None = None) -> FilePolicy: defaults: dict[str, tuple[str, int, int]] = { "max_files": ("INGESTION_MAX_FILES", 5, _HARD_MAX_FILES), "max_file_bytes": ("INGESTION_MAX_FILE_BYTES", 10 * _MIB, _HARD_MAX_FILE_BYTES), "max_total_bytes": ("INGESTION_MAX_TOTAL_BYTES", 25 * _MIB, _HARD_MAX_TOTAL_BYTES), "max_redirects": ("INGESTION_MAX_REDIRECTS", 2, _HARD_MAX_REDIRECTS), "max_text_chars_per_file": ( "INGESTION_MAX_TEXT_CHARS_PER_FILE", 100_000, _HARD_MAX_TEXT_CHARS, ), "max_total_text_chars": ( "INGESTION_MAX_TOTAL_TEXT_CHARS", 250_000, _HARD_MAX_TEXT_CHARS, ), "max_pdf_pages": ("INGESTION_MAX_PDF_PAGES", 50, 200), "max_pdf_decompressed_bytes": ( "INGESTION_MAX_PDF_DECOMPRESSED_BYTES", 10 * _MIB, 20 * _MIB, ), "max_docx_entries": ("INGESTION_MAX_DOCX_ENTRIES", 256, 512), "max_docx_relationships": ("INGESTION_MAX_DOCX_RELATIONSHIPS", 256, 1024), "max_docx_xml_elements": ("INGESTION_MAX_DOCX_XML_ELEMENTS", 100_000, 200_000), "max_docx_uncompressed_bytes": ( "INGESTION_MAX_DOCX_UNCOMPRESSED_BYTES", 20 * _MIB, _HARD_MAX_DOCX_UNCOMPRESSED, ), "max_docx_compression_ratio": ("INGESTION_MAX_DOCX_COMPRESSION_RATIO", 100, 500), } values: dict[str, int] = {} for field_name, (env_name, int_default, int_upper_bound) in defaults.items(): raw = os.environ.get(env_name) try: value = int_default if raw is None else int(raw, 10) except ValueError as exc: raise ValueError(f"{env_name} must be an integer") from exc if not 0 <= value <= int_upper_bound: raise ValueError(f"{env_name} is outside its safe range") values[field_name] = value time_defaults = { "download_timeout_seconds": ("INGESTION_DOWNLOAD_TIMEOUT_SECONDS", 15.0, 60.0), "scan_timeout_seconds": ("INGESTION_SCAN_TIMEOUT_SECONDS", 20.0, 60.0), "extraction_timeout_seconds": ("INGESTION_EXTRACTION_TIMEOUT_SECONDS", 15.0, 30.0), } time_values: dict[str, float] = {} for field_name, (env_name, time_default, time_upper_bound) in time_defaults.items(): raw = os.environ.get(env_name) try: timeout_value = time_default if raw is None else float(raw) except ValueError as exc: raise ValueError(f"{env_name} must be a number") from exc if not 0.1 <= timeout_value <= time_upper_bound: raise ValueError(f"{env_name} is outside its safe range") time_values[field_name] = timeout_value selected_scanner = scanner if scanner is not None else scanner_from_environment() return cls( max_files=values["max_files"], max_file_bytes=values["max_file_bytes"], max_total_bytes=values["max_total_bytes"], max_redirects=values["max_redirects"], download_timeout_seconds=time_values["download_timeout_seconds"], scan_timeout_seconds=time_values["scan_timeout_seconds"], extraction_timeout_seconds=time_values["extraction_timeout_seconds"], max_text_chars_per_file=values["max_text_chars_per_file"], max_total_text_chars=values["max_total_text_chars"], max_pdf_pages=values["max_pdf_pages"], max_pdf_decompressed_bytes=values["max_pdf_decompressed_bytes"], max_docx_entries=values["max_docx_entries"], max_docx_relationships=values["max_docx_relationships"], max_docx_xml_elements=values["max_docx_xml_elements"], max_docx_uncompressed_bytes=values["max_docx_uncompressed_bytes"], max_docx_compression_ratio=values["max_docx_compression_ratio"], scanner=selected_scanner, ) @dataclass(frozen=True, slots=True) class ProcessedSource: file_id: str file_name: str mime_type: str sha256: str text: str = field(repr=False) byte_size: int content: bytes = field(repr=False) @dataclass(frozen=True, slots=True) class DownloadedSource: file_id: str file_name: str path: Path = field(repr=False) byte_size: int sha256: str content_type: str | None = field(repr=False) def secure_cleanup(self) -> None: try: file_stat = self.path.lstat() if stat.S_ISREG(file_stat.st_mode) and not stat.S_ISLNK(file_stat.st_mode): with self.path.open("r+b", buffering=0) as source: remaining = file_stat.st_size zeroes = bytes(64 * 1024) while remaining: count = min(remaining, len(zeroes)) source.write(zeroes[:count]) remaining -= count source.flush() os.fsync(source.fileno()) except OSError: pass with suppress(OSError): self.path.unlink(missing_ok=True) def scanner_from_environment() -> ClamdScanner | None: host = os.environ.get("INGESTION_SCANNER_HOST", "").strip() if not host: return None raw_port = os.environ.get("INGESTION_SCANNER_PORT", "3310") try: port = int(raw_port, 10) except ValueError as exc: raise ValueError("INGESTION_SCANNER_PORT must be an integer") from exc return ClamdScanner(host, port=port) @dataclass(frozen=True, slots=True) class ClamdScanner: host: str = field(repr=False) port: int = 3310 chunk_bytes: int = 64 * 1024 max_response_bytes: int = 4096 def __post_init__(self) -> None: if ( not isinstance(self.host, str) or not self.host or len(self.host) > 253 or any(character.isspace() for character in self.host) or "://" in self.host or "/" in self.host ): raise ValueError("scanner host is invalid") if ( isinstance(self.port, bool) or not isinstance(self.port, int) or not 1 <= self.port <= 65535 ): raise ValueError("scanner port is invalid") if self.chunk_bytes != 64 * 1024 or self.max_response_bytes != 4096: raise ValueError("scanner limits are fixed") async def require_clean(self, path: Path, *, deadline: float) -> None: loop = asyncio.get_running_loop() remaining = deadline - loop.time() if remaining <= 0: raise FileIngestionError("scan_failed") writer: asyncio.StreamWriter | None = None try: async with asyncio.timeout(remaining): reader, writer = await asyncio.open_connection( self.host, self.port, limit=self.max_response_bytes, ) writer.write(b"zINSTREAM\0") with path.open("rb") as source: while chunk := source.read(self.chunk_bytes): writer.write(struct.pack("!I", len(chunk))) writer.write(chunk) await writer.drain() writer.write(b"\0\0\0\0") await writer.drain() response = await reader.readuntil(b"\0") if len(response) > self.max_response_bytes: raise FileIngestionError("scan_failed") if response == b"stream: OK\0": return if response.startswith(b"stream: ") and response.endswith(b" FOUND\0"): raise FileIngestionError("malware_detected") raise FileIngestionError("scan_failed") except FileIngestionError: raise except Exception: raise FileIngestionError("scan_failed") from None finally: if writer is not None: writer.close() with suppress(Exception): await writer.wait_closed() async def _system_resolver(host: str) -> Sequence[str]: try: records = await asyncio.get_running_loop().getaddrinfo( host, 443, type=socket.SOCK_STREAM, proto=socket.IPPROTO_TCP, ) return tuple(dict.fromkeys(str(record[4][0]) for record in records)) except Exception: raise FileIngestionError("url_not_allowed") from None def _canonical_host(host: str) -> str: if not host or len(host) > 253 or "%" in host: raise FileIngestionError("url_not_allowed") try: address = ipaddress.ip_address(host) except ValueError: pass else: _require_public_address(address) return str(address) candidate = host.rstrip(".") if not candidate or "." not in candidate: raise FileIngestionError("url_not_allowed") try: ascii_host = candidate.encode("idna").decode("ascii").lower() except UnicodeError: raise FileIngestionError("url_not_allowed") from None if len(ascii_host) > 253: raise FileIngestionError("url_not_allowed") labels = ascii_host.split(".") if any(not _DNS_LABEL.fullmatch(label) for label in labels): raise FileIngestionError("url_not_allowed") if ascii_host == "localhost" or ascii_host.endswith((".localhost", ".local", ".internal")): raise FileIngestionError("url_not_allowed") if ascii_host in {"metadata", "metadata.google.internal"}: raise FileIngestionError("url_not_allowed") return ascii_host def _require_public_address(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> None: if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None: raise FileIngestionError("url_not_allowed") if ( not address.is_global or address.is_private or address.is_loopback or address.is_link_local or address.is_multicast or address.is_reserved or address.is_unspecified ): raise FileIngestionError("url_not_allowed") async def _resolve_public_addresses( host: str, resolver: Resolver, ) -> tuple[ipaddress.IPv4Address | ipaddress.IPv6Address, ...]: canonical = _canonical_host(host) try: literal = ipaddress.ip_address(canonical) except ValueError: literal = None if literal is not None: _require_public_address(literal) return (literal,) try: raw_addresses = tuple(await resolver(canonical)) except FileIngestionError: raise except Exception: raise FileIngestionError("url_not_allowed") from None if not raw_addresses or len(raw_addresses) > 32: raise FileIngestionError("url_not_allowed") addresses: list[ipaddress.IPv4Address | ipaddress.IPv6Address] = [] for raw_address in raw_addresses: try: address = ipaddress.ip_address(raw_address) except (TypeError, ValueError): raise FileIngestionError("url_not_allowed") from None _require_public_address(address) if address not in addresses: addresses.append(address) return tuple(addresses) async def validate_download_url(url: str, *, resolver: Resolver = _system_resolver) -> str: if ( not isinstance(url, str) or not url or len(url) > _URL_MAX_CHARS or any(ord(character) < 0x20 or ord(character) == 0x7F for character in url) ): raise FileIngestionError("url_not_allowed") try: parts = urlsplit(url) host = parts.hostname port = parts.port except ValueError: raise FileIngestionError("url_not_allowed") from None if ( parts.scheme.lower() != "https" or not parts.netloc or host is None or parts.username is not None or parts.password is not None or parts.fragment or (port is not None and port != 443) ): raise FileIngestionError("url_not_allowed") await _resolve_public_addresses(host, resolver) return url class _AsyncioNetworkStream(httpcore.AsyncNetworkStream): def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None: self._reader = reader self._writer = writer async def read(self, max_bytes: int, timeout: float | None = None) -> bytes: if timeout is None: return await self._reader.read(max_bytes) async with asyncio.timeout(timeout): return await self._reader.read(max_bytes) async def write(self, buffer: bytes, timeout: float | None = None) -> None: self._writer.write(buffer) if timeout is None: await self._writer.drain() else: async with asyncio.timeout(timeout): await self._writer.drain() async def aclose(self) -> None: self._writer.close() try: async with asyncio.timeout(_TLS_SHUTDOWN_GRACE_SECONDS): await self._writer.wait_closed() except TimeoutError: self._abort_transport() except asyncio.CancelledError: self._abort_transport() raise except (OSError, RuntimeError): self._abort_transport() def _abort_transport(self) -> None: with suppress(AttributeError, OSError, RuntimeError): self._writer.transport.abort() async def start_tls( self, ssl_context: ssl.SSLContext, server_hostname: str | None = None, timeout: float | None = None, ) -> _AsyncioNetworkStream: if timeout is None: await self._writer.start_tls(ssl_context, server_hostname=server_hostname) else: async with asyncio.timeout(timeout): await self._writer.start_tls( ssl_context, server_hostname=server_hostname, ssl_handshake_timeout=timeout, ) return self def get_extra_info(self, info: str) -> Any: if info == "ssl_object": return self._writer.get_extra_info("ssl_object") if info == "client_addr": return self._writer.get_extra_info("sockname") if info == "server_addr": return self._writer.get_extra_info("peername") if info == "socket": return self._writer.get_extra_info("socket") return None class PinnedNetworkBackend(httpcore.AsyncNetworkBackend): """Resolve and validate immediately before connecting, then dial the IP literal.""" def __init__(self, *, resolver: Resolver = _system_resolver) -> None: self._resolver = resolver async def connect_tcp( self, host: str, port: int, timeout: float | None = None, local_address: str | None = None, socket_options: Iterable[SocketOption] | None = None, ) -> httpcore.AsyncNetworkStream: if port != 443: raise FileIngestionError("url_not_allowed") addresses = await _resolve_public_addresses(host, self._resolver) for address in addresses: family = socket.AF_INET6 if address.version == 6 else socket.AF_INET sock = socket.socket(family, socket.SOCK_STREAM, socket.IPPROTO_TCP) sock.setblocking(False) try: if local_address is not None: sock.bind((local_address, 0)) for option in socket_options or (): sock.setsockopt(*option) sockaddr: tuple[Any, ...] = ( (str(address), port, 0, 0) if address.version == 6 else (str(address), port) ) operation = asyncio.get_running_loop().sock_connect(sock, sockaddr) if timeout is None: await operation else: async with asyncio.timeout(timeout): await operation reader, writer = await asyncio.open_connection(sock=sock) return _AsyncioNetworkStream(reader, writer) except FileIngestionError: sock.close() raise except OSError: sock.close() raise FileIngestionError("download_failed") from None async def connect_unix_socket( self, path: str, timeout: float | None = None, socket_options: Iterable[SocketOption] | None = None, ) -> httpcore.AsyncNetworkStream: raise FileIngestionError("url_not_allowed") async def sleep(self, seconds: float) -> None: await asyncio.sleep(seconds) class _HttpcoreResponseStream(httpx.AsyncByteStream): def __init__(self, stream: Any) -> None: self._stream = stream async def __aiter__(self) -> AsyncIterator[bytes]: async for part in self._stream: yield part async def aclose(self) -> None: await self._stream.aclose() class _PinnedHTTPTransport(httpx.AsyncBaseTransport): def __init__(self, *, resolver: Resolver) -> None: self._pool = httpcore.AsyncConnectionPool( ssl_context=ssl.create_default_context(), max_connections=1, max_keepalive_connections=0, keepalive_expiry=0, http1=True, http2=False, retries=0, network_backend=PinnedNetworkBackend(resolver=resolver), ) async def handle_async_request(self, request: httpx.Request) -> httpx.Response: core_request = httpcore.Request( method=request.method, url=httpcore.URL( scheme=request.url.raw_scheme, host=request.url.raw_host, port=request.url.port, target=request.url.raw_path, ), headers=request.headers.raw, content=request.stream, extensions=request.extensions, ) response = await self._pool.handle_async_request(core_request) return httpx.Response( status_code=response.status, headers=response.headers, stream=_HttpcoreResponseStream(response.stream), extensions=response.extensions, request=request, ) async def aclose(self) -> None: await self._pool.aclose() @dataclass(frozen=True, slots=True) class _ValidatedInput: url: str = field(repr=False) file_id: str mime_type: str | None file_name: str def _safe_file_name(value: object) -> str: if value is None: return "source" if not isinstance(value, str) or len(value) > 1024: raise FileIngestionError("invalid_input", input_reason="attachment_file_name_invalid") name = value.replace("\\", "/").rsplit("/", maxsplit=1)[-1] filtered = ( "".join( character for character in name if character.isprintable() and character not in {"/", "\\", "\x00"} ) .strip() .rstrip(". ") ) if filtered in {"", ".", ".."}: return "source" return filtered[:128] def _file_id_input_reason(value: object) -> str | None: if not isinstance(value, str): return "attachment_file_id_invalid" if not value: return "attachment_file_id_empty" if len(value) > _FILE_ID_MAX_CHARS: return "attachment_file_id_too_long" if not value.isprintable(): return "attachment_file_id_non_printable" if not value.strip(): return "attachment_file_id_empty" return None def _validate_file_input(item: object) -> _ValidatedInput: if not isinstance(item, dict): raise FileIngestionError("invalid_input", input_reason="attachment_not_object") if set(item) - { "download_url", "file_id", "mime_type", "file_name", }: raise FileIngestionError("invalid_input", input_reason="attachment_unknown_fields") if "download_url" not in item or "file_id" not in item: raise FileIngestionError( "invalid_input", input_reason="attachment_required_fields_missing" ) file_id = item["file_id"] file_id_reason = _file_id_input_reason(file_id) if file_id_reason is not None: raise FileIngestionError("invalid_input", input_reason=file_id_reason) mime_type = item.get("mime_type") if mime_type is not None and (not isinstance(mime_type, str) or not mime_type.strip()): raise FileIngestionError("invalid_input", input_reason="attachment_mime_type_invalid") if isinstance(mime_type, str): mime_type = mime_type.split(";", maxsplit=1)[0].strip().lower() return _ValidatedInput( url=item["download_url"] if isinstance(item["download_url"], str) else "", file_id=file_id, mime_type=mime_type, file_name=_safe_file_name(item.get("file_name")), ) def _normalized_mime(value: str | None) -> str | None: if value is None: return None return value.split(";", maxsplit=1)[0].strip().lower() def _mime_matches(claim: str | None, detected: str) -> bool: normalized = _normalized_mime(claim) if normalized is None or normalized == "application/octet-stream": return True if detected == TEXT_MIME: return normalized in _TEXT_MIME_TYPES return normalized == detected def _read_zip_member(archive: zipfile.ZipFile, name: str, limit: int) -> bytes: try: info = archive.getinfo(name) except KeyError: raise FileIngestionError("unsupported_format") from None if info.file_size > limit or info.compress_size < 0: raise FileIngestionError("extraction_limit") if info.flag_bits & 0x1: raise FileIngestionError("unsupported_format") with archive.open(info, "r") as entry: value = entry.read(limit + 1) if len(value) > limit: raise FileIngestionError("extraction_limit") return value def _parse_xml_bounded(data: bytes, *, element_limit: int) -> tuple[Any, int]: if b" element_limit: raise FileIngestionError("extraction_limit") if root is None: root = element except FileIngestionError: raise except Exception: raise FileIngestionError("extraction_failed") from None if root is None: raise FileIngestionError("extraction_failed") return root, element_count def _extract_docx(data: bytes, policy: FilePolicy) -> tuple[str, bool]: try: with zipfile.ZipFile(BytesIO(data), "r") as archive: infos = archive.infolist() if not infos or len(infos) > policy.max_docx_entries: raise FileIngestionError("extraction_limit") names: set[str] = set() expanded_bytes = 0 for info in infos: name = info.filename if ( name.startswith(("/", "\\")) or ".." in name.replace("\\", "/").split("/") or len(name) > 512 or name in names ): raise FileIngestionError("unsupported_format") names.add(name) if info.flag_bits & 0x1 or info.file_size < 0: raise FileIngestionError("unsupported_format") expanded_bytes += info.file_size if expanded_bytes > policy.max_docx_uncompressed_bytes: raise FileIngestionError("extraction_limit") if info.file_size and ( info.compress_size == 0 or info.file_size / max(info.compress_size, 1) > policy.max_docx_compression_ratio ): raise FileIngestionError("extraction_limit") content_types_data = _read_zip_member( archive, "[Content_Types].xml", min(policy.max_docx_uncompressed_bytes, 1 * _MIB), ) document_data = _read_zip_member( archive, "word/document.xml", min(policy.max_docx_uncompressed_bytes, 10 * _MIB), ) content_root, _ = _parse_xml_bounded( content_types_data, element_limit=policy.max_docx_xml_elements, ) declares_docx = any( element.attrib.get("PartName") == "/word/document.xml" and element.attrib.get("ContentType") == _DOCX_MAIN_TYPE for element in content_root.iter() ) if not declares_docx: raise FileIngestionError("unsupported_format") rels_name = "word/_rels/document.xml.rels" if rels_name in names: relationships_data = _read_zip_member( archive, rels_name, min(policy.max_docx_uncompressed_bytes, 2 * _MIB), ) relationships_root, relation_elements = _parse_xml_bounded( relationships_data, element_limit=policy.max_docx_xml_elements, ) relationships = [ element for element in relationships_root.iter() if element.tag.endswith("}Relationship") ] if len(relationships) > policy.max_docx_relationships: raise FileIngestionError("extraction_limit") if any( element.attrib.get("TargetMode", "").lower() == "external" for element in relationships ): raise FileIngestionError("unsupported_format") if relation_elements > policy.max_docx_xml_elements: raise FileIngestionError("extraction_limit") if len(document_data) > policy.max_docx_uncompressed_bytes: raise FileIngestionError("extraction_limit") if b" policy.max_docx_xml_elements: raise FileIngestionError("extraction_limit") if element.tag.endswith(namespace_text) and element.text: text_parts.append(element.text) text_length += len(element.text) elif element.tag.endswith(namespace_paragraph): text_parts.append("\n") text_length += 1 if text_length > policy.max_text_chars_per_file: raise FileIngestionError("extraction_limit") element.clear() except FileIngestionError: raise except Exception: raise FileIngestionError("extraction_failed") from None return "".join(text_parts), True except FileIngestionError: raise except (OSError, zipfile.BadZipFile, RuntimeError, ValueError): raise FileIngestionError("unsupported_format") from None def _extract_pdf(data: bytes, policy: FilePolicy) -> str: decompressed_limit = min(policy.max_pdf_decompressed_bytes, 20 * _MIB) try: with _pdf_configuration( maximum_declared_stream_length=decompressed_limit, array_based_stream_maximum_output_length=decompressed_limit, jbig2_maximum_output_length=decompressed_limit, lzw_maximum_output_length=decompressed_limit, run_length_maximum_output_length=decompressed_limit, zlib_maximum_output_length=decompressed_limit, zlib_maximum_recovery_input_length=min(decompressed_limit, 512 * 1024), flate_maximum_columns=100_000, flate_maximum_row_length=min(decompressed_limit, 2 * _MIB), image_maximum_buffer_size=min(decompressed_limit, 2 * _MIB), xmp_maximum_input_length=min(decompressed_limit, 1 * _MIB), xmp_maximum_element_count=10_000, outline_maximum_entries=5_000, outline_maximum_depth=32, page_tree_maximum_entries=policy.max_pdf_pages + 1, page_tree_maximum_depth=32, xform_maximum_invocations_per_extraction=500, jbig2dec_binary=None, ): reader = PdfReader(BytesIO(data), strict=True, root_object_recovery_limit=10_000) try: if reader.is_encrypted: raise FileIngestionError("unsupported_format") pages = reader.pages if len(pages) > policy.max_pdf_pages: raise FileIngestionError("extraction_limit") text_parts: list[str] = [] text_length = 0 decompressed_total = 0 for page in pages: contents = page.get_contents() if contents is not None: decompressed_total += len(contents.get_data()) if decompressed_total > decompressed_limit: raise FileIngestionError("extraction_limit") text = page.extract_text() or "" text_parts.append(text) text_length += len(text) if text_length > policy.max_text_chars_per_file: raise FileIngestionError("extraction_limit") return "\n".join(text_parts) finally: reader.close() except FileIngestionError: raise except Exception: raise FileIngestionError("extraction_failed") from None @contextmanager def _pdf_configuration(**settings: Any) -> Iterator[None]: """Use pypdf's scoped limits, with a locked compatibility path for older installs.""" if _pypdf_apply_configuration is not None: with _pypdf_apply_configuration(**settings): yield return import pypdf.filters as filters legacy_names = { "maximum_declared_stream_length": "MAX_DECLARED_STREAM_LENGTH", "array_based_stream_maximum_output_length": "MAX_ARRAY_BASED_STREAM_OUTPUT_LENGTH", "jbig2_maximum_output_length": "JBIG2_MAX_OUTPUT_LENGTH", "lzw_maximum_output_length": "LZW_MAX_OUTPUT_LENGTH", "run_length_maximum_output_length": "RUN_LENGTH_MAX_OUTPUT_LENGTH", "zlib_maximum_output_length": "ZLIB_MAX_OUTPUT_LENGTH", "zlib_maximum_recovery_input_length": "ZLIB_MAX_RECOVERY_INPUT_LENGTH", "flate_maximum_columns": "FLATE_MAX_COLUMNS", "flate_maximum_row_length": "FLATE_MAX_ROW_LENGTH", "image_maximum_buffer_size": "FLATE_MAX_BUFFER_SIZE", "jbig2dec_binary": "JBIG2DEC_BINARY", } with _LEGACY_PDF_CONFIGURATION_LOCK: previous: dict[str, Any] = {} for setting_name, legacy_name in legacy_names.items(): if setting_name in settings: previous[legacy_name] = getattr(filters, legacy_name) setattr(filters, legacy_name, settings[setting_name]) try: yield finally: for legacy_name, value in previous.items(): setattr(filters, legacy_name, value) def _sniff_and_extract( data: bytes, *, declared_types: Sequence[str | None], policy: FilePolicy ) -> tuple[str, str]: if data.startswith(b"%PDF-"): detected = PDF_MIME extracted = _extract_pdf(data, policy) elif data.startswith(b"PK\x03\x04"): extracted, is_docx = _extract_docx(data, policy) if not is_docx: raise FileIngestionError("unsupported_format") detected = DOCX_MIME else: try: extracted = data.decode("utf-8-sig") except UnicodeDecodeError: raise FileIngestionError("unsupported_format") from None if "\x00" in extracted: raise FileIngestionError("unsupported_format") if any(ord(character) < 0x20 and character not in "\t\n\r\f" for character in extracted): raise FileIngestionError("unsupported_format") if len(extracted) > policy.max_text_chars_per_file: raise FileIngestionError("extraction_limit") detected = TEXT_MIME if any(not _mime_matches(declared, detected) for declared in declared_types): raise FileIngestionError("mime_mismatch") if len(extracted) > policy.max_text_chars_per_file: raise FileIngestionError("extraction_limit") if not extracted.strip(): raise FileIngestionError("extraction_failed") return detected, extracted class FileIngestionService: def __init__( self, policy: FilePolicy, *, resolver: Resolver = _system_resolver, transport: httpx.AsyncBaseTransport | None = None, ) -> None: self.policy = policy self.resolver = resolver self.transport = transport async def acquire( self, files: Sequence[FileInput], *, actor: ActorContext, ) -> tuple[ProcessedSource, ...]: if actor is None: raise FileIngestionError("invalid_input", input_reason="actor_context_missing") scanner = self.policy.scanner if scanner is None: raise FileIngestionError("scan_unavailable") raw_files: object = files if not isinstance(raw_files, Sequence) or isinstance(raw_files, str | bytes): raise FileIngestionError("invalid_input", input_reason="files_container_invalid") if not raw_files or len(raw_files) > self.policy.max_files: raise FileIngestionError("source_limit") validated = tuple(_validate_file_input(item) for item in raw_files) file_ids = [item.file_id for item in validated] if len(set(file_ids)) != len(file_ids): raise FileIngestionError("invalid_input", input_reason="duplicate_file_ids") sources: list[ProcessedSource] = [] total_download_bytes = 0 total_text_chars = 0 temp_root: str | None = None client: httpx.AsyncClient | None = None try: temp_root = tempfile.mkdtemp( prefix="itworx-ingestion-", dir=str(self.policy.temp_directory) if self.policy.temp_directory else None, ) os.chmod(temp_root, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) transport = self.transport or _PinnedHTTPTransport(resolver=self.resolver) client = httpx.AsyncClient( transport=transport, follow_redirects=False, trust_env=False, timeout=httpx.Timeout( self.policy.download_timeout_seconds, connect=min(5.0, self.policy.download_timeout_seconds), ), headers={ "Accept": ( "application/pdf, " "application/vnd.openxmlformats-officedocument.wordprocessingml.document, " "text/plain, text/markdown, text/csv, application/json, " "application/octet-stream" ), "Accept-Encoding": "identity", "User-Agent": "ITWorx-Knowledge-Intake/1.0", }, ) for source_input in validated: remaining_bytes = self.policy.max_total_bytes - total_download_bytes if remaining_bytes <= 0: raise FileIngestionError("download_limit") per_file_limit = min(self.policy.max_file_bytes, remaining_bytes) downloaded = await self._download( client, source_input, directory=Path(temp_root), byte_limit=per_file_limit, ) try: deadline = asyncio.get_running_loop().time() + self.policy.scan_timeout_seconds try: async with asyncio.timeout(self.policy.scan_timeout_seconds): await scanner.require_clean(downloaded.path, deadline=deadline) except Exception: raise FileIngestionError("scan_failed") from None content = downloaded.path.read_bytes() if ( len(content) != downloaded.byte_size or hashlib.sha256(content).hexdigest() != downloaded.sha256 ): raise FileIngestionError("download_failed") extracted_task = asyncio.create_task( asyncio.to_thread( _sniff_and_extract, content, declared_types=(source_input.mime_type, downloaded.content_type), policy=self.policy, ) ) try: detected_mime, text = await asyncio.wait_for( extracted_task, timeout=self.policy.extraction_timeout_seconds, ) except TimeoutError: extracted_task.cancel() raise FileIngestionError("extraction_timeout") from None total_download_bytes += downloaded.byte_size total_text_chars += len(text) if total_text_chars > self.policy.max_total_text_chars: raise FileIngestionError("extraction_limit") sources.append( ProcessedSource( file_id=downloaded.file_id, file_name=downloaded.file_name, mime_type=detected_mime, sha256=downloaded.sha256, text=text, byte_size=downloaded.byte_size, content=content, ) ) finally: downloaded.secure_cleanup() return tuple(sources) except FileIngestionError: raise except TimeoutError: raise FileIngestionError("download_timeout") from None except Exception: raise FileIngestionError("download_failed") from None finally: if client is not None: with suppress(Exception): await client.aclose() if temp_root is not None: shutil.rmtree(temp_root, ignore_errors=True) async def _download( self, client: httpx.AsyncClient, source: _ValidatedInput, *, directory: Path, byte_limit: int, ) -> DownloadedSource: url = source.url redirect_count = 0 digest = hashlib.sha256() received = 0 response_type: str | None = None descriptor = -1 destination: Path | None = None completed = False try: async with asyncio.timeout(self.policy.download_timeout_seconds): while True: url = await validate_download_url(url, resolver=self.resolver) async with client.stream( "GET", url, headers={"Accept-Encoding": "identity"} ) as response: if response.status_code in _ALLOWED_REDIRECT_CODES: location = response.headers.get("location") client.cookies.clear() if location is None or redirect_count >= self.policy.max_redirects: raise FileIngestionError("url_not_allowed") url = urljoin(url, location) redirect_count += 1 continue if response.status_code != 200: raise FileIngestionError("download_failed") content_encoding = ( response.headers.get("content-encoding", "identity").strip().lower() ) if content_encoding not in {"", "identity"}: raise FileIngestionError("unsupported_format") length = response.headers.get("content-length") if length is not None: try: advertised = int(length, 10) except ValueError: raise FileIngestionError("download_failed") from None if advertised < 0 or advertised > byte_limit: raise FileIngestionError("download_limit") response_type = _normalized_mime(response.headers.get("content-type")) descriptor, raw_path = tempfile.mkstemp( prefix="source-", suffix=".bin", dir=directory, ) destination = Path(raw_path) os.fchmod(descriptor, stat.S_IRUSR | stat.S_IWUSR) with os.fdopen(descriptor, "wb", closefd=True) as output: descriptor = -1 async for chunk in response.aiter_bytes(): received += len(chunk) if received > byte_limit: raise FileIngestionError("download_limit") digest.update(chunk) output.write(chunk) client.cookies.clear() if received <= 0: raise FileIngestionError("unsupported_format") downloaded = DownloadedSource( file_id=source.file_id, file_name=source.file_name, path=destination, byte_size=received, sha256=digest.hexdigest(), content_type=response_type, ) completed = True return downloaded except FileIngestionError: raise except TimeoutError: raise FileIngestionError("download_timeout") from None except Exception: raise FileIngestionError("download_failed") from None finally: if descriptor >= 0: with suppress(OSError): os.close(descriptor) if not completed and destination is not None and destination.exists(): secure = DownloadedSource( file_id=source.file_id, file_name=source.file_name, path=destination, byte_size=received, sha256="", content_type=None, ) secure.secure_cleanup() async def acquire_scan_extract( files: Sequence[FileInput], *, actor: ActorContext, policy: FilePolicy, ) -> tuple[ProcessedSource, ...]: return await FileIngestionService(policy).acquire(files, actor=actor) async def scan_text_source(text: str, file_name: str, *, policy: FilePolicy) -> ProcessedSource: """Bound and malware-scan caller-provided text before accepting it as a source.""" if not isinstance(text, str) or not isinstance(file_name, str): raise FileIngestionError("invalid_input") scanner = policy.scanner if scanner is None: raise FileIngestionError("scan_unavailable") if ( not text.strip() or len(text) > policy.max_text_chars_per_file or len(text) > policy.max_total_text_chars or "\x00" in text or any(ord(character) < 0x20 and character not in "\t\n\r\f" for character in text) ): raise FileIngestionError("extraction_limit" if text.strip() else "invalid_input") name = _safe_file_name(file_name) if name == "source": raise FileIngestionError("invalid_input") content = text.encode("utf-8") if len(content) > min(policy.max_file_bytes, policy.max_total_bytes): raise FileIngestionError("download_limit") root: str | None = None source_path: Path | None = None descriptor = -1 digest = hashlib.sha256(content).hexdigest() try: root = tempfile.mkdtemp( prefix="itworx-ingestion-text-", dir=str(policy.temp_directory) if policy.temp_directory else None, ) os.chmod(root, stat.S_IRUSR | stat.S_IWUSR | stat.S_IXUSR) descriptor, raw_path = tempfile.mkstemp(prefix="source-", suffix=".txt", dir=root) source_path = Path(raw_path) os.fchmod(descriptor, stat.S_IRUSR | stat.S_IWUSR) with os.fdopen(descriptor, "wb", closefd=True) as output: descriptor = -1 output.write(content) deadline = asyncio.get_running_loop().time() + policy.scan_timeout_seconds try: async with asyncio.timeout(policy.scan_timeout_seconds): await scanner.require_clean(source_path, deadline=deadline) except TimeoutError: raise FileIngestionError("scan_failed") from None except FileIngestionError: raise except Exception: raise FileIngestionError("scan_failed") from None return ProcessedSource( file_id="repo-text-source", file_name=name, mime_type=TEXT_MIME, sha256=digest, text=text, byte_size=len(content), content=content, ) except FileIngestionError: raise except Exception: raise FileIngestionError("scan_failed") from None finally: if descriptor >= 0: with suppress(OSError): os.close(descriptor) if source_path is not None and source_path.exists(): DownloadedSource( file_id="repo-text-source", file_name=name, path=source_path, byte_size=len(content), sha256=digest, content_type=TEXT_MIME, ).secure_cleanup() if root is not None: shutil.rmtree(root, ignore_errors=True)