ImpTune/imptune/api/packages.py

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

from fastapi import APIRouter
from fastapi.responses import PlainTextResponse, Response

import imptune.config as cfg
from imptune.db.models import Icon, 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")


def _get_printer_and_driver(printer_id: int):
    """Fetch printer and validate driver — returns (printer, driver, driver_name) or PlainTextResponse error."""
    printer = Printer.get_or_none(Printer.id == printer_id)
    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))


@router.get("/{printer_id}/packages/ninja")
def get_ninja_package(printer_id: int):
    """Download a NinjaRMM-ready ZIP containing install.ps1 and the driver files."""
    result, error = _get_printer_and_driver(printer_id)
    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 ZIP in-memory
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        # Add install script
        zf.writestr(f"{safe_name}/install.ps1", install_script)

        # Extract and re-add driver files from driver ZIP
        with zipfile.ZipFile(driver_zip_path, "r") as driver_zf:
            for member in driver_zf.namelist():
                member_data = driver_zf.read(member)
                zf.writestr(f"{safe_name}/drivers/{member}", member_data)

    return Response(
        content=buf.getvalue(),
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{safe_name}_ninja.zip"'},
    )


@router.get("/{printer_id}/packages/intunewin")
def get_intunewin_package(printer_id: int):
    """Download a Microsoft Intune .intunewin deployment package."""
    result, error = _get_printer_and_driver(printer_id)
    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)

    with tempfile.TemporaryDirectory(prefix="imptune_") as tmpdir:
        # Write scripts
        with open(os.path.join(tmpdir, "install.ps1"), "w", encoding="utf-8") as f:
            f.write(install_script)
        with open(os.path.join(tmpdir, "uninstall.ps1"), "w", encoding="utf-8") as f:
            f.write(uninstall_script)
        with open(os.path.join(tmpdir, "detect.ps1"), "w", encoding="utf-8") as f:
            f.write(detect_script)

        # Extract driver ZIP contents into tmpdir/drivers/
        drivers_subdir = os.path.join(tmpdir, "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(tmpdir, "icon.png"))

        # Build .intunewin
        output_path = os.path.join(tmpdir, "out.intunewin")
        build_intunewin(tmpdir, "install.ps1", output_path)

        with open(output_path, "rb") as f:
            content = f.read()

    return Response(
        content=content,
        media_type="application/octet-stream",
        headers={"Content-Disposition": f'attachment; filename="{safe_name}.intunewin"'},
    )