Commit initial

imptune/api/drivers.py

@@ -0,0 +1,132 @@
+"""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"<div id='driver-list' class='error'><p>{message}</p></div>",
+        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,
+        },
+    )