"""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" )