cesnimda
19c5251612
Run untrusted document decoders in a secret-free, resource-bounded child process and terminate its process tree on deadline. Harden the production container and enforce parser and lint gates in CI.
254 lines
8.9 KiB
Python
254 lines
8.9 KiB
Python
"""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())
|