"""Resource-bounded CV document decoder invoked only as a child process.""" from __future__ import annotations import argparse import io import json import os from pathlib import Path import re import sys import zipfile MAX_EXTRACT_CHARS = 200_000 MAX_PDF_PAGES = 40 MAX_DOCX_ENTRIES = 256 MAX_DOCX_UNCOMPRESSED_BYTES = 32 * 1024 * 1024 MAX_DOCX_ENTRY_BYTES = 8 * 1024 * 1024 MAX_DOCX_COMPRESSION_RATIO = 100 MAX_IMAGE_DIMENSION = 12_000 MAX_IMAGE_PIXELS = 40_000_000 MAX_PDF_PAGE_PIXELS = 12_000_000 MAX_PDF_OCR_PIXELS = 120_000_000 OCR_LANGUAGES = "eng" IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp"} class ParserError(Exception): def __init__(self, code: str): super().__init__(code) self.code = code def _apply_resource_limits(cpu_seconds: int, memory_bytes: int) -> None: if os.name != "posix": return import resource resource.setrlimit(resource.RLIMIT_CPU, (cpu_seconds, cpu_seconds + 1)) resource.setrlimit(resource.RLIMIT_AS, (memory_bytes, memory_bytes)) resource.setrlimit(resource.RLIMIT_FSIZE, (1024 * 1024, 1024 * 1024)) resource.setrlimit(resource.RLIMIT_NOFILE, (64, 64)) resource.setrlimit(resource.RLIMIT_CORE, (0, 0)) def normalize_text(value: str) -> str: value = value.replace("\x00", " ").replace("\r\n", "\n").replace("\r", "\n") lines: list[str] = [] blank = False for raw_line in value.split("\n"): line = re.sub(r"[\t\f\v]+", " ", raw_line).strip() if not line: if lines and not blank: lines.append("") blank = True continue lines.append(line) blank = False return "\n".join(lines).strip() def _check_extracted_size(text: str, limit: int = MAX_EXTRACT_CHARS) -> str: if len(text) > limit: raise ParserError("extracted_text_limit") return text def _check_image_size(image, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> None: width, height = image.size if width <= 0 or height <= 0 or width > MAX_IMAGE_DIMENSION or height > MAX_IMAGE_DIMENSION: raise ParserError("image_dimensions") if width * height > pixel_limit: raise ParserError("image_pixel_limit") if getattr(image, "n_frames", 1) != 1: raise ParserError("image_frames") def validate_docx_container( data: bytes, *, entry_limit: int = MAX_DOCX_ENTRIES, total_limit: int = MAX_DOCX_UNCOMPRESSED_BYTES, entry_size_limit: int = MAX_DOCX_ENTRY_BYTES, ratio_limit: int = MAX_DOCX_COMPRESSION_RATIO, ) -> None: try: with zipfile.ZipFile(io.BytesIO(data)) as archive: entries = archive.infolist() names = {entry.filename for entry in entries} if "[Content_Types].xml" not in names or "word/document.xml" not in names: raise ParserError("docx_signature") if len(entries) > entry_limit: raise ParserError("docx_entry_limit") total_size = 0 for entry in entries: total_size += entry.file_size if entry.file_size > entry_size_limit: raise ParserError("docx_entry_size") if entry.file_size > 0 and entry.compress_size == 0: raise ParserError("docx_ratio") if entry.compress_size > 0 and entry.file_size / entry.compress_size > ratio_limit: raise ParserError("docx_ratio") if total_size > total_limit: raise ParserError("docx_expanded_size") except ParserError: raise except (zipfile.BadZipFile, OSError) as exc: raise ParserError("docx_signature") from exc def extract_plain_text(data: bytes) -> str: try: decoded = data.decode("utf-8-sig") except UnicodeDecodeError as exc: raise ParserError("text_encoding") from exc return _check_extracted_size(normalize_text(decoded)) def extract_docx_text(data: bytes) -> str: validate_docx_container(data) from docx import Document document = Document(io.BytesIO(data)) parts: list[str] = [] blocks = document.iter_inner_content() if hasattr(document, "iter_inner_content") else document.paragraphs for block in blocks: if hasattr(block, "rows"): for row in block.rows: cells = [cell.text.strip() for cell in row.cells if cell.text and cell.text.strip()] if cells: parts.append(" | ".join(cells)) elif getattr(block, "text", "").strip(): text = block.text.strip() style = (getattr(getattr(block, "style", None), "name", "") or "").lower() if "role" in style and parts: parts.append("") parts.append(f"- {text}" if "bullet" in style else text) return _check_extracted_size(normalize_text("\n".join(parts))) def _ocr_image(image) -> str: import pytesseract if image.mode not in ("RGB", "L"): image = image.convert("RGB") return normalize_text(pytesseract.image_to_string(image, lang=OCR_LANGUAGES)) def extract_image_text(data: bytes, *, pixel_limit: int = MAX_IMAGE_PIXELS) -> str: from PIL import Image with Image.open(io.BytesIO(data)) as image: _check_image_size(image, pixel_limit=pixel_limit) return _check_extracted_size(_ocr_image(image)) def extract_pdf_text(data: bytes, *, page_limit: int = MAX_PDF_PAGES) -> tuple[str, bool, int]: from pypdf import PdfReader page_count = 0 extracted_pages: list[str] = [] try: reader = PdfReader(io.BytesIO(data)) page_count = len(reader.pages) if page_count > page_limit: raise ParserError("pdf_page_limit") for page in reader.pages: try: extracted_pages.append(page.extract_text(extraction_mode="layout") or "") except TypeError: extracted_pages.append(page.extract_text() or "") except ParserError: raise except Exception: extracted_pages = [] text = _check_extracted_size(normalize_text("\n".join(extracted_pages))) if len(text) >= 80: return text, False, page_count import fitz from PIL import Image document = fitz.open(stream=data, filetype="pdf") try: page_count = document.page_count if page_count > page_limit: raise ParserError("pdf_page_limit") ocr_pages: list[str] = [] total_pixels = 0 for page in document: pixmap = page.get_pixmap(matrix=fitz.Matrix(2, 2), alpha=False) page_pixels = pixmap.width * pixmap.height total_pixels += page_pixels if page_pixels > MAX_PDF_PAGE_PIXELS or total_pixels > MAX_PDF_OCR_PIXELS: raise ParserError("pdf_ocr_limit") with Image.open(io.BytesIO(pixmap.tobytes("png"))) as image: _check_image_size(image, pixel_limit=MAX_PDF_PAGE_PIXELS) ocr_pages.append(_ocr_image(image)) _check_extracted_size("\n".join(ocr_pages)) return _check_extracted_size(normalize_text("\n".join(ocr_pages))), True, page_count finally: document.close() def extract(extension: str, data: bytes) -> dict[str, object]: if extension in {".txt", ".md"}: text = extract_plain_text(data) return {"text": text, "ocr_used": False, "page_count": None} if extension == ".docx": text = extract_docx_text(data) return {"text": text, "ocr_used": False, "page_count": None} if extension == ".pdf": text, ocr_used, page_count = extract_pdf_text(data) return {"text": text, "ocr_used": ocr_used, "page_count": page_count} if extension in IMAGE_EXTENSIONS: text = extract_image_text(data) return {"text": text, "ocr_used": True, "page_count": 1} raise ParserError("unsupported_type") def _write_result(path: Path, result: dict[str, object]) -> None: path.write_text(json.dumps(result, ensure_ascii=False, separators=(",", ":")), encoding="utf-8") def main() -> int: parser = argparse.ArgumentParser(add_help=False) parser.add_argument("--input", required=True) parser.add_argument("--output", required=True) parser.add_argument("--extension", required=True) parser.add_argument("--cpu-seconds", type=int, default=20) parser.add_argument("--memory-bytes", type=int, default=512 * 1024 * 1024) args = parser.parse_args() output_path = Path(args.output) try: _apply_resource_limits(args.cpu_seconds, args.memory_bytes) data = Path(args.input).read_bytes() result = extract(args.extension.lower(), data) if not result.get("text"): raise ParserError("no_text") _write_result(output_path, {"ok": True, **result}) return 0 except ParserError as exc: _write_result(output_path, {"ok": False, "code": exc.code}) return 2 except BaseException: _write_result(output_path, {"ok": False, "code": "parser_failed"}) return 3 if __name__ == "__main__": sys.exit(main())