"""Driver upload API — POST /drivers/upload.""" from __future__ import annotations import io import json import zipfile from pathlib import Path from fastapi import APIRouter, Form, Request, UploadFile from fastapi.responses import HTMLResponse from fastapi.templating import Jinja2Templates import imptune.config as _cfg from imptune.db.models import Driver from imptune.services.inf_parser import _detect_encoding, parse_inf from imptune.storage.driver_store import DriverStore router = APIRouter(prefix="/drivers") templates = Jinja2Templates( directory=str(Path(__file__).parent.parent / "templates") ) MAX_UPLOAD_BYTES = 100 * 1024 * 1024 # 100 MB def _error_response(message: str, status_code: int = 400) -> HTMLResponse: """Return an HTMX-friendly error fragment swapped into #driver-list.""" return HTMLResponse( content=f"

{message}

", status_code=status_code, ) @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 drivers = list(Driver.select().order_by(Driver.uploaded_at.desc())) driver_data = [] for d in drivers: names = json.loads(d.driver_desc) if d.driver_desc else [] driver_data.append({"driver": d, "names": names}) # 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, }, )