feat(fleet): write the queued images to every drive that is plugged in

Arm via the queue block toggle with a typed ARM confirmation. A session
lives only for the current run, expires after an hour idle, stamps each
flashed drive by serial so re-inserted sticks are skipped, and requires
every queued image to fit the drive's reported capacity. A failed flash
disarms the session with an error instead of silently retrying.

Author: gowthvm

core/fleet.py

@@ -0,0 +1,100 @@
+"""Fleet flashing: write the queued images to every drive that is plugged in.
+
+A fleet session is never persisted: arming is an explicit, per-run action
+that survives only while the app is open and expires after an hour without
+any activity, so an unattended machine cannot silently keep flashing
+drives forever.
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+IDLE_EXPIRY_SECONDS = 60 * 60
+BYTES_PER_GB = 1_000_000_000
+
+
+@dataclass
+class FleetSession:
+    """One armed session: the images to write and what has been done."""
+
+    images: list[str]
+    arm_ts: float = field(default_factory=time.monotonic)
+    last_activity: float = field(default_factory=time.monotonic)
+    done_serials: set[str] = field(default_factory=set)
+    done_count: int = 0
+    failed_count: int = 0
+
+    def image_sizes(self) -> dict[str, int]:
+        return {
+            path: os.path.getsize(path)
+            for path in self.images
+            if os.path.isfile(path)
+        }
+
+    def fits_on_drive(self, image: str, drive: dict[str, Any]) -> bool:
+        """The drive's reported capacity must hold the whole image.
+
+        The picker reports sizes in decimal gigabytes (size_gb * 1e9),
+        so fleet applies the same convention for the capacity check.
+        """
+        size = self.image_sizes().get(image)
+        if size is None:
+            return False
+        capacity = drive.get("size_gb")
+        if not isinstance(capacity, (int, float)):
+            return False
+        return size <= capacity * BYTES_PER_GB
+
+    def mark_flashed(self, drive: dict[str, Any]) -> None:
+        fp = drive_fingerprint(drive)
+        if fp is not None:
+            self.done_serials.add(fp)
+        self.done_count += 1
+        self.last_activity = time.monotonic()
+
+    def mark_failed(self) -> None:
+        self.failed_count += 1
+        self.last_activity = time.monotonic()
+
+    def expired(self, now: float | None = None) -> bool:
+        """Idle expiry: armed but nothing happened for an hour."""
+        return (now if now is not None else time.monotonic()) - self.last_activity > IDLE_EXPIRY_SECONDS
+
+
+def drive_fingerprint(drive: dict[str, Any]) -> str | None:
+    """Stable per-stick identity for a session.
+
+    The serial is preferred; physical paths are stable enough for a
+    session when the stick reports no serial at all.
+    """
+    serial = drive.get("serial")
+    if serial:
+        return str(serial)
+    path = drive.get("physical_path")
+    return str(path) if path else None
+
+
+def pick_candidate(
+    drives: list[dict[str, Any]],
+    session: FleetSession,
+    now: float | None = None,
+) -> dict[str, Any] | None:
+    """First drive that has not been flashed yet and fits every image."""
+    if session.expired(now):
+        return None
+    sizes = session.image_sizes()
+    if not sizes or len(sizes) != len(session.images):
+        return None
+    for drive in drives:
+        fp = drive_fingerprint(drive)
+        if fp is not None and fp in session.done_serials:
+            continue
+        if all(
+            session.fits_on_drive(image, drive) for image in session.images
+        ):
+            return drive
+    return None

tests/test_fleet.py

@@ -0,0 +1,132 @@
+"""Fleet-mode policy tests: capacity gating, per-stick tracking, expiry.
+
+No real drives and no Qt are involved; every behaviour here is pure
+policy logic from core.fleet.
+"""
+
+from core import fleet
+
+
+def _drive(serial="SN123", path=r"\\.\PHYSICALDRIVE1", size_gb=32):
+    return {"serial": serial, "physical_path": path, "size_gb": size_gb}
+
+
+def _image(tmp_path, name="ubuntu.iso", size=1_000):
+    p = tmp_path / name
+    p.write_bytes(b"\x00" * size)
+    return str(p)
+
+
+class TestDriveFingerprint:
+    def test_serial_preferred(self):
+        assert fleet.drive_fingerprint(_drive()) == "SN123"
+
+    def test_path_fallback_without_serial(self):
+        assert (
+            fleet.drive_fingerprint(_drive(serial="")) == r"\\.\PHYSICALDRIVE1"
+        )
+
+    def test_none_when_neither_known(self):
+        assert fleet.drive_fingerprint({"size_gb": 8}) is None
+
+
+class TestCapacityGate:
+    def test_image_fits_with_room(self, tmp_path):
+        img = _image(tmp_path, size=2_000)
+        session = fleet.FleetSession(images=[img])
+        assert session.fits_on_drive(img, _drive(size_gb=8))
+
+    def test_capacity_boundary_is_inclusive(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        # 0.000001 * 1e9 = exactly 1000 bytes of claimed capacity
+        session = fleet.FleetSession(images=[img])
+        assert session.fits_on_drive(img, _drive(size_gb=0.000001))
+
+    def test_image_larger_than_drive_rejected(self, tmp_path):
+        img = _image(tmp_path, size=2_000)
+        session = fleet.FleetSession(images=[img])
+        assert not session.fits_on_drive(img, _drive(size_gb=0.000001))
+
+    def test_zero_capacity_rejected(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        assert not session.fits_on_drive(img, _drive(size_gb=0))
+
+    def test_missing_image_rejected(self):
+        session = fleet.FleetSession(images=[r"C:\nope.iso"])
+        assert not session.fits_on_drive(r"C:\nope.iso", _drive())
+
+    def test_drive_without_capacity_rejected(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        assert not session.fits_on_drive(img, {"serial": "SN"})
+
+    def test_all_images_must_fit_for_candidate(self, tmp_path):
+        small = _image(tmp_path, "a.iso", size=1_000)
+        big = _image(tmp_path, "b.iso", size=2_000)
+        # capacity is exactly enough for `small` but not both
+        session = fleet.FleetSession(images=[small, big])
+        drive = _drive(size_gb=0.000001)
+        assert not fleet.pick_candidate([drive], session)
+
+    def test_missing_image_blocks_candidate(self, tmp_path):
+        good = _image(tmp_path, "a.iso", size=1_000)
+        session = fleet.FleetSession(images=[good, r"C:\gone.iso"])
+        assert fleet.pick_candidate([_drive()], session) is None
+
+
+class TestSessionTracking:
+    def test_flashed_drive_is_skipped(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        d1 = _drive(serial="SN1")
+        d2 = _drive(serial="SN2")
+        assert fleet.pick_candidate([d1, d2], session) is d1
+        session.mark_flashed(d1)
+        assert session.done_count == 1
+        assert fleet.pick_candidate([d1, d2], session) is d2
+
+    def test_failed_drive_can_be_retried(self, tmp_path):
+        """A failure is recorded but does not blacklist the stick: the
+        operator may re-insert it for another attempt."""
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        d = _drive(serial="SN1")
+        session.mark_failed()
+        assert fleet.pick_candidate([d], session) is d
+        assert session.failed_count == 1
+
+    def test_finished_session_sweeps_every_stick_once(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        ds = [_drive(serial=f"SN{i}") for i in range(4)]
+        for d in ds:
+            assert fleet.pick_candidate(ds, session) is d
+            session.mark_flashed(d)
+        assert fleet.pick_candidate(ds, session) is None
+        assert session.done_count == 4
+
+    def test_sticks_without_serial_tracked_by_path(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        d = _drive(serial="", path=r"\\.\PHYSICALDRIVE9")
+        assert fleet.pick_candidate([d], session) is d
+        session.mark_flashed(d)
+        assert fleet.pick_candidate([d], session) is None
+
+
+class TestExpiry:
+    def test_expired_session_blocks_candidates(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        now = 1_000.0
+        session.last_activity = now - fleet.IDLE_EXPIRY_SECONDS - 1
+        assert fleet.pick_candidate([_drive()], session, now=now) is None
+        assert session.expired(now)
+
+    def test_active_session_keeps_candidates(self, tmp_path):
+        img = _image(tmp_path, size=1_000)
+        session = fleet.FleetSession(images=[img])
+        now = 1_000.0
+        session.last_activity = now - 10
+        assert fleet.pick_candidate([_drive()], session, now=now) is not None