"""Icon upload API — POST /printers/{printer_id}/icon.""" from __future__ import annotations import hashlib import io from pathlib import Path from fastapi import APIRouter, Request, UploadFile from fastapi.responses import FileResponse, HTMLResponse, Response from PIL import Image import imptune.config as cfg from imptune.db.models import Icon, Printer router = APIRouter(prefix="/printers") MAX_ICON_BYTES = 750 * 1024 # 750 KB @router.post("/{printer_id}/icon", response_class=HTMLResponse) def upload_icon(request: Request, printer_id: int, file: UploadFile) -> HTMLResponse: """Accept a printer icon PNG, validate it, store it, and update the Icon record. Validation rules: - Format: PNG only - Dimensions: exactly 256x256 pixels - Size: at most 750 KB Replaces any previously uploaded icon for this printer. Returns an HTMX-friendly HTML fragment. """ # Check printer exists and belongs to this owner printer = Printer.get_or_none( (Printer.id == printer_id) & (Printer.owner == request.state.owner) ) if printer is None: return HTMLResponse( content="

Printer not found.

", status_code=404, ) # Read file (read one byte extra to detect oversized files) data = file.file.read(MAX_ICON_BYTES + 1) if len(data) > MAX_ICON_BYTES: return HTMLResponse( content="

Icon exceeds 750 KB limit.

", status_code=422, ) # Validate with Pillow try: img = Image.open(io.BytesIO(data)) except Exception: return HTMLResponse( content="

Icon must be PNG format.

", status_code=422, ) if img.format != "PNG": return HTMLResponse( content="

Icon must be PNG format.

", status_code=422, ) if img.size != (256, 256): return HTMLResponse( content=f"

Icon must be 256x256 pixels, got {img.size}.

", status_code=422, ) # Store SHA256-addressed on disk. cfg.ICONS_DIR is the same path the # .intunewin export reads from — deriving it a second time from DATA_DIR # would silently diverge if the layout ever changes. sha256 = hashlib.sha256(data).hexdigest() icons_dir = Path(cfg.ICONS_DIR) icons_dir.mkdir(parents=True, exist_ok=True) icon_path = icons_dir / sha256 icon_path.write_bytes(data) # Replace existing Icon record for this printer Icon.delete().where(Icon.printer == printer_id).execute() Icon.create( printer=printer_id, sha256=sha256, original_filename=file.filename or "icon.png", size_bytes=len(data), ) return HTMLResponse( content=( '

Icon uploaded successfully

' f'
{img.size[0]}×{img.size[1]} PNG
' ), status_code=200, ) @router.get("/{printer_id}/icon") def get_icon(request: Request, printer_id: int) -> Response: """Serve the stored 256x256 PNG so the UI can show what was uploaded. Owner-scoped: a printer belonging to another owner reads as missing. """ printer = Printer.get_or_none( (Printer.id == printer_id) & (Printer.owner == request.state.owner) ) if printer is None: return Response(status_code=404) icon = Icon.get_or_none(Icon.printer == printer_id) if icon is None: return Response(status_code=404) icon_path = Path(cfg.ICONS_DIR) / icon.sha256 if not icon_path.exists(): return Response(status_code=404) return FileResponse( icon_path, media_type="image/png", headers={"Cache-Control": "private, max-age=300"}, )