pdf-quiz-generator/backend/app/services/file_intake.py

"""What a person may upload, decided by the bytes rather than by the name.

Three questions, in this order, and all three before anything else touches the
file:

1. **Is it small enough.** Read one byte past the cap and refuse if that byte
   exists, so an enormous file is never held in memory to be measured.
2. **What is it, actually.** Sniffed from the leading bytes. An extension is a
   claim the uploader makes, and `report.pdf` holding a shell script is not a
   PDF; the type we act on is the one the file's own header states.
3. **What text is in it.** PDF and DOCX are parsed for their text; an image is
   read by the tool model. Nothing is executed, ever, and none of it is stored.

On code injection specifically, because it is the question people mean when
they ask about uploads: there is no path from an uploaded file to code that
runs. The extracted text is used as a *search query* — passed as a bound
parameter to Postgres, never concatenated into SQL — and as the content of a
message to a language model. It is never rendered as HTML (the reader escapes
raw HTML in Markdown), never written to disk here, never passed to a shell, and
never `eval`'d. A DOCX is a zip file, so it is also checked against the one
attack that shape does allow: a member that expands to far more than the cap.
"""
import io
import re
import zipfile

#: Two megabytes. Not a technical limit — a PDF of a lecture is comfortably
#: under it and a scanned textbook chapter is not, and the second one is a
#: different feature.
MAX_BYTES = 2 * 1024 * 1024

#: How far a DOCX may expand once unzipped. A 40 KB zip that becomes a
#: gigabyte of XML is the only interesting thing about the format's shape.
MAX_UNZIPPED = 40 * 1024 * 1024

PDF = "pdf"
DOCX = "docx"
IMAGE = "image"

IMAGE_TYPES = {
    b"\x89PNG\r\n\x1a\n": "image/png",
    b"\xff\xd8\xff": "image/jpeg",
    b"GIF87a": "image/gif",
    b"GIF89a": "image/gif",
}

#: What the upload control offers, and what the error message names. Kept here
#: so the two cannot drift apart.
ACCEPT = ".pdf,.docx,.png,.jpg,.jpeg,.gif,.webp"
ACCEPT_HUMAN = "a PDF, a Word document (.docx), or an image"


class Rejected(ValueError):
    """The file may not be taken, with the sentence to show the person."""


def media_type(data: bytes) -> str | None:
    """The image type these bytes are, or None if they are not an image."""
    for signature, mime in IMAGE_TYPES.items():
        if data.startswith(signature):
            return mime
    # RIFF....WEBP — the size sits between the two markers.
    if data[:4] == b"RIFF" and data[8:12] == b"WEBP":
        return "image/webp"
    return None


def kind_of(data: bytes) -> str | None:
    """pdf, docx, image — or None for anything else at all."""
    if data.startswith(b"%PDF-"):
        return PDF
    if media_type(data):
        return IMAGE
    # Every OOXML file is a zip whose first member is [Content_Types].xml. A
    # plain zip, a jar and an xlsx all start the same way, so the manifest is
    # what separates a Word document from them.
    if data[:2] == b"PK" and b"word/document.xml" in data[:4096] + data[-4096:]:
        return DOCX
    if data[:2] == b"PK":
        try:
            with zipfile.ZipFile(io.BytesIO(data)) as bundle:
                if "word/document.xml" in bundle.namelist():
                    return DOCX
        except zipfile.BadZipFile:
            return None
    return None


def read(upload) -> bytes:
    """The file's bytes, refused if there are too many of them.

    One byte past the cap is read deliberately: its presence is the test, and
    reading it costs one byte rather than the whole file.
    """
    data = upload.file.read(MAX_BYTES + 1)
    if len(data) > MAX_BYTES:
        raise Rejected(
            f"That file is over {MAX_BYTES // (1024 * 1024)} MB. Upload a "
            "smaller one — a single chapter or handout rather than a whole book.")
    if not data:
        raise Rejected("That file is empty.")
    return data


def _docx_text(data: bytes, limit: int) -> str:
    try:
        with zipfile.ZipFile(io.BytesIO(data)) as bundle:
            total = sum(info.file_size for info in bundle.infolist())
            if total > MAX_UNZIPPED:
                raise Rejected("That Word document unpacks to far more than its size suggests.")
            xml = bundle.read("word/document.xml").decode("utf-8", errors="ignore")
    except Rejected:
        raise
    except Exception:
        raise Rejected("That Word document could not be read.")
    # Paragraph and line breaks become newlines before the tags go, so the text
    # does not run together into one sentence.
    xml = re.sub(r"</w:p>", "\n", xml)
    xml = re.sub(r"<w:br[^>]*/?>", "\n", xml)
    xml = re.sub(r"<[^>]+>", "", xml)
    return clean(xml)[:limit]


def _pdf_text(data: bytes, limit: int) -> str:
    try:
        import fitz

        with fitz.open(stream=data, filetype="pdf") as document:
            return clean(" ".join(page.get_text() for page in document))[:limit]
    except Exception:
        raise Rejected("That PDF could not be read.")


def clean(text: str) -> str:
    """Printable text, with the control characters gone.

    Not a security measure — nothing downstream interprets them — but a file
    full of NULs makes a query that matches nothing and an error nobody can
    read, and stripping them is one line.
    """
    text = text.replace("\x00", " ")
    text = re.sub(r"[\x01-\x08\x0b\x0c\x0e-\x1f\x7f]", " ", text)
    return re.sub(r"[ \t]{2,}", " ", text).strip()


def text_from(data: bytes, kind: str, limit: int, *, describe=None) -> str:
    """The words in this file.

    `describe` is how an image becomes text: a callable taking (bytes, mime)
    and returning what the tool model read in it. Passed in rather than
    imported so this module stays a file-handling module — and so the caller,
    which knows whether a tool model is even configured, decides.
    """
    if kind == PDF:
        return _pdf_text(data, limit)
    if kind == DOCX:
        return _docx_text(data, limit)
    if kind == IMAGE:
        if describe is None:
            raise Rejected(
                "Reading an image needs a tool model, and no model is configured "
                "for that job. Upload a PDF or a Word document instead.")
        return clean(describe(data, media_type(data) or "image/png"))[:limit]
    raise Rejected(f"Upload {ACCEPT_HUMAN}.")