fix(cv): isolate document parsing
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.
This commit is contained in:
cesnimda
@@ -0,0 +1,253 @@
|
||||
"""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())
|
||||
Reference in New Issue
Block a user