tests/test_verify_archive.py

from __future__ import annotations

import csv
import hashlib
import json
from pathlib import Path

import pytest
from verify_archive import verify_archive


def _write_archive(root: Path) -> None:
    for relative, content in {
        "images/example.png": b"image",
        "results/metrics.csv": b"condition,success\npng,1\n",
        "samples/example.png": b"sample",
        "weights/model.png": b"weights",
    }.items():
        path = root / relative
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_bytes(content)
    rows = []
    for relative in ("weights/model.png",):
        path = root / relative
        rows.append(
            {
                "path": relative,
                "bytes": str(path.stat().st_size),
                "sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
            }
        )
    with (root / "weights/manifest.csv").open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=["path", "bytes", "sha256"])
        writer.writeheader()
        writer.writerows(rows)
    manifest = {
        "files": [
            {
                "path": relative,
                "bytes": (root / relative).stat().st_size,
                "sha256": hashlib.sha256((root / relative).read_bytes()).hexdigest(),
            }
            for relative in (
                "images/example.png",
                "results/metrics.csv",
                "samples/example.png",
                "weights/manifest.csv",
            )
        ]
    }
    (root / "results/archive_manifest.json").write_text(json.dumps(manifest), encoding="utf-8")


def _refresh_manifest(root: Path) -> None:
    manifest_path = root / "results/archive_manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    manifest.pop("files", None)
    manifest["artifacts"] = []
    for path in sorted(p for directory in ("images", "results", "samples", "weights")
                       for p in (root / directory).rglob("*") if p.is_file() and p.name != "archive_manifest.json"):
        relative = path.relative_to(root).as_posix()
        manifest["artifacts"].append({"path": relative, "bytes": path.stat().st_size,
                                      "sha256": hashlib.sha256(path.read_bytes()).hexdigest()})
    manifest.update({"profile": "test", "prompt_count": 1, "conditions": ["png_baseline", "webp_lossless"]})
    manifest_path.write_text(json.dumps(manifest), encoding="utf-8")


def _add_generation_rows(root: Path) -> None:
    sample = root / "samples/example.png"
    rows = [
        {"row_type": "generation", "condition": "png_baseline", "prompt_id": "p1", "success": True,
         "sample_path": "samples/example.png", "sample_bytes": sample.stat().st_size,
         "sample_hash": hashlib.sha256(sample.read_bytes()).hexdigest(), "sample_sha256": None},
        {"row_type": "generation", "condition": "webp_lossless", "prompt_id": "p1", "success": False,
         "sample_path": None, "sample_bytes": None, "sample_hash": None, "sample_sha256": None},
    ]
    (root / "results/raw_rows.jsonl").write_text("".join(json.dumps(row) + "\n" for row in rows), encoding="utf-8")
    _refresh_manifest(root)


def test_verify_archive_accepts_hash_and_containment_valid_archive(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)

    report = verify_archive(tmp_path)

    assert report["checked_files"] == 7
    assert report["checked_weight_entries"] == 1
    assert report["generation_rows"] == 2
    assert report["successful_samples"] == 1


def test_verify_archive_rejects_path_escape(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    manifest_path = tmp_path / "results/archive_manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    manifest["artifacts"].append({"path": "../outside", "bytes": 0, "sha256": ""})
    manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

    with pytest.raises(ValueError, match="containment"):
        verify_archive(tmp_path)


@pytest.mark.parametrize("mutation, message", [
    ("missing", "sample"),
    ("hash", "SHA-256"),
    ("bytes", "byte count"),
])
def test_verify_archive_rejects_invalid_success_provenance(tmp_path: Path, mutation: str, message: str) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    raw = tmp_path / "results/raw_rows.jsonl"
    row = json.loads(raw.read_text(encoding="utf-8").splitlines()[0])
    if mutation == "missing":
        row["sample_path"] = None
    elif mutation == "hash":
        row["sample_hash"] = "0" * 64
    else:
        row["sample_bytes"] += 1
    raw.write_text(json.dumps(row) + "\n" + raw.read_text(encoding="utf-8").splitlines()[1] + "\n", encoding="utf-8")
    _refresh_manifest(tmp_path)

    with pytest.raises(ValueError, match=message):
        verify_archive(tmp_path)


def test_verify_archive_rejects_failed_row_with_sample(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    raw = tmp_path / "results/raw_rows.jsonl"
    rows = [json.loads(line) for line in raw.read_text(encoding="utf-8").splitlines()]
    rows[1]["sample_path"] = "samples/example.png"
    raw.write_text("\n".join(json.dumps(row) for row in rows) + "\n", encoding="utf-8")
    _refresh_manifest(tmp_path)

    with pytest.raises(ValueError, match="failed row"):
        verify_archive(tmp_path)


def test_verify_archive_rejects_raw_rows_symlink(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    raw = tmp_path / "results/raw_rows.jsonl"
    raw.unlink()
    raw.symlink_to(tmp_path / "outside.jsonl")
    (tmp_path / "outside.jsonl").write_text("", encoding="utf-8")

    with pytest.raises(ValueError, match="symlink"):
        verify_archive(tmp_path)


def test_verify_archive_rejects_duplicate_manifest_path(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    manifest_path = tmp_path / "results/archive_manifest.json"
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    manifest["artifacts"].append(manifest["artifacts"][0])
    manifest_path.write_text(json.dumps(manifest), encoding="utf-8")

    with pytest.raises(ValueError, match="duplicate archive path"):
        verify_archive(tmp_path)


def test_verify_archive_rejects_listed_hash_tamper(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    path = tmp_path / "samples/example.png"
    path.write_bytes(b"tampeR")

    with pytest.raises(ValueError, match="SHA-256"):
        verify_archive(tmp_path)


def test_verify_archive_rejects_archived_tree_symlink_and_unlisted_file(tmp_path: Path) -> None:
    _write_archive(tmp_path)
    _add_generation_rows(tmp_path)
    extra = tmp_path / "samples/unlisted.bin"
    extra.write_bytes(b"extra")
    with pytest.raises(ValueError, match="unlisted"):
        verify_archive(tmp_path)
    extra.unlink()
    extra.symlink_to(tmp_path / "weights/model.png")

    with pytest.raises(ValueError, match="symlink"):
        verify_archive(tmp_path)