PacketStation/server/packetstation/domain/assets.py

"""Stored files. One row per file in `assets`; the bytes live at ASSETS_DIR/<id[:2]>/<id>.

An upload is streamed to a temp file in ASSETS_DIR in chunks (never held in memory whole), hashed as it goes,
typed by its first bytes (not by its name or the browser's claim), fsynced and renamed into place. The same bytes
uploaded twice (same sha256 and size) return the existing row.
"""

from __future__ import annotations

import contextlib
import hashlib
import logging
import os
import re
import tempfile
import uuid
from pathlib import Path
from typing import Any, BinaryIO

from psycopg.rows import dict_row
from psycopg.types.json import Jsonb

from .. import audit
from ..api.errors import ApiError
from ..settings import Settings
from ..text import clean_text

log = logging.getLogger(__name__)

CHUNK = 1024 * 1024
HEAD = 16
PDF = "application/pdf"
MAGIC: tuple[tuple[bytes, str], ...] = (
    (b"%PDF-", PDF),
    (b"\x89PNG\r\n\x1a\n", "image/png"),
    (b"\xff\xd8\xff", "image/jpeg"),
)
EXT = {PDF: ".pdf", "image/png": ".png", "image/jpeg": ".jpg"}
NAME_MAX = 160
_UNSAFE_NAME = re.compile(r'[\\/:*?"<>|\x00-\x1f]')


def sniff(head: bytes) -> str | None:
    """The mime type from the first bytes, or None when it is not a type Packet Station stores."""
    for magic, mime in MAGIC:
        if head.startswith(magic):
            return mime
    return None


def safe_name(name: Any) -> str | None:
    """The original file name, cleaned: no path, no separators or reserved characters, at most 160 characters."""
    text = clean_text(name, 1000)
    text = text.replace("\\", "/").rsplit("/", 1)[-1]
    text = _UNSAFE_NAME.sub("_", text).strip(" .")
    return text[:NAME_MAX] or None


def storage_key(asset_id: uuid.UUID | str) -> str:
    s = str(asset_id)
    return f"{s[:2]}/{s}"


def path_of(settings: Settings, key: str) -> Path:
    root = Path(settings.assets_dir).resolve()
    p = (root / key).resolve()
    if not p.is_relative_to(root):
        raise ApiError(404, "not_found", "No such file.")
    return p


def pdf_pages(path: Path) -> int:
    """Page count through pypdfium2, or pypdf when pypdfium2 is missing or cannot read the file.

    422 bad_pdf when neither parses it or it has no pages.
    """
    n = _pages_pdfium(path)
    if n is None:
        n = _pages_pypdf(path)
    if n < 1:
        raise ApiError(422, "bad_pdf", "That PDF has no pages.")
    return n


def _pages_pdfium(path: Path) -> int | None:
    try:
        import pypdfium2
    except ImportError:
        return None
    try:
        doc = pypdfium2.PdfDocument(str(path))
    except Exception as e:  # any pdfium failure falls back to pypdf
        log.info("pypdfium2 could not read the PDF (%s); trying pypdf", type(e).__name__)
        return None
    try:
        return len(doc)
    finally:
        doc.close()


def _pages_pypdf(path: Path) -> int:
    from pypdf import PdfReader
    from pypdf.errors import PyPdfError

    try:
        n = len(PdfReader(str(path)).pages)
    except (PyPdfError, ValueError, KeyError, TypeError, OSError, IndexError, AttributeError) as e:
        log.info("upload is not a readable PDF: %s", type(e).__name__)
        raise ApiError(422, "bad_pdf", "That PDF could not be read.") from None
    return n


def public(row: dict[str, Any]) -> dict[str, Any]:
    return {
        "id": str(row["id"]),
        "mime": row["mime"],
        "size": int(row["size"]),
        "sha256": bytes(row["sha256"]).hex(),
        "pages": row["pages"],
        "original_name": row["original_name"],
        "url": f"/_blob/{row['id']}",
    }


def get(conn: Any, asset_id: Any) -> dict[str, Any] | None:
    try:
        aid = uuid.UUID(str(asset_id))
    except ValueError:
        return None
    with conn.cursor(row_factory=dict_row) as cur:
        r = cur.execute("SELECT * FROM assets WHERE id = %s", (aid,)).fetchone()
    return dict(r) if r else None


def _fsync_dir(path: Path) -> None:
    if os.name != "posix":
        return
    fd = os.open(str(path), os.O_RDONLY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def _existing(conn: Any, sha: bytes, size: int) -> dict[str, Any] | None:
    """The stored asset with the same bytes (same sha256 and size), oldest first."""
    with conn.cursor(row_factory=dict_row) as cur:
        r = cur.execute(
            "SELECT * FROM assets WHERE sha256 = %s AND size = %s ORDER BY uploaded_at LIMIT 1", (sha, size)
        ).fetchone()
    return dict(r) if r else None


def _place(settings: Settings, tmp: Path) -> tuple[uuid.UUID, str, Path]:
    """Rename a fsynced file to ASSETS_DIR/<id[:2]>/<id>, then fsync the directory. Returns (id, key, path)."""
    aid = uuid.uuid4()
    key = storage_key(aid)
    final = path_of(settings, key)
    final.parent.mkdir(parents=True, exist_ok=True)
    os.replace(tmp, final)
    _fsync_dir(final.parent)
    return aid, key, final


def _insert(conn: Any, values: tuple[Any, ...]) -> dict[str, Any]:
    """Insert one assets row from (id, sha256, size, mime, name, key, source, source_ref, pages, by, packet_id)."""
    v = list(values)
    if v[7] is not None and not isinstance(v[7], Jsonb):
        v[7] = Jsonb(v[7])
    with conn.cursor(row_factory=dict_row) as cur:
        row = cur.execute(
            "INSERT INTO assets (id, sha256, size, mime, original_name, storage_key, source, source_ref, pages,"
            " uploaded_by, packet_id) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) RETURNING *",
            v,
        ).fetchone()
    return dict(row)


def store_upload(
    conn: Any,
    settings: Settings,
    stream: BinaryIO,
    *,
    original_name: Any,
    user: Any,
    packet_id: int | None = None,
    source: str = "upload",
    source_ref: Any = None,
    actor: Any = None,
) -> dict[str, Any]:
    """Store one upload and return its public row. Raises 413 too_large, 415 unsupported_type, 422 bad_pdf.

    `user` is the uploader's user id; `actor` (the session Principal) is used for the audit row when given.
    """
    root = Path(settings.assets_dir)
    root.mkdir(parents=True, exist_ok=True)
    limit = int(settings.upload_max_bytes)
    fd, tmp_name = tempfile.mkstemp(prefix=".up-", dir=str(root))
    tmp = Path(tmp_name)
    keep = False
    placed: Path | None = None
    try:
        digest = hashlib.sha256()
        size = 0
        head = b""
        mime: str | None = None
        with os.fdopen(fd, "wb") as out:
            while True:
                chunk = stream.read(CHUNK)
                if not chunk:
                    break
                size += len(chunk)
                if size > limit:
                    raise ApiError(413, "too_large", f"The file is over the {limit // (1024 * 1024)} MB limit.")
                if mime is None and len(head) < HEAD:
                    head += chunk[: HEAD - len(head)]
                    if len(head) >= HEAD:
                        mime = sniff(head)
                        if mime is None:
                            raise ApiError(415, "unsupported_type", "Only PDF, PNG and JPEG files can be stored.")
                digest.update(chunk)
                out.write(chunk)
            out.flush()
            os.fsync(out.fileno())
        if mime is None:
            mime = sniff(head)
        if mime is None or size == 0:
            raise ApiError(415, "unsupported_type", "Only PDF, PNG and JPEG files can be stored.")
        sha = digest.digest()
        existing = _existing(conn, sha, size)
        if existing is not None:
            return public(existing)
        pages = pdf_pages(tmp) if mime == PDF else None
        aid, key, placed = _place(settings, tmp)
        keep = True
        name = safe_name(original_name)
        with conn.transaction():
            row = _insert(conn, (aid, sha, size, mime, name, key, source, source_ref, pages, user, packet_id))
            audit.append(
                conn,
                action="asset.uploaded",
                user_id=user,
                session_id=getattr(actor, "session_id", None),
                role_snapshot=getattr(actor, "role_snapshot", None),
                subject_type="asset",
                subject_id=str(aid),
                packet_id=packet_id,
                detail={"mime": mime, "size": size, "pages": pages, "sha256": sha.hex()},
                ip=getattr(actor, "ip", None),
            )
        return public(dict(row))
    except BaseException:
        if placed is not None:  # the row was not written: do not leave an orphan file
            with contextlib.suppress(OSError):
                placed.unlink()
        raise
    finally:
        if not keep:
            with contextlib.suppress(OSError):
                tmp.unlink()


def store_file(
    conn: Any,
    settings: Settings,
    path: Path,
    *,
    mime: str = PDF,
    original_name: Any,
    source: str,
    source_ref: Any = None,
    uploaded_by: Any = None,
    packet_id: int | None = None,
) -> dict[str, Any]:
    """Store a file the server wrote itself (the worker's M-Files download). Returns the full assets row.

    The file is consumed: renamed into ASSETS_DIR when its bytes are new, deleted when the same bytes are already
    stored (same sha256 and size; that row comes back). Write it under ASSETS_DIR (ASSETS_DIR/tmp) so the rename
    stays on one file system. It is hashed in chunks, its first bytes must match `mime`, and it is fsynced before
    the rename. The row carries one extra key, `reused`. No audit row: the caller writes its own.
    Raises 413 too_large, 415 unsupported_type, 422 bad_pdf.
    """
    path = Path(path)
    limit = int(settings.upload_max_bytes)
    placed: Path | None = None
    done = False
    try:
        digest = hashlib.sha256()
        size = 0
        head = b""
        with path.open("r+b") as f:
            while True:
                chunk = f.read(CHUNK)
                if not chunk:
                    break
                if len(head) < HEAD:
                    head += chunk[: HEAD - len(head)]
                size += len(chunk)
                if size > limit:
                    raise ApiError(413, "too_large", f"The file is over the {limit // (1024 * 1024)} MB limit.")
                digest.update(chunk)
            os.fsync(f.fileno())
        if size == 0 or sniff(head) != mime:
            raise ApiError(415, "unsupported_type", f"The file is not {EXT.get(mime, mime)}.")
        sha = digest.digest()
        existing = _existing(conn, sha, size)
        if existing is not None:
            return {**existing, "reused": True}
        pages = pdf_pages(path) if mime == PDF else None
        aid, key, placed = _place(settings, path)
        name = safe_name(original_name)
        with conn.transaction():
            row = _insert(conn, (aid, sha, size, mime, name, key, source, source_ref, pages, uploaded_by, packet_id))
        done = True
        return {**row, "reused": False}
    except BaseException:
        if placed is not None and not done:  # the row was not written: do not leave an orphan file
            with contextlib.suppress(OSError):
                placed.unlink()
        raise
    finally:
        with contextlib.suppress(OSError):
            path.unlink()


def open_blob(settings: Settings, row: dict[str, Any]) -> Path:
    """The stored file of an asset row; 404 when it is missing on disk."""
    p = path_of(settings, str(row["storage_key"]))
    if not p.is_file():
        log.error("asset %s is missing on disk", row["id"])
        raise ApiError(404, "not_found", "No such file.")
    return p


def download_name(row: dict[str, Any]) -> str:
    """A safe file name for Content-Disposition (ASCII only, quotes removed)."""
    name = row.get("original_name") or ""
    name = name.encode("ascii", "ignore").decode("ascii").replace('"', "").replace(";", "_").strip()
    if not name:
        name = str(row["id"]) + EXT.get(str(row["mime"]), "")
    return name