generated from kgod/ai-review-template
131 lines
4.8 KiB
Python
131 lines
4.8 KiB
Python
"""Safe text extraction for the supported import formats."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
from multiprocessing import get_context
|
|
from queue import Empty
|
|
from pathlib import Path
|
|
from zipfile import ZipFile
|
|
|
|
from docx import Document
|
|
from pypdf import PdfReader
|
|
|
|
|
|
class ImportExtractionError(ValueError):
|
|
pass
|
|
|
|
|
|
# A docx is a zip: a tiny compressed upload can expand into huge XML and burn
|
|
# minutes of parser CPU (measured: 90 KB -> ~30 MB -> 86 s, ~3 s per MB). A real
|
|
# resume decompresses to well under 1 MB, so 10 MB is generous and still bounds
|
|
# worst-case parse time to seconds.
|
|
_MAX_DECOMPRESSED_BYTES = 10 * 1024 * 1024
|
|
_MAX_PDF_PAGES = 20
|
|
_MAX_PDF_TEXT_CHARACTERS = 100_000
|
|
_PDF_EXTRACTION_TIMEOUT_SECONDS = 10.0
|
|
|
|
|
|
def _extract_pdf_text_worker(content: bytes, result_queue: object) -> None:
|
|
"""Run pypdf in an isolated process so the parent can enforce a CPU deadline."""
|
|
try:
|
|
reader = PdfReader(BytesIO(content))
|
|
if len(reader.pages) > _MAX_PDF_PAGES:
|
|
raise ImportExtractionError("import_file_too_complex")
|
|
parts: list[str] = []
|
|
characters = 0
|
|
for page in reader.pages:
|
|
page_text = page.extract_text() or ""
|
|
characters += len(page_text)
|
|
if characters > _MAX_PDF_TEXT_CHARACTERS:
|
|
raise ImportExtractionError("import_file_too_complex")
|
|
if page_text:
|
|
parts.append(page_text)
|
|
result_queue.put(("ok", "\n".join(parts).strip()))
|
|
except ImportExtractionError as exc:
|
|
result_queue.put(("error", str(exc)))
|
|
except Exception:
|
|
result_queue.put(("error", "ocr_required"))
|
|
|
|
|
|
def _extract_pdf_text(content: bytes) -> str:
|
|
context = get_context("spawn")
|
|
result_queue = context.Queue(maxsize=1)
|
|
process = context.Process(target=_extract_pdf_text_worker, args=(content, result_queue))
|
|
process.start()
|
|
process.join(_PDF_EXTRACTION_TIMEOUT_SECONDS)
|
|
if process.is_alive():
|
|
process.terminate()
|
|
process.join()
|
|
raise ImportExtractionError("import_file_too_complex")
|
|
try:
|
|
status, value = result_queue.get(timeout=1.0)
|
|
except Empty as exc:
|
|
raise ImportExtractionError("ocr_required") from exc
|
|
finally:
|
|
result_queue.close()
|
|
result_queue.join_thread()
|
|
if status != "ok":
|
|
raise ImportExtractionError(value)
|
|
if not value:
|
|
raise ImportExtractionError("ocr_required")
|
|
return value
|
|
|
|
|
|
def _reject_decompression_bomb(content: bytes) -> None:
|
|
try:
|
|
with ZipFile(BytesIO(content)) as archive:
|
|
total = sum(info.file_size for info in archive.infolist())
|
|
except Exception as exc:
|
|
raise ImportExtractionError("invalid_import_file") from exc
|
|
if total > _MAX_DECOMPRESSED_BYTES:
|
|
raise ImportExtractionError("import_file_too_large")
|
|
|
|
|
|
def normalize_upload_name(file_name: str) -> tuple[str, str]:
|
|
safe_name = Path(file_name or "upload").name
|
|
if not safe_name or safe_name in {".", ".."}:
|
|
raise ImportExtractionError("invalid_file_name")
|
|
extension = Path(safe_name).suffix.lower()
|
|
if extension == ".doc":
|
|
raise ImportExtractionError("legacy_doc_unsupported")
|
|
if extension not in {".pdf", ".docx"}:
|
|
raise ImportExtractionError("unsupported_import_format")
|
|
return safe_name, extension
|
|
|
|
|
|
def validate_upload(*, extension: str, declared_mime: str | None, content: bytes) -> str:
|
|
if not content:
|
|
raise ImportExtractionError("empty_import_file")
|
|
if extension == ".pdf":
|
|
if not content.startswith(b"%PDF-"):
|
|
raise ImportExtractionError("invalid_import_file")
|
|
return "application/pdf"
|
|
if not content.startswith(b"PK"):
|
|
raise ImportExtractionError("invalid_import_file")
|
|
if declared_mime and declared_mime not in {
|
|
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
|
"application/octet-stream",
|
|
}:
|
|
raise ImportExtractionError("invalid_import_mime")
|
|
return "application/vnd.openxmlformats-officedocument.wordprocessingml.document"
|
|
|
|
|
|
def extract_text(*, extension: str, content: bytes) -> str:
|
|
if extension == ".pdf":
|
|
return _extract_pdf_text(content)
|
|
_reject_decompression_bomb(content)
|
|
try:
|
|
document = Document(BytesIO(content))
|
|
except Exception as exc:
|
|
raise ImportExtractionError("invalid_import_file") from exc
|
|
parts = [paragraph.text.strip() for paragraph in document.paragraphs if paragraph.text.strip()]
|
|
for table in document.tables:
|
|
for row in table.rows:
|
|
values = [cell.text.strip() for cell in row.cells if cell.text.strip()]
|
|
if values:
|
|
parts.append(" | ".join(values))
|
|
text = "\n".join(parts).strip()
|
|
if not text:
|
|
raise ImportExtractionError("empty_import_text")
|
|
return text |