ImpTune/tests/test_packages.py

"""Integration tests for package export endpoints — NinjaRMM ZIP and .intunewin."""
import io
import json
import zipfile

import pytest


# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def driver_zip_bytes():
    """Create a minimal valid driver ZIP with a fake INF and CAT file."""
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("printer.inf", "[Version]\nSignature=$WINDOWS NT$\n")
        zf.writestr("printer.cat", "FAKE_CAT")
    return buf.getvalue()


@pytest.fixture
def setup_printer_with_driver(tmp_data_dir, driver_zip_bytes, owner):
    """Create Driver record (with ZIP on disk) and Printer record linked to it."""
    import hashlib
    import os

    import imptune.config as cfg
    from imptune.db.models import Driver, Printer

    # Write driver ZIP to DRIVERS_DIR
    sha = hashlib.sha256(driver_zip_bytes).hexdigest()
    drivers_dir = cfg.DRIVERS_DIR
    os.makedirs(drivers_dir, exist_ok=True)
    zip_path = os.path.join(drivers_dir, f"{sha}.zip")
    with open(zip_path, "wb") as f:
        f.write(driver_zip_bytes)

    driver = Driver.create(
        sha256=sha,
        original_filename="printer_driver.zip",
        size_bytes=len(driver_zip_bytes),
        driver_desc=json.dumps(["HP LaserJet Pro"]),
        inf_filename="printer.inf",
    )
    printer = Printer.create(
        name="Test Printer",
        ip_address="192.168.1.100",
        port_name="IP_192.168.1.100",
        driver=driver,
        owner=owner,
        duplex_mode="OneSided",
        color_mode=True,
        paper_size="A4",
        collate=True,
    )
    return printer, driver


@pytest.fixture
def printer_no_driver(tmp_data_dir, owner):
    """Create Printer record with no driver assigned."""
    from imptune.db.models import Printer

    return Printer.create(
        name="No Driver Printer",
        ip_address="10.0.0.1",
        port_name="IP_10.0.0.1",
        driver=None,
        owner=owner,
    )


# ---------------------------------------------------------------------------
# NinjaRMM ZIP endpoint
# ---------------------------------------------------------------------------


class TestNinjaDownload:
    def test_returns_zip(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/ninja")
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/zip"
        assert "attachment" in resp.headers["content-disposition"]
        assert ".zip" in resp.headers["content-disposition"]

    def test_zip_contains_install_script(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/ninja")
        assert resp.status_code == 200
        safe_name = printer.name.replace(" ", "_")
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
        assert f"{safe_name}/install.ps1" in names

    def test_zip_contains_driver_files(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/ninja")
        assert resp.status_code == 200
        safe_name = printer.name.replace(" ", "_")
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
        # Driver ZIP contained printer.inf and printer.cat
        assert f"{safe_name}/drivers/printer.inf" in names
        assert f"{safe_name}/drivers/printer.cat" in names

    def test_404_missing_printer(self, client, tmp_data_dir):
        resp = client.get("/printers/9999/packages/ninja")
        assert resp.status_code == 404

    def test_422_no_driver(self, client, printer_no_driver):
        resp = client.get(f"/printers/{printer_no_driver.id}/packages/ninja")
        assert resp.status_code == 422


# ---------------------------------------------------------------------------
# .intunewin endpoint
# ---------------------------------------------------------------------------


class TestIntunewinDownload:
    def test_returns_intunewin(self, client, setup_printer_with_driver):
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/intunewin")
        assert resp.status_code == 200
        assert resp.headers["content-type"] == "application/octet-stream"
        assert "attachment" in resp.headers["content-disposition"]
        assert ".intunewin" in resp.headers["content-disposition"]

    def test_intunewin_is_valid_zip(self, client, setup_printer_with_driver):
        """Outer .intunewin file must be a valid ZIP with IntuneWinPackage/ structure."""
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/intunewin")
        assert resp.status_code == 200
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
        assert any(n.startswith("IntuneWinPackage/") for n in names)

    def test_404_missing_printer(self, client, tmp_data_dir):
        resp = client.get("/printers/9999/packages/intunewin")
        assert resp.status_code == 404

    def test_422_no_driver(self, client, printer_no_driver):
        resp = client.get(f"/printers/{printer_no_driver.id}/packages/intunewin")
        assert resp.status_code == 422


# ---------------------------------------------------------------------------
# Printer detail page — command preview and export links
# ---------------------------------------------------------------------------


class TestCommandPreview:
    def test_detail_page_shows_commands(self, client, setup_printer_with_driver):
        """Printer detail page shows install and uninstall command strings."""
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}")
        assert resp.status_code == 200
        html = resp.text
        assert "install-cmd" in html
        assert "uninstall-cmd" in html
        assert "install.ps1" in html
        assert "uninstall.ps1" in html

    def test_detail_page_shows_export_links(self, client, setup_printer_with_driver):
        """Printer detail page shows NinjaRMM ZIP and .intunewin download links."""
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}")
        assert resp.status_code == 200
        html = resp.text
        assert f"/printers/{printer.id}/packages/ninja" in html
        assert f"/printers/{printer.id}/packages/intunewin" in html

    def test_detail_page_hides_commands_without_driver(self, client, printer_no_driver):
        """Detail page hides command section and export links when no driver assigned."""
        resp = client.get(f"/printers/{printer_no_driver.id}")
        assert resp.status_code == 200
        html = resp.text
        assert "install-cmd" not in html
        assert "packages/ninja" not in html
        assert "packages/intunewin" not in html

    def test_detail_page_shows_icon_upload_form(self, client, setup_printer_with_driver):
        """Printer detail page always shows icon upload form."""
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}")
        assert resp.status_code == 200
        html = resp.text
        assert f"/printers/{printer.id}/icon" in html
        assert "icon-status" in html

    def test_detail_page_shows_script_links(self, client, setup_printer_with_driver):
        """Printer detail page shows 3 direct .ps1 script download links when driver assigned."""
        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}")
        assert resp.status_code == 200
        html = resp.text
        assert f"/printers/{printer.id}/scripts/install.ps1" in html
        assert f"/printers/{printer.id}/scripts/uninstall.ps1" in html
        assert f"/printers/{printer.id}/scripts/detect.ps1" in html


# ---------------------------------------------------------------------------
# Icon inclusion in .intunewin export
# ---------------------------------------------------------------------------


class TestIntunewinIconInclusion:
    def test_intunewin_includes_icon(self, client, setup_printer_with_driver, tmp_data_dir, monkeypatch):
        """When a printer has an uploaded icon, icon.png must appear in the .intunewin staging directory."""
        import io as _io
        import os

        from PIL import Image

        printer, _ = setup_printer_with_driver

        # Upload a 256x256 PNG icon for the printer
        buf = _io.BytesIO()
        Image.new("RGBA", (256, 256), color="red").save(buf, format="PNG")
        buf.seek(0)
        upload_resp = client.post(
            f"/printers/{printer.id}/icon",
            files={"file": ("icon.png", buf, "image/png")},
        )
        assert upload_resp.status_code == 200

        # Monkeypatch build_intunewin to capture staged files and write a fake output
        staged_files: list[str] = []

        def fake_build(source_dir, setup_file, output_path):
            staged_files.extend(os.listdir(source_dir))
            with open(output_path, "wb") as f:
                f.write(b"FAKE")

        monkeypatch.setattr("imptune.api.packages.build_intunewin", fake_build)

        resp = client.get(f"/printers/{printer.id}/packages/intunewin")
        assert resp.status_code == 200
        assert "icon.png" in staged_files

    def test_intunewin_without_icon_succeeds(self, client, setup_printer_with_driver, monkeypatch):
        """When a printer has no icon, .intunewin export must succeed without error."""
        import os

        # Monkeypatch build_intunewin to write a fake output
        def fake_build(source_dir, setup_file, output_path):
            with open(output_path, "wb") as f:
                f.write(b"FAKE")

        monkeypatch.setattr("imptune.api.packages.build_intunewin", fake_build)

        printer, _ = setup_printer_with_driver
        resp = client.get(f"/printers/{printer.id}/packages/intunewin")
        assert resp.status_code == 200


# ---------------------------------------------------------------------------
# Streamed exports — packages are assembled on disk, not buffered in RAM
# ---------------------------------------------------------------------------


class TestStreamedExportCleanup:
    @pytest.fixture
    def bulky_printer(self, tmp_data_dir, owner):
        """Printer whose driver ZIP is large enough to prove the body streamed."""
        import hashlib
        import os

        import imptune.config as cfg
        from imptune.db.models import Driver, Printer

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w") as zf:
            zf.writestr("printer.inf", "[Version]\nSignature=$WINDOWS NT$\n")
            zf.writestr("sub/payload.bin", os.urandom(200_000))
        data = buf.getvalue()

        sha = hashlib.sha256(data).hexdigest()
        os.makedirs(cfg.DRIVERS_DIR, exist_ok=True)
        with open(os.path.join(cfg.DRIVERS_DIR, f"{sha}.zip"), "wb") as f:
            f.write(data)

        driver = Driver.create(
            sha256=sha,
            original_filename="bulky.zip",
            size_bytes=len(data),
            driver_desc=json.dumps(["Acme SuperPrint"]),
            inf_filename="printer.inf",
        )
        return Printer.create(
            name="Bulky Printer",
            ip_address="10.1.1.1",
            port_name="IP_10.1.1.1",
            owner=owner,
            driver=driver,
        )

    @pytest.mark.parametrize("kind", ["ninja", "intunewin"])
    def test_export_streams_and_removes_its_temp_dir(self, client, bulky_printer, kind):
        """The response outlives the handler's temp dir, so a background task
        deletes it — a leak here fills the container's disk one download at a time."""
        import glob
        import os
        import tempfile

        pattern = os.path.join(tempfile.gettempdir(), "imptune_*")
        before = set(glob.glob(pattern))

        resp = client.get(f"/printers/{bulky_printer.id}/packages/{kind}")

        assert resp.status_code == 200
        assert len(resp.content) > 150_000, "payload looks truncated"
        assert set(glob.glob(pattern)) == before

    def test_ninja_zip_preserves_nested_driver_paths(self, client, bulky_printer):
        """Members are copied chunk-wise via ZipFile.open — nested paths must survive."""
        resp = client.get(f"/printers/{bulky_printer.id}/packages/ninja")
        with zipfile.ZipFile(io.BytesIO(resp.content)) as zf:
            names = zf.namelist()
            assert "Bulky_Printer/drivers/sub/payload.bin" in names
            assert zf.getinfo("Bulky_Printer/drivers/sub/payload.bin").file_size == 200_000