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 content (not by its name or the browser's claim), fsynced, scanned when ClamAV is configured and
renamed into place. The same bytes uploaded twice (same sha256 and size) return the existing row.

Each upload has a kind (its purpose) that sets the types it may be and its size limit. See KINDS.
"""

from __future__ import annotations

import codecs
import contextlib
import hashlib
import logging
import os
import re
import shutil
import tempfile
import uuid
import zipfile
from dataclasses import dataclass
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
from . import scan

log = logging.getLogger(__name__)

CHUNK = 1024 * 1024
HEAD = 16
MIB = 1024 * 1024
PDF = "application/pdf"
PNG = "image/png"
JPEG = "image/jpeg"
XLSX = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
CSV = "text/csv"
ZIP_MAGIC = b"PK\x03\x04"
MAGIC: tuple[tuple[bytes, str], ...] = (
    (b"%PDF-", PDF),
    (b"\x89PNG\r\n\x1a\n", PNG),
    (b"\xff\xd8\xff", JPEG),
    (ZIP_MAGIC, XLSX),
)
EXT = {PDF: ".pdf", PNG: ".png", JPEG: ".jpg", XLSX: ".xlsx", CSV: ".csv"}
# Served as a download only, never shown in the browser.
ATTACHMENT_ONLY = frozenset({XLSX, CSV})
XLSX_PARTS = ("[Content_Types].xml", "xl/workbook.xml")
CSV_SAMPLE = 64 * 1024
_CSV_ALLOWED_C0 = frozenset("\t\r\n")
NAME_MAX = 160
_UNSAFE_NAME = re.compile(r'[\\/:*?"<>|\x00-\x1f]')
MSG_DOCUMENT_TYPE = "Only PDF, PNG and JPEG files can be stored."


@dataclass(frozen=True)
class Kind:
    """An upload kind: the types it may be (by content) and the setting that holds its size limit."""

    name: str
    mimes: frozenset[str]
    limit_setting: str
    wrong_type: str


KINDS: dict[str, Kind] = {
    "drawing": Kind("drawing", frozenset({PDF}), "upload_max_bytes", "A drawing is a PDF file."),
    "mtr": Kind("mtr", frozenset({PDF, PNG, JPEG}), "upload_max_bytes", "An MTR is a PDF, PNG or JPG file."),
    "document": Kind("document", frozenset({PDF, PNG, JPEG}), "upload_max_bytes", MSG_DOCUMENT_TYPE),
    "bom": Kind("bom", frozenset({XLSX, CSV}), "upload_max_bytes_bom", "A BOM is an Excel (.xlsx) or CSV file."),
}
DEFAULT_KIND = "document"
# Older client names for a kind.
KIND_ALIASES = {"doc": "document", "docs": "document"}


def kind_of(purpose: Any) -> Kind:
    """The upload kind for a `purpose` value. Missing or blank is `document`; anything unknown is 422."""
    if purpose is None:
        return KINDS[DEFAULT_KIND]
    name = str(purpose).strip().lower()
    if not name:
        return KINDS[DEFAULT_KIND]
    name = KIND_ALIASES.get(name, name)
    found = KINDS.get(name)
    if found is None:
        raise ApiError(
            422, "invalid_request", f"purpose must be one of {', '.join(KINDS)}.", extra={"field": "purpose"}
        )
    return found


def kind_limit(settings: Settings, kind: Kind) -> int:
    return int(getattr(settings, kind.limit_setting))


def max_upload_limit(settings: Settings) -> int:
    """The largest limit of any kind (the multipart pre-check, before the purpose field is read)."""
    return max(kind_limit(settings, k) for k in KINDS.values())


def limit_text(limit: int) -> str:
    if limit >= MIB:
        mb = limit / MIB
        return f"{mb:.0f} MB" if mb == int(mb) else f"{mb:.1f} MB"
    return f"{max(1, limit // 1024)} KB"


def too_large(limit: int) -> ApiError:
    return ApiError(413, "too_large", f"The file is over the {limit_text(limit)} limit.", extra={"limit": limit})


def _magic(head: bytes) -> str | None:
    for magic, mime in MAGIC:
        if head.startswith(magic):
            return mime
    return None


def is_xlsx(source: Path | bytes) -> bool:
    """A zip that lists the parts of an Excel workbook. A plain zip (or a .docx) is not one."""
    try:
        if isinstance(source, bytes | bytearray):
            import io

            zf = zipfile.ZipFile(io.BytesIO(bytes(source)))
        else:
            zf = zipfile.ZipFile(source)
        with zf:
            names = set(zf.namelist())
    except (zipfile.BadZipFile, zipfile.LargeZipFile, OSError, ValueError, EOFError, KeyError):
        return False
    return all(part in names for part in XLSX_PARTS)


def is_csv_text(sample: bytes, *, complete: bool) -> bool:
    """UTF-8 text (a byte order mark is allowed) with no NUL and no control character other than tab, CR, LF.

    `complete` is False when the sample is a cut from a longer file: a character split at the end is allowed.
    """
    if not sample or b"\x00" in sample:
        return False
    if sample.startswith(codecs.BOM_UTF8):
        sample = sample[len(codecs.BOM_UTF8) :]
    try:
        text = codecs.getincrementaldecoder("utf-8")("strict").decode(sample, final=complete)
    except UnicodeDecodeError:
        return False
    return not any(ord(c) < 0x20 and c not in _CSV_ALLOWED_C0 for c in text)


def _read_sample(path: Path) -> tuple[bytes, bool]:
    with Path(path).open("rb") as f:
        sample = f.read(CSV_SAMPLE + 1)
    return sample[:CSV_SAMPLE], len(sample) <= CSV_SAMPLE


def sniff(data: bytes, path: Path | None = None, *, purpose: Any = None) -> str | None:
    """The mime type from the content, or None when it is not a type Packet Station stores.

    `data` is the first bytes (or all of them). PDF, PNG and JPEG go by their magic. A zip is XLSX only when it lists
    the workbook parts: the file at `path` is opened, or `data` when no path is given. CSV has no magic: it is
    considered only for purpose `bom`, from the first 64 KB of `path` (or of `data`).
    """
    mime = _magic(data)
    if mime == XLSX:
        return XLSX if is_xlsx(path if path is not None else data) else None
    if mime is not None:
        return mime
    if purpose is None or kind_of(purpose).name != "bom":
        return None
    if path is not None:
        sample, complete = _read_sample(path)
    else:
        sample, complete = data[:CSV_SAMPLE], len(data) <= CSV_SAMPLE
    return CSV if is_csv_text(sample, complete=complete) else 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 _type_error(kind: Kind) -> ApiError:
    return ApiError(415, "unsupported_type", kind.wrong_type, extra={"allowed": sorted(EXT[m] for m in kind.mimes)})


def _content_type(path: Path, head: bytes, kind: Kind) -> str | None:
    """The mime of a complete file for this kind, or None when it is not one of the kind's types."""
    mime = sniff(head, path, purpose=kind.name if CSV in kind.mimes else None)
    return mime if mime in kind.mimes else None


def _audit_rejected(
    conn: Any, err: ApiError, *, user: Any, actor: Any, packet_id: int | None, size: int, mime: str | None
) -> None:
    """The audit row for an upload the virus scan refused. A failure here never hides the refusal."""
    try:
        audit.append(
            conn,
            action="asset.rejected",
            user_id=user,
            session_id=getattr(actor, "session_id", None),
            role_snapshot=getattr(actor, "role_snapshot", None),
            subject_type="asset",
            packet_id=packet_id,
            detail={"reason": "infected", "signature": err.extra.get("signature"), "size": size, "mime": mime},
            ip=getattr(actor, "ip", None),
        )
    except Exception:
        log.exception("could not write the asset.rejected audit row")


def _scan(
    conn: Any,
    settings: Settings,
    path: Path,
    *,
    user: Any,
    actor: Any,
    packet_id: int | None,
    size: int,
    mime: str | None,
) -> None:
    try:
        scan.check(settings, path)
    except ApiError as e:
        if e.code == "infected":
            _audit_rejected(conn, e, user=user, actor=actor, packet_id=packet_id, size=size, mime=mime)
        raise


def _with_purpose(source_ref: Any, kind: Kind) -> Any:
    if source_ref is None:
        return {"purpose": kind.name}
    if isinstance(source_ref, dict):
        return {"purpose": kind.name, **source_ref}
    return source_ref


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,
    purpose: Any = None,
) -> dict[str, Any]:
    """Store one upload and return its public row.

    `purpose` picks the kind (see KINDS; missing is `document`: PDF, PNG or JPEG up to UPLOAD_MAX_BYTES). The type
    is decided by content. When ClamAV is configured the complete file is scanned before the dedupe lookup and the
    rename. Raises 422 invalid_request (unknown purpose), 413 too_large, 415 unsupported_type, 422 bad_pdf,
    422 infected and 503 scan_unavailable. `user` is the uploader's user id; `actor` (the session Principal) is used
    for the audit row when given. `source_ref` records the kind as `{"purpose": kind}`.
    """
    kind = kind_of(purpose)
    root = Path(settings.assets_dir)
    root.mkdir(parents=True, exist_ok=True)
    limit = kind_limit(settings, kind)
    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""
        checked = False
        with os.fdopen(fd, "wb") as out:
            while True:
                chunk = stream.read(CHUNK)
                if not chunk:
                    break
                size += len(chunk)
                if size > limit:
                    raise too_large(limit)
                if not checked and len(head) < HEAD:
                    head += chunk[: HEAD - len(head)]
                    if len(head) >= HEAD:
                        checked = True
                        early = _magic(head)
                        # Refuse a wrong type before the rest is read. CSV has no magic, so it waits for the end.
                        if early not in kind.mimes and not (early is None and CSV in kind.mimes):
                            raise _type_error(kind)
                digest.update(chunk)
                out.write(chunk)
            out.flush()
            os.fsync(out.fileno())
        if size == 0:
            raise _type_error(kind)
        mime = _content_type(tmp, head, kind)
        if mime is None:
            raise _type_error(kind)
        _scan(conn, settings, tmp, user=user, actor=actor, packet_id=packet_id, size=size, mime=mime)
        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)
        ref = _with_purpose(source_ref, kind)
        with conn.transaction():
            row = _insert(conn, (aid, sha, size, mime, name, key, 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(), "purpose": kind.name},
                ip=getattr(actor, "ip", None),
            )
        placed = None  # the row is written: the file stays
    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()
    quota_after_store(conn, settings)
    return public(dict(row))


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 content must match `mime`, it is fsynced, and it is
    scanned when ClamAV is configured, before the dedupe lookup and the rename. The row carries one extra key,
    `reused`. No audit row for a stored file (the caller writes its own); a virus hit writes `asset.rejected`.
    Raises 413 too_large, 415 unsupported_type, 422 bad_pdf, 422 infected, 503 scan_unavailable.
    """
    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 too_large(limit)
                digest.update(chunk)
            os.fsync(f.fileno())
        found = sniff(head, path, purpose="bom" if mime == CSV else None) if size else None
        if found != mime:
            raise ApiError(415, "unsupported_type", f"The file is not {EXT.get(mime, mime)}.")
        _scan(conn, settings, path, user=uploaded_by, actor=None, packet_id=packet_id, size=size, 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
    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()
    quota_after_store(conn, settings)
    return {**row, "reused": False}


# ---------------------------------------------------------------- quota

QUOTA_KIND = "assets_quota"
QUOTA_SUBJECT = "assets"


def _stored_bytes(conn: Any) -> int:
    r = conn.execute("SELECT coalesce(sum(size), 0) FROM assets").fetchone()
    return int(r[0]) if r else 0


def volume_usage(settings: Settings, conn: Any = None) -> dict[str, Any]:
    """How full the assets store is: {"total", "used", "free", "pct", "warn"} in bytes and percent.

    With ASSETS_QUOTA_BYTES at 0 this is the assets volume (shutil.disk_usage). Otherwise the total is the quota
    and used is the sum of assets.size (read through `conn`; without one, the bytes on disk under ASSETS_DIR).
    `warn` is pct at or over ASSETS_QUOTA_WARN_PCT.
    """
    root = Path(settings.assets_dir)
    probe = root if root.exists() else root.parent
    quota = int(settings.assets_quota_bytes)
    if quota <= 0:
        du = shutil.disk_usage(probe)
        total, used, free = int(du.total), int(du.used), int(du.free)
    else:
        used = _stored_bytes(conn) if conn is not None else _bytes_on_disk(root)
        total = quota
        free = max(0, quota - used)
    pct = round(used * 100.0 / total, 1) if total > 0 else 0.0
    return {"total": total, "used": used, "free": free, "pct": pct, "warn": pct >= int(settings.assets_quota_warn_pct)}


def _bytes_on_disk(root: Path) -> int:
    n = 0
    for dirpath, _dirs, files in os.walk(root):
        for f in files:
            with contextlib.suppress(OSError):
                n += os.stat(os.path.join(dirpath, f)).st_size
    return n


def usage_for(conn: Any, settings: Settings) -> dict[str, Any]:
    """volume_usage, given the connection only when the quota is a byte count (the sum of assets.size needs it)."""
    if int(settings.assets_quota_bytes) > 0:
        return volume_usage(settings, conn)
    return volume_usage(settings)


def open_quota_alert(conn: Any) -> dict[str, Any] | None:
    with conn.cursor(row_factory=dict_row) as cur:
        r = cur.execute(
            "SELECT * FROM notification_events WHERE kind = %s AND coalesce(subject, '') = %s"
            " AND resolved_at IS NULL ORDER BY id LIMIT 1",
            (QUOTA_KIND, QUOTA_SUBJECT),
        ).fetchone()
    return dict(r) if r else None


def check_quota(conn: Any, settings: Settings) -> dict[str, Any]:
    """Open or resolve the storage alert. Returns the usage with `alert`: "opened", "resolved" or None.

    Over the threshold with no open `assets_quota` row: insert one (severity warning, subject assets, detail the
    usage) and audit `storage.quota_warning`. Under it with one open: set its resolved_at.
    """
    usage = usage_for(conn, settings)
    change: str | None = None
    with conn.transaction():
        open_row = open_quota_alert(conn)
        if usage["warn"] and open_row is None:
            r = conn.execute(
                "INSERT INTO notification_events (kind, severity, subject, detail) VALUES (%s, 'warning', %s, %s)"
                " ON CONFLICT (kind, coalesce(subject, '')) WHERE resolved_at IS NULL DO NOTHING RETURNING id",
                (QUOTA_KIND, QUOTA_SUBJECT, Jsonb(usage)),
            ).fetchone()
            if r is not None:
                change = "opened"
                audit.append(
                    conn,
                    action="storage.quota_warning",
                    subject_type="notification",
                    subject_id=str(r[0]),
                    detail={**usage, "warn_pct": int(settings.assets_quota_warn_pct)},
                )
                log.warning(
                    "assets storage is %.1f%% full (warning at %s%%)", usage["pct"], settings.assets_quota_warn_pct
                )
        elif not usage["warn"] and open_row is not None:
            conn.execute("UPDATE notification_events SET resolved_at = now() WHERE id = %s", (open_row["id"],))
            change = "resolved"
            log.info("assets storage is back under the warning: %.1f%%", usage["pct"])
    return {**usage, "alert": change}


def quota_after_store(conn: Any, settings: Settings) -> None:
    """check_quota after a store. A failure is logged and never fails the store."""
    try:
        check_quota(conn, settings)
    except Exception:
        log.exception("the storage quota check failed")


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


def disposition(row: dict[str, Any], download: bool = False) -> str:
    """inline for PDF and images unless a download is asked for; XLSX and CSV are always attachment."""
    kind = "attachment" if download or str(row["mime"]) in ATTACHMENT_ONLY else "inline"
    return f'{kind}; filename="{download_name(row)}"'