"""Stored files. One row per file in `assets`; the bytes live at ASSETS_DIR//. 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//, 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