ImpTune/imptune/api/packages.py

"""Package export endpoints — serves deployment packages for NinjaRMM and Microsoft Intune."""
import json
import os
import shutil
import tempfile
import zipfile

from fastapi import APIRouter, Request
from fastapi.responses import FileResponse, PlainTextResponse
from starlette.background import BackgroundTask

import imptune.config as cfg
from imptune.db.models import Icon, Owner, Printer
from imptune.generators.intunewin_builder import build_intunewin
from imptune.generators.script_generator import render_detect, render_install, render_uninstall
from imptune.storage.driver_store import DriverStore

router = APIRouter(prefix="/printers")

# Driver payloads run to ~100 MB. Packages are assembled in a temp dir and
# streamed from disk instead of being held in memory, so N concurrent
# downloads cost N file handles rather than N × package size of RAM.
_CHUNK = 1024 * 1024


def _get_printer_and_driver(printer_id: int, owner: Owner):
    """Fetch printer (scoped to owner) and validate driver — returns (printer, driver, driver_name) or PlainTextResponse error."""
    printer = Printer.get_or_none((Printer.id == printer_id) & (Printer.owner == owner))
    if printer is None:
        return None, PlainTextResponse("Printer not found", status_code=404)

    driver = printer.driver
    if driver is None:
        return None, PlainTextResponse("No driver assigned", status_code=422)

    if not driver.inf_filename:
        return None, PlainTextResponse("Driver has no INF file", status_code=422)

    if not driver.driver_desc:
        return None, PlainTextResponse("Driver has no description", status_code=422)

    try:
        desc_list = json.loads(driver.driver_desc)
        driver_name = desc_list[0]
    except (json.JSONDecodeError, IndexError, TypeError):
        return None, PlainTextResponse("Driver description is invalid", status_code=422)

    return (printer, driver, driver_name), None


def _get_driver_zip_path(driver) -> str:
    return str(DriverStore(cfg.DRIVERS_DIR).get_path(driver.sha256))


def _streamed_download(path: str, tmpdir: str, media_type: str, filename: str) -> FileResponse:
    """Serve a built package off disk, deleting its temp dir once sent."""
    return FileResponse(
        path,
        media_type=media_type,
        filename=filename,
        background=BackgroundTask(shutil.rmtree, tmpdir, True),
    )


@router.get("/{printer_id}/packages/ninja")
def get_ninja_package(request: Request, printer_id: int):
    """Download a NinjaRMM-ready ZIP containing install.ps1 and the driver files."""
    result, error = _get_printer_and_driver(printer_id, request.state.owner)
    if error is not None:
        return error

    printer, driver, driver_name = result

    # Validate driver file exists on disk
    driver_zip_path = _get_driver_zip_path(driver)
    if not os.path.isfile(driver_zip_path):
        return PlainTextResponse("Driver file not found on disk", status_code=422)

    safe_name = printer.name.replace(" ", "_")

    # Render install script
    install_script = render_install(
        printer_name=printer.name,
        ip_address=printer.ip_address,
        port_name=printer.port_name,
        driver_name=driver_name,
        inf_filename=driver.inf_filename,
        duplex_mode=printer.duplex_mode,
        color_mode=printer.color_mode,
        paper_size=printer.paper_size,
        collate=printer.collate,
    )

    # Build the ZIP on disk, copying driver members through in chunks so a
    # 100 MB driver never lands in memory whole.
    tmpdir = tempfile.mkdtemp(prefix="imptune_ninja_")
    try:
        # Fixed on-disk name — the printer name only shapes the download name,
        # so a "/" or ":" in it can't break the temp path.
        out_path = os.path.join(tmpdir, "package.zip")
        with zipfile.ZipFile(out_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
            zf.writestr(f"{safe_name}/install.ps1", install_script)

            with zipfile.ZipFile(driver_zip_path, "r") as driver_zf:
                for member in driver_zf.infolist():
                    target = f"{safe_name}/drivers/{member.filename}"
                    if member.is_dir():
                        zf.writestr(target, b"")
                        continue
                    with driver_zf.open(member) as src, zf.open(target, "w") as dst:
                        shutil.copyfileobj(src, dst, _CHUNK)
    except BaseException:
        shutil.rmtree(tmpdir, ignore_errors=True)
        raise

    return _streamed_download(
        out_path, tmpdir, "application/zip", f"{safe_name}_ninja.zip"
    )


@router.get("/{printer_id}/packages/intunewin")
def get_intunewin_package(request: Request, printer_id: int):
    """Download a Microsoft Intune .intunewin deployment package."""
    result, error = _get_printer_and_driver(printer_id, request.state.owner)
    if error is not None:
        return error

    printer, driver, driver_name = result

    # Validate driver file exists on disk
    driver_zip_path = _get_driver_zip_path(driver)
    if not os.path.isfile(driver_zip_path):
        return PlainTextResponse("Driver file not found on disk", status_code=422)

    safe_name = printer.name.replace(" ", "_")

    # Render all three scripts
    install_script = render_install(
        printer_name=printer.name,
        ip_address=printer.ip_address,
        port_name=printer.port_name,
        driver_name=driver_name,
        inf_filename=driver.inf_filename,
        duplex_mode=printer.duplex_mode,
        color_mode=printer.color_mode,
        paper_size=printer.paper_size,
        collate=printer.collate,
    )
    uninstall_script = render_uninstall(
        printer_name=printer.name,
        driver_name=driver_name,
        port_name=printer.port_name,
    )
    detect_script = render_detect(printer_name=printer.name)

    # The build output is streamed straight off disk, so the temp tree has to
    # outlive this handler — the response's background task removes it.
    tmpdir = tempfile.mkdtemp(prefix="imptune_")
    try:
        staging = os.path.join(tmpdir, "staging")
        os.makedirs(staging, exist_ok=True)

        # Write scripts
        for filename, script in (
            ("install.ps1", install_script),
            ("uninstall.ps1", uninstall_script),
            ("detect.ps1", detect_script),
        ):
            with open(os.path.join(staging, filename), "w", encoding="utf-8") as f:
                f.write(script)

        # Extract driver ZIP contents into staging/drivers/
        drivers_subdir = os.path.join(staging, "drivers")
        os.makedirs(drivers_subdir, exist_ok=True)
        with zipfile.ZipFile(driver_zip_path, "r") as driver_zf:
            driver_zf.extractall(drivers_subdir)

        # Copy icon into staging if one exists for this printer
        icon_record = Icon.get_or_none(Icon.printer == printer.id)
        if icon_record is not None:
            icon_src = os.path.join(cfg.ICONS_DIR, icon_record.sha256)
            if os.path.isfile(icon_src):
                shutil.copy2(icon_src, os.path.join(staging, "icon.png"))

        # Build .intunewin outside the staging dir — an output file written into
        # the tree being packaged would end up inside its own package.
        output_path = os.path.join(tmpdir, "package.intunewin")
        build_intunewin(staging, "install.ps1", output_path)
    except BaseException:
        shutil.rmtree(tmpdir, ignore_errors=True)
        raise

    return _streamed_download(
        output_path, tmpdir, "application/octet-stream", f"{safe_name}.intunewin"
    )