PacketStation/tests/server/test_upload_pipeline.py

"""Phase 4 upload hardening: content sniffing (PDF, PNG, JPEG, XLSX, CSV), the purpose and type matrix, the limit
per kind, safe names, the blob headers, the optional ClamAV scan (with tests/server/fakeclam.py), the storage quota
alert, and the admin storage and alerts endpoints."""

from __future__ import annotations

import io
import struct
import zipfile
import zlib
from collections.abc import Iterator
from pathlib import Path
from typing import Any

import pytest
from fakeclam import EICAR, SIGNATURE, FakeClam
from fakedir import FakeDirectory
from fastapi.testclient import TestClient
from pkthelpers import as_user, phase2_client, seed_users, tiny_pdf
from pshelpers import PG_URL, app_conn, make_settings

from packetstation.api.errors import ApiError
from packetstation.domain import assets, scan
from packetstation.settings import Settings

needs_pg = pytest.mark.skipif(not PG_URL, reason="TEST_PG_SUPERUSER_URL not set (no PostgreSQL)")

XLSX_MIME = "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
WARN = {"total": 1000, "used": 900, "free": 100, "pct": 90.0, "warn": True}
OK = {"total": 1000, "used": 100, "free": 900, "pct": 10.0, "warn": False}


# ---------------------------------------------------------------- sample files


def png_bytes(extra: bytes = b"") -> bytes:
    def chunk(kind: bytes, data: bytes) -> bytes:
        return struct.pack(">I", len(data)) + kind + data + struct.pack(">I", zlib.crc32(kind + data) & 0xFFFFFFFF)

    ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)
    body = chunk(b"IHDR", ihdr) + chunk(b"IDAT", zlib.compress(b"\x00\x10\x20\x30"))
    if extra:
        body += chunk(b"tEXt", b"Comment\x00" + extra)
    return b"\x89PNG\r\n\x1a\n" + body + chunk(b"IEND", b"")


def jpeg_bytes() -> bytes:
    return b"\xff\xd8\xff\xe0\x00\x10JFIF\x00\x01\x01\x00\x00\x01\x00\x01\x00\x00" + b"\x00" * 32 + b"\xff\xd9"


def xlsx_bytes() -> bytes:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as z:
        z.writestr("[Content_Types].xml", '<?xml version="1.0"?><Types/>')
        z.writestr("xl/workbook.xml", '<?xml version="1.0"?><workbook/>')
        z.writestr("xl/worksheets/sheet1.xml", '<?xml version="1.0"?><worksheet/>')
    return buf.getvalue()


def zip_bytes() -> bytes:
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w") as z:
        z.writestr("readme.txt", "just a zip")
    return buf.getvalue()


CSV = b"item,qty,part\r\n1,4,PIPE-4\r\n2,2,FLG-4\r\n"
CSV_BOM = b"\xef\xbb\xbf" + CSV


# ---------------------------------------------------------------- sniffing (no database)


def test_sniff_by_content() -> None:
    assert assets.sniff(tiny_pdf()) == "application/pdf"
    assert assets.sniff(png_bytes()) == "image/png"
    assert assets.sniff(jpeg_bytes()) == "image/jpeg"
    assert assets.sniff(xlsx_bytes()) == XLSX_MIME
    assert assets.sniff(zip_bytes()) is None  # a plain zip is not a workbook
    assert assets.sniff(b"just text") is None


def test_sniff_xlsx_from_a_path(tmp_path: Path) -> None:
    good = tmp_path / "a.bin"
    good.write_bytes(xlsx_bytes())
    plain = tmp_path / "b.bin"
    plain.write_bytes(zip_bytes())
    assert assets.sniff(good.read_bytes()[:16], good) == XLSX_MIME
    assert assets.sniff(plain.read_bytes()[:16], plain) is None


def test_sniff_csv_only_for_a_bom() -> None:
    assert assets.sniff(CSV, purpose="bom") == "text/csv"
    assert assets.sniff(CSV_BOM, purpose="bom") == "text/csv"
    assert assets.sniff(b"a\tb\n1\t2\n", purpose="bom") == "text/csv"
    assert assets.sniff(CSV) is None and assets.sniff(CSV, purpose="document") is None
    assert assets.sniff(b"a,b\n\x00\x01,2\n", purpose="bom") is None
    assert assets.sniff(b"a,b\n\x07,2\n", purpose="bom") is None
    assert assets.sniff(b"a,b\n\xff\xfe,2\n", purpose="bom") is None  # not UTF-8


def test_safe_name_and_download_name() -> None:
    name = assets.safe_name("..\\..\\x<y>.pdf")
    assert name is not None and name.endswith(".pdf")
    assert not any(ch in name for ch in '\\/<>:"|?*') and ".." not in name
    row = {"id": "0f0e", "mime": XLSX_MIME, "original_name": None}
    assert assets.download_name(row).endswith(".xlsx")
    assert assets.download_name({"id": "0f0e", "mime": "text/csv", "original_name": None}).endswith(".csv")
    assert assets.disposition({"id": "1", "mime": XLSX_MIME, "original_name": "b.xlsx"}).startswith("attachment;")
    assert assets.disposition({"id": "1", "mime": "text/csv", "original_name": "b.csv"}).startswith("attachment;")
    assert assets.disposition({"id": "1", "mime": "application/pdf", "original_name": "d.pdf"}).startswith("inline;")


# ---------------------------------------------------------------- the scanner (no database)


def test_fake_clamd_and_scanner(tmp_path: Path) -> None:
    clean = tmp_path / "clean.png"
    clean.write_bytes(png_bytes())
    bad = tmp_path / "bad.png"
    bad.write_bytes(png_bytes(EICAR))
    with FakeClam() as clam:
        s = scan.ClamdScanner(clam.host, clam.port, timeout=5)
        assert s.ping() is True
        assert s.scan_path(clean) == scan.ScanResult(clean=True, signature=None)
        assert s.scan_path(bad) == scan.ScanResult(clean=False, signature=SIGNATURE)
        assert clam.commands == ["PING", "INSTREAM", "INSTREAM"]
        assert clam.scanned[0] == clean.read_bytes()
    with FakeClam(refuse=True) as down:
        s = scan.ClamdScanner(down.host, down.port, timeout=2)
        assert s.ping() is False
        with pytest.raises(scan.ScanError):
            s.scan_path(clean)
    with FakeClam(mode="garbage") as odd:
        with pytest.raises(scan.ScanError):
            scan.ClamdScanner(odd.host, odd.port, timeout=2).scan_path(clean)


def test_scan_check(tmp_path: Path) -> None:
    clean = tmp_path / "clean.png"
    clean.write_bytes(png_bytes())
    bad = tmp_path / "bad.png"
    bad.write_bytes(png_bytes(EICAR))
    off = make_settings(tmp_path)
    assert off.clamav_host == "" and off.clamav_port == 3310 and off.clamav_timeout_seconds == 30
    assert scan.scanner_for(off) is None and scan.check(off, bad) is None
    with FakeClam() as clam:
        on = off.model_copy(update={"clamav_host": clam.host, "clamav_port": clam.port})
        assert isinstance(scan.scanner_for(on), scan.ClamdScanner)
        assert scan.check(on, clean) == scan.ScanResult(clean=True, signature=None)
        with pytest.raises(ApiError) as e:
            scan.check(on, bad)
        assert e.value.status == 422 and e.value.code == "infected"
        assert e.value.detail == "That file failed the virus scan and was not stored."
        assert e.value.extra == {"signature": SIGNATURE}
    with FakeClam(refuse=True) as down:
        dead = off.model_copy(update={"clamav_host": down.host, "clamav_port": down.port, "clamav_timeout_seconds": 2})
        with pytest.raises(ApiError) as e:
            scan.check(dead, clean)
        assert e.value.status == 503 and e.value.code == "scan_unavailable"
        assert e.value.detail == "The virus scanner is not reachable. Try again in a minute."


# ---------------------------------------------------------------- the API (PostgreSQL)


@pytest.fixture
def pg_client2(pg_settings: Settings, fake_dir: FakeDirectory, dev_visual: None) -> Iterator[TestClient]:
    with phase2_client(pg_settings, fake_dir) as client:
        seed_users(client)
        as_user(client, "clerk.clara")
        yield client


def up(
    c: TestClient, data: bytes, name: str, purpose: str | None = None, mime: str = "application/octet-stream"
) -> Any:
    form = {} if purpose is None else {"purpose": purpose}
    return c.post("/api/assets", files={"file": (name, data, mime)}, data=form)


def _one(settings: Settings, sql: str, params: Any = ()) -> Any:
    with app_conn(settings) as conn:
        return conn.execute(sql, params).fetchone()


def stored_files(settings: Settings) -> list[Path]:
    root = Path(settings.assets_dir)
    return sorted(p for p in root.rglob("*") if p.is_file()) if root.exists() else []


SAMPLES = {
    "pdf": (tiny_pdf, "d.pdf"),
    "png": (png_bytes, "p.png"),
    "jpeg": (jpeg_bytes, "j.jpg"),
    "xlsx": (xlsx_bytes, "b.xlsx"),
    "csv": (lambda: CSV, "b.csv"),
    "zip": (zip_bytes, "z.zip"),
}
# purpose -> what is accepted (everything else in SAMPLES is 415)
ACCEPTS = {
    "drawing": {"pdf"},
    "mtr": {"pdf", "png", "jpeg"},
    "document": {"pdf", "png", "jpeg"},
    None: {"pdf", "png", "jpeg"},
    "bom": {"xlsx", "csv"},
}
WRONG_TYPE = {
    "drawing": "A drawing is a PDF file.",
    "mtr": "An MTR is a PDF, PNG or JPG file.",
    "bom": "A BOM is an Excel (.xlsx) or CSV file.",
}


@needs_pg
@pytest.mark.parametrize("purpose", ["drawing", "mtr", "document", None, "bom"])
def test_purpose_and_type_matrix(pg_client2: TestClient, pg_settings: Settings, purpose: str | None) -> None:
    for kind, (make, name) in SAMPLES.items():
        r = up(pg_client2, make(), name, purpose)
        if kind in ACCEPTS[purpose]:
            assert r.status_code == 201, (purpose, kind, r.text)
            ref = _one(pg_settings, "SELECT source_ref FROM assets WHERE id = %s", (r.json()["id"],))[0]
            # the same bytes may have been stored before under another purpose (dedupe returns that row)
            assert ref is None or "purpose" in ref, ref
            assert r.json()["mime"] in ("application/pdf", "image/png", "image/jpeg", XLSX_MIME, "text/csv")
        else:
            assert r.status_code == 415 and r.json()["code"] == "unsupported_type", (purpose, kind, r.text)
            if purpose in WRONG_TYPE:
                assert r.json()["detail"] == WRONG_TYPE[purpose]
            assert isinstance(r.json()["allowed"], list) and r.json()["allowed"]


@needs_pg
def test_purpose_recorded_and_unknown_refused(pg_client2: TestClient, pg_settings: Settings) -> None:
    r = up(pg_client2, png_bytes(b"mtr one"), "m.png", "mtr")
    assert r.status_code == 201
    ref = _one(pg_settings, "SELECT source_ref FROM assets WHERE id = %s", (r.json()["id"],))[0]
    assert ref == {"purpose": "mtr"}
    r = up(pg_client2, png_bytes(b"doc one"), "d.png")
    ref = _one(pg_settings, "SELECT source_ref FROM assets WHERE id = %s", (r.json()["id"],))[0]
    assert ref == {"purpose": "document"}
    r = up(pg_client2, xlsx_bytes(), "b.xlsx", "bom")
    assert r.status_code == 201 and r.json()["mime"] == XLSX_MIME
    r = up(pg_client2, CSV_BOM, "b.csv", "bom")
    assert r.status_code == 201 and r.json()["mime"] == "text/csv"
    r = up(pg_client2, png_bytes(b"doc alias"), "a.png", "doc")
    assert r.status_code == 201, r.text
    ref = _one(pg_settings, "SELECT source_ref FROM assets WHERE id = %s", (r.json()["id"],))[0]
    assert ref == {"purpose": "document"}
    bad = up(pg_client2, tiny_pdf(), "d.pdf", "selfie")
    assert bad.status_code == 422 and bad.json()["code"] == "invalid_request", bad.text
    nul = up(pg_client2, b"a,b\n\x00,1\n", "n.csv", "bom")
    assert nul.status_code == 415 and nul.json()["code"] == "unsupported_type"


@needs_pg
def test_limits_per_kind(pg_settings: Settings, fake_dir: FakeDirectory, dev_visual: None) -> None:
    small = pg_settings.model_copy(update={"upload_max_bytes": 4096, "upload_max_bytes_bom": 1024})
    with phase2_client(small, fake_dir) as c:
        seed_users(c, ("clerk.clara",))
        as_user(c, "clerk.clara")
        big_png = png_bytes(b"x" * 2000)
        assert 1024 < len(big_png) < 4096
        assert up(c, big_png, "p.png", "mtr").status_code == 201
        csv = b"item,qty\n" + b"1,4\n" * 400
        assert 1024 < len(csv) < 4096
        r = up(c, csv, "b.csv", "bom")
        assert r.status_code == 413 and r.json()["code"] == "too_large", r.text
        assert r.json()["limit"] == 1024 and "limit" in r.json()["detail"]
        assert up(c, b"item,qty\n1,4\n", "s.csv", "bom").status_code == 201
        huge = png_bytes(b"y" * 6000)
        r = up(c, huge, "h.png", "document")
        assert r.status_code == 413 and r.json()["code"] == "too_large" and r.json()["limit"] == 4096
        r = up(c, b"%PDF-1.4\n" + b"0" * 200_000, "h.pdf", "drawing")  # over the pre-check too
        assert r.status_code == 413 and r.json()["code"] == "too_large"
        assert not [p for p in stored_files(small) if p.name.startswith(".up-")]


@needs_pg
def test_names_and_blob_headers(pg_client2: TestClient) -> None:
    r = up(pg_client2, tiny_pdf(2), "..\\..\\x<y>.pdf", "drawing", "application/pdf")
    assert r.status_code == 201, r.text
    a = r.json()
    name = a["original_name"]
    assert name and name.endswith(".pdf") and not any(ch in name for ch in "\\/<>") and not name.startswith(".")
    b = pg_client2.get(a["url"])
    assert b.status_code == 200 and b.headers["content-disposition"].startswith("inline;")
    assert b.headers["x-content-type-options"] == "nosniff"
    x = up(pg_client2, xlsx_bytes(), "bom list.xlsx", "bom").json()
    b = pg_client2.get(x["url"])
    assert b.status_code == 200 and b.headers["content-disposition"] == 'attachment; filename="bom list.xlsx"'
    assert b.headers["x-content-type-options"] == "nosniff" and b.headers["content-type"].startswith(XLSX_MIME)
    cv = up(pg_client2, CSV, "...", "bom").json()  # no usable name: the id and the extension
    b = pg_client2.get(cv["url"])
    assert b.headers["content-disposition"].startswith("attachment;")
    assert b.headers["content-disposition"].endswith('.csv"') and b.headers["x-content-type-options"] == "nosniff"


# ---------------------------------------------------------------- ClamAV through the API


@pytest.fixture
def clam() -> Iterator[FakeClam]:
    with FakeClam() as fake:
        yield fake


@needs_pg
def test_scan_on_upload(pg_settings: Settings, fake_dir: FakeDirectory, dev_visual: None, clam: FakeClam) -> None:
    on = pg_settings.model_copy(update={"clamav_host": clam.host, "clamav_port": clam.port})
    with phase2_client(on, fake_dir) as c:
        seed_users(c, ("clerk.clara",))
        as_user(c, "clerk.clara")
        clean = png_bytes(b"clean one")
        r = up(c, clean, "ok.png", "mtr")
        assert r.status_code == 201, r.text
        assert "INSTREAM" in clam.commands and clean in clam.scanned
        before = stored_files(on)
        bad = png_bytes(EICAR)
        r = up(c, bad, "bad.png", "mtr")
        assert r.status_code == 422 and r.json()["code"] == "infected", r.text
        assert r.json()["detail"] == "That file failed the virus scan and was not stored."
        assert r.json()["signature"] == SIGNATURE
        assert stored_files(on) == before  # nothing placed, no temp file left
        assert _one(on, "SELECT count(*) FROM assets WHERE size = %s", (len(bad),))[0] == 0
        au = _one(on, "SELECT detail, user_id FROM audit_events WHERE action = 'asset.rejected' ORDER BY id DESC")
        assert au is not None
        assert au[0]["reason"] == "infected" and au[0]["signature"] == SIGNATURE
        assert au[0]["size"] == len(bad) and au[0]["mime"] == "image/png" and au[1] is not None
        # The same infected bytes are never deduped against a stored row: a second try is refused again
        assert up(c, bad, "bad2.png", "document").status_code == 422


@needs_pg
def test_scanner_down_fails_closed(pg_settings: Settings, fake_dir: FakeDirectory, dev_visual: None) -> None:
    with FakeClam(refuse=True) as down:
        dead = pg_settings.model_copy(
            update={"clamav_host": down.host, "clamav_port": down.port, "clamav_timeout_seconds": 2}
        )
        with phase2_client(dead, fake_dir) as c:
            seed_users(c, ("clerk.clara",))
            as_user(c, "clerk.clara")
            r = up(c, png_bytes(b"down"), "p.png", "mtr")
            assert r.status_code == 503 and r.json()["code"] == "scan_unavailable", r.text
            assert r.json()["detail"] == "The virus scanner is not reachable. Try again in a minute."
            assert stored_files(dead) == []


@needs_pg
def test_no_scanner_means_no_call(pg_client2: TestClient, clam: FakeClam) -> None:
    r = up(pg_client2, png_bytes(EICAR), "p.png", "mtr")
    assert r.status_code == 201, r.text  # ClamAV is off: nothing scans
    assert clam.calls == 0


@needs_pg
def test_store_file_is_scanned(pg_settings: Settings, clam: FakeClam) -> None:
    on = pg_settings.model_copy(update={"clamav_host": clam.host, "clamav_port": clam.port})
    tmp = Path(on.assets_dir) / "tmp"
    tmp.mkdir(parents=True, exist_ok=True)
    bad = tmp / "dl-bad.png"
    bad.write_bytes(png_bytes(EICAR))
    with app_conn(on) as conn:
        with pytest.raises(ApiError) as e:
            assets.store_file(conn, on, bad, mime="image/png", original_name="m.png", source="mfiles")
        assert e.value.code == "infected"
        assert not bad.exists()
        assert conn.execute("SELECT count(*) FROM assets").fetchone()[0] == 0
        assert conn.execute("SELECT count(*) FROM audit_events WHERE action = 'asset.rejected'").fetchone()[0] == 1
        good = tmp / "dl-good.pdf"
        good.write_bytes(tiny_pdf())
        row = assets.store_file(conn, on, good, original_name="d.pdf", source="mfiles")
        assert row["reused"] is False and row["mime"] == "application/pdf"
    assert clam.commands.count("INSTREAM") == 2
    assert [p for p in stored_files(on) if p.parent.name != "tmp"]


# ---------------------------------------------------------------- the quota alert


@needs_pg
def test_quota_alert_opens_and_resolves(
    pg_client2: TestClient, pg_settings: Settings, monkeypatch: pytest.MonkeyPatch
) -> None:
    monkeypatch.setattr(assets, "volume_usage", lambda *a, **k: dict(WARN))
    assert up(pg_client2, png_bytes(b"q1"), "q1.png", "mtr").status_code == 201
    row = _one(
        pg_settings,
        "SELECT id, kind, severity, subject, detail, resolved_at FROM notification_events WHERE kind = 'assets_quota'",
    )
    assert row is not None and row[1:4] == ("assets_quota", "warning", "assets") and row[5] is None
    assert row[4]["pct"] == 90.0 and row[4]["warn"] is True
    assert _one(pg_settings, "SELECT count(*) FROM audit_events WHERE action = 'storage.quota_warning'")[0] == 1
    # A second upload while it is open adds nothing
    assert up(pg_client2, png_bytes(b"q2"), "q2.png", "mtr").status_code == 201
    assert _one(pg_settings, "SELECT count(*) FROM notification_events WHERE resolved_at IS NULL")[0] == 1
    assert _one(pg_settings, "SELECT count(*) FROM audit_events WHERE action = 'storage.quota_warning'")[0] == 1

    as_user(pg_client2, "admin.ann")
    st = pg_client2.get("/api/admin/storage")
    assert st.status_code == 200, st.text
    body = st.json()
    assert body["usage"] == WARN and body["quota_warn_pct"] == 85
    assert body["alert"]["id"] == row[0] and body["alert"]["acknowledged_at"] is None
    assert body["clamav"]["configured"] is False and body["clamav"]["reachable"] is None
    assert body["clamav"]["host"] is None
    al = pg_client2.get("/api/admin/alerts")
    assert al.status_code == 200 and [a["id"] for a in al.json()["alerts"]] == [row[0]]
    ack = pg_client2.post(f"/api/admin/alerts/{row[0]}/ack")
    assert ack.status_code == 200, ack.text
    assert ack.json()["acknowledged_at"] and ack.json()["acknowledged_by"]["display_name"] == "Ann Admin"
    assert _one(pg_settings, "SELECT count(*) FROM audit_events WHERE action = 'alert.acknowledged'")[0] == 1
    assert pg_client2.post("/api/admin/alerts/999999/ack").status_code == 404

    monkeypatch.setattr(assets, "volume_usage", lambda *a, **k: dict(OK))
    with app_conn(pg_settings) as conn:
        assets.check_quota(conn, pg_settings)
    resolved = _one(pg_settings, "SELECT resolved_at FROM notification_events WHERE id = %s", (row[0],))[0]
    assert resolved is not None
    assert pg_client2.get("/api/admin/storage").json()["alert"] is None
    assert pg_client2.get("/api/admin/alerts").json()["alerts"] == []


@needs_pg
def test_quota_check_failure_never_fails_the_upload(
    pg_client2: TestClient, pg_settings: Settings, monkeypatch: pytest.MonkeyPatch
) -> None:
    def boom(*a: Any, **k: Any) -> dict[str, Any]:
        raise OSError("disk usage unavailable")

    monkeypatch.setattr(assets, "volume_usage", boom)
    r = up(pg_client2, png_bytes(b"still stored"), "s.png", "mtr")
    assert r.status_code == 201, r.text


@needs_pg
def test_admin_storage_and_alerts_are_admin_only(pg_client2: TestClient) -> None:
    for path in ("/api/admin/storage", "/api/admin/alerts"):
        r = pg_client2.get(path)
        assert r.status_code == 403 and r.json()["code"] == "forbidden", path
    assert pg_client2.post("/api/admin/alerts/1/ack").status_code == 403
    as_user(pg_client2, "admin.ann")
    assert pg_client2.get("/api/admin/storage").status_code == 200
    assert pg_client2.get("/api/admin/alerts").json() == {"alerts": []}


@needs_pg
def test_admin_storage_shows_clamav(
    pg_settings: Settings, fake_dir: FakeDirectory, dev_visual: None, clam: FakeClam
) -> None:
    on = pg_settings.model_copy(update={"clamav_host": clam.host, "clamav_port": clam.port})
    with phase2_client(on, fake_dir) as c:
        seed_users(c, ("admin.ann",))
        as_user(c, "admin.ann")
        cl = c.get("/api/admin/storage").json()["clamav"]
        assert cl["configured"] is True and cl["host"] == clam.host and cl["reachable"] is True
        assert "PING" in clam.commands
    with FakeClam(refuse=True) as down:
        off = pg_settings.model_copy(
            update={"clamav_host": down.host, "clamav_port": down.port, "clamav_timeout_seconds": 2}
        )
        with phase2_client(off, fake_dir) as c:
            as_user(c, "admin.ann")
            cl = c.get("/api/admin/storage").json()["clamav"]
            assert cl["configured"] is True and cl["reachable"] is False