generated from kgod/ai-review-template
87 lines
3.2 KiB
Python
87 lines
3.2 KiB
Python
"""Safe text extraction for the supported import formats."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from io import BytesIO
|
|
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
|
|
|
|
|
|
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":
|
|
try:
|
|
reader = PdfReader(BytesIO(content))
|
|
text = "\n".join(page.extract_text() or "" for page in reader.pages).strip()
|
|
except Exception as exc:
|
|
raise ImportExtractionError("ocr_required") from exc
|
|
if not text:
|
|
raise ImportExtractionError("ocr_required")
|
|
return text
|
|
_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 |