"""Driver API — upload, rename, and driver-library icons."""
from __future__ import annotations
import io
import json
import zipfile
from fastapi import APIRouter, Form, Request, UploadFile
from fastapi.responses import FileResponse, HTMLResponse, Response
import imptune.config as _cfg
from imptune.api.pages import build_driver_data, printer_counts_by_driver
from imptune.db.models import Driver, DriverIcon
from imptune.services import websearch
from imptune.services.icons import (
MAX_ICON_BYTES,
ImageError,
IconTooLarge,
icon_path,
set_driver_icon,
)
from imptune.services.inf_parser import _detect_encoding, parse_inf
from imptune.storage.driver_store import DriverStore
from imptune.templating import templates
router = APIRouter(prefix="/drivers")
MAX_UPLOAD_BYTES = 100 * 1024 * 1024 # 100 MB
MAX_DISPLAY_NAME = 120
def _error_response(message: str, status_code: int = 400) -> HTMLResponse:
"""Return an HTMX-friendly error fragment swapped into #driver-list."""
return HTMLResponse(
content=f"",
status_code=status_code,
)
def _render_driver_list(request: Request, **extra) -> HTMLResponse:
"""Re-render the whole `#driver-list` table — the target every form swaps."""
drivers = Driver.select().order_by(Driver.uploaded_at.desc())
context = {
"driver_data": build_driver_data(drivers),
"usage": printer_counts_by_driver(request.state.owner),
}
context.update(extra)
return templates.TemplateResponse(
request=request, name="partials/driver_list.html", context=context
)
@router.post("/upload", response_class=HTMLResponse)
def upload_driver(
request: Request,
file: UploadFile,
caller: str = Form(""),
) -> HTMLResponse:
"""Accept a driver ZIP, parse its INF, persist via DriverStore + Peewee ORM.
Returns an HTMX partial (partials/driver_list.html) on success, or an
inline error fragment with HTTP 400 on validation failure.
"""
data = file.file.read(MAX_UPLOAD_BYTES + 1)
if len(data) > MAX_UPLOAD_BYTES:
return _error_response("File exceeds 100 MB limit.")
# Must end with .zip
filename = file.filename or ""
if not filename.lower().endswith(".zip"):
return _error_response("Only .zip files are accepted.")
# Must be a valid ZIP archive
if not zipfile.is_zipfile(io.BytesIO(data)):
return _error_response("Uploaded file is not a valid ZIP archive.")
with zipfile.ZipFile(io.BytesIO(data)) as zf:
zip_names = zf.namelist()
# Reject zip-slip paths
for name in zip_names:
if ".." in name or name.startswith("/"):
return _error_response("ZIP contains unsafe paths.")
# Find .inf files
inf_names = [n for n in zip_names if n.lower().endswith(".inf")]
if not inf_names:
return _error_response("No .inf file found in the uploaded ZIP.")
# Prefer amd64/x64 INF when multiple exist; fall back to alphabetical first
preferred = [
n for n in inf_names if "amd64" in n.lower() or "x64" in n.lower()
]
chosen_inf = preferred[0] if preferred else sorted(inf_names)[0]
raw_inf = zf.read(chosen_inf)
# Decode INF
encoding = _detect_encoding(raw_inf)
inf_text = raw_inf.decode(encoding)
# Parse INF
parsed = parse_inf(inf_text, inf_filename=chosen_inf, zip_names=zip_names)
# Persist file (content-addressed, dedup automatic)
# Read DRIVERS_DIR at call time so tests can monkeypatch imptune.config.DRIVERS_DIR
store = DriverStore(_cfg.DRIVERS_DIR)
sha256 = store.save(data)
# Upsert Driver record (no duplicate if same SHA256)
new_driver, _created = Driver.get_or_create(
sha256=sha256,
defaults={
"original_filename": filename,
"size_bytes": len(data),
"driver_desc": json.dumps(parsed.driver_names),
"inf_filename": parsed.inf_filename,
"architecture": parsed.architecture,
"has_cat_file": parsed.has_cat_file,
},
)
# Build driver_data for template
driver_data = build_driver_data(Driver.select().order_by(Driver.uploaded_at.desc()))
# When called from the printer form, emit primary fragment + OOB select refresh
if caller == "printer_form":
return templates.TemplateResponse(
request=request,
name="partials/driver_upload_with_oob.html",
context={
"driver_data": driver_data,
"new_driver_id": new_driver.id,
"parsed": parsed,
},
)
# Default: existing behavior — driver list fragment only
return templates.TemplateResponse(
request=request,
name="partials/driver_list.html",
context={
"driver_data": driver_data,
"parsed": parsed,
"usage": printer_counts_by_driver(request.state.owner),
},
)
@router.patch("/{driver_id}", response_class=HTMLResponse)
def rename_driver(
request: Request, driver_id: int, display_name: str = Form("")
) -> HTMLResponse:
"""Set or clear a driver's display name.
Drivers are global/shared, so this rename is what every Owner sees — the
same as the driver row itself. An empty value clears the rename and the
listing falls back to the ZIP filename.
"""
driver = Driver.get_or_none(Driver.id == driver_id)
if driver is None:
return _error_response(f"Driver {driver_id} not found.", status_code=404)
name = display_name.strip()
if len(name) > MAX_DISPLAY_NAME:
return _error_response(
f"Name must be at most {MAX_DISPLAY_NAME} characters."
)
driver.display_name = name or None
driver.save()
return _render_driver_list(request)
def _driver_icon_status(driver_id: int, sha256: str) -> HTMLResponse:
"""Inline confirmation for the rename/icon dialog, plus an OOB row refresh.
The dialog stays open after picking an icon, so the thumbnail in the table
row behind it is swapped out of band instead of re-rendering the table and
tearing the open `