133 lines
4.3 KiB
Python
133 lines
4.3 KiB
Python
"""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,
|
|
},
|
|
)
|