1380 lines
55 KiB
Python
1380 lines
55 KiB
Python
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"<!DOCTYPE" in data.upper() or b"<!ENTITY" in data.upper():
|
|
raise FileIngestionError("unsupported_format")
|
|
root: Any | None = None
|
|
element_count = 0
|
|
try:
|
|
iterator = ET.iterparse(
|
|
BytesIO(data),
|
|
events=("start", "end"),
|
|
forbid_dtd=True,
|
|
forbid_entities=True,
|
|
forbid_external=True,
|
|
)
|
|
for event, element in iterator:
|
|
if event == "start":
|
|
element_count += 1
|
|
if element_count > 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"<!DOCTYPE" in document_data.upper() or b"<!ENTITY" in document_data.upper():
|
|
raise FileIngestionError("unsupported_format")
|
|
text_parts: list[str] = []
|
|
text_length = 0
|
|
element_count = 0
|
|
namespace_text = "}t"
|
|
namespace_paragraph = "}p"
|
|
try:
|
|
iterator = ET.iterparse(
|
|
BytesIO(document_data),
|
|
events=("end",),
|
|
forbid_dtd=True,
|
|
forbid_entities=True,
|
|
forbid_external=True,
|
|
)
|
|
for _, element in iterator:
|
|
element_count += 1
|
|
if element_count > 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)
|