kota-technician-tool/kota_tech/sd/win_format.py

"""Elevated Windows SD-card formatter (runs as: python -m kota_tech.sd.win_format ...).

Replaces Clear-Disk / New-Partition / Format-Volume: every Storage cmdlet
goes through the storage provider, which stalls ~13 s on an empty card
reader slot, several times per format. This helper touches only the chosen
disk through its own driver:

  1. guard: bus is USB/SD/MMC, size is exactly the one the technician
     confirmed, and the disk does not hold the Windows drive;
  2. lock + dismount the card's volumes;
  3. raw writes: zero old MBR/GPT, FAT32 structures (kota_tech.sd.fat32),
     then the new MBR last;
  4. unlock, let Windows re-read the disk, wait for the FAT32 'KOTA' volume;
  5. read-only `chkdsk` of the new volume.

The result (with per-step timings) goes to a JSON file for the non-elevated
tool, which then writes ENV.TXT through the normal file API.
"""

from __future__ import annotations

import argparse
import ctypes
import json
import os
import random
import subprocess
import sys
import time
from ctypes import wintypes

from . import fat32, win_disks

GENERIC_RW, SHARE_RW, OPEN_EXISTING = 0xC0000000, 0x3, 3
FSCTL_LOCK_VOLUME, FSCTL_UNLOCK_VOLUME, FSCTL_DISMOUNT_VOLUME = 0x00090018, 0x0009001C, 0x00090020
IOCTL_DISK_UPDATE_PROPERTIES = 0x00070140
CARD_BUSES = {"USB", "SD", "MMC"}
CHUNK = 1024 * 1024

k32 = win_disks._k32
k32.WriteFile.argtypes = [wintypes.HANDLE, wintypes.LPCVOID, wintypes.DWORD,
                          ctypes.POINTER(wintypes.DWORD), wintypes.LPVOID]
k32.SetFilePointerEx.argtypes = [wintypes.HANDLE, ctypes.c_longlong, ctypes.POINTER(ctypes.c_longlong), wintypes.DWORD]
k32.VirtualAlloc.restype = wintypes.LPVOID
k32.VirtualAlloc.argtypes = [wintypes.LPVOID, ctypes.c_size_t, wintypes.DWORD, wintypes.DWORD]


class FormatError(RuntimeError):
    pass


def _open_rw(path: str):
    h = k32.CreateFileW(path, GENERIC_RW, SHARE_RW, None, OPEN_EXISTING, 0, None)
    if h in (None, win_disks._INVALID):
        raise FormatError(f"cannot open {path} (error {ctypes.get_last_error()})")
    return h


def _ioctl(h, code: int) -> bool:
    got = wintypes.DWORD()
    return bool(k32.DeviceIoControl(h, code, None, 0, None, 0, ctypes.byref(got), None))


class RawDisk:
    """Sector-aligned writes through a page-aligned buffer (device alignment rules)."""

    def __init__(self, number: int, bps: int):
        self.h = _open_rw(rf"\\.\PhysicalDrive{number}")
        self.bps = bps
        self.buf = k32.VirtualAlloc(None, CHUNK, 0x3000, 0x04)      # MEM_COMMIT|RESERVE, RW

    def write(self, offset: int, data: bytes) -> None:
        if offset % self.bps or len(data) % self.bps:
            raise FormatError("unaligned write")
        pos = 0
        while pos < len(data):
            n = min(CHUNK, len(data) - pos)
            ctypes.memmove(self.buf, data[pos:pos + n], n)
            self._raw(offset + pos, n)
            pos += n

    def zero(self, offset: int, length: int) -> None:
        ctypes.memset(self.buf, 0, CHUNK)
        pos = 0
        while pos < length:
            n = min(CHUNK, length - pos)
            self._raw(offset + pos, n)
            pos += n

    def _raw(self, offset: int, n: int) -> None:
        if not k32.SetFilePointerEx(self.h, offset, None, 0):
            raise FormatError(f"seek failed (error {ctypes.get_last_error()})")
        done = wintypes.DWORD()
        if not k32.WriteFile(self.h, self.buf, n, ctypes.byref(done), None) or done.value != n:
            raise FormatError(f"write at {offset} failed (error {ctypes.get_last_error()}) - write-protected card?")

    def close(self) -> None:
        k32.FlushFileBuffers(self.h)
        _ioctl(self.h, IOCTL_DISK_UPDATE_PROPERTIES)                 # re-read the new partition table
        k32.CloseHandle(self.h)


def format_card(number: int, expected_size: int, label: str, timings: dict) -> str:
    t0 = time.perf_counter()
    disk = next((d for d in win_disks.list_disks(CARD_BUSES) if d.id == str(number)), None)
    if disk is None:
        raise FormatError("Refusing: the disk is not a USB/SD/MMC card reader.")
    if disk.size != expected_size:
        raise FormatError("Refusing: the disk changed since it was selected (card swapped?).")
    system = os.environ.get("SystemDrive", "C:").upper().rstrip("\\") + "\\"
    if any(v.mount.upper() == system for v in disk.volumes):
        raise FormatError("Refusing: this disk holds the Windows drive.")
    timings["guard"] = time.perf_counter() - t0

    t = time.perf_counter()
    locked = []
    for v in disk.volumes:
        h = _open_rw(rf"\\.\{v.mount[0]}:")
        for _ in range(20):
            if _ioctl(h, FSCTL_LOCK_VOLUME):
                break
            time.sleep(0.25)
        else:
            k32.CloseHandle(h)
            raise FormatError(f"{v.mount} is in use by another program - close it and try again.")
        _ioctl(h, FSCTL_DISMOUNT_VOLUME)
        locked.append(h)
    timings["lock_volumes"] = time.perf_counter() - t

    t = time.perf_counter()
    lay = fat32.plan(disk.size)
    zeros, writes = fat32.build(lay, label, random.getrandbits(32), random.getrandbits(32), disk.size)
    raw = RawDisk(number, lay.bytes_per_sector)
    try:
        for off, n in zeros:
            raw.zero(off, n)
        for off, data in writes:                       # MBR is last: the partition appears complete
            raw.write(off, data)
    finally:
        for h in locked:
            _ioctl(h, FSCTL_UNLOCK_VOLUME)
            k32.CloseHandle(h)
        raw.close()
    timings["write"] = time.perf_counter() - t

    t = time.perf_counter()
    drive = ""
    for _ in range(80):                                # Windows mounts the new volume by itself
        d = next((x for x in win_disks.list_disks(CARD_BUSES) if x.id == str(number)), None)
        vol = next((v for v in (d.volumes if d else []) if v.filesystem.upper() == "FAT32"), None)
        if vol and vol.label.upper() == label.upper():
            drive = vol.mount
            break
        time.sleep(0.25)
    timings["mount"] = time.perf_counter() - t
    if not drive:
        raise FormatError("Formatted, but Windows did not mount the card - remove and re-insert it.")

    t = time.perf_counter()
    chk = subprocess.run(["chkdsk", drive.rstrip("\\")], capture_output=True, text=True, timeout=120,
                         creationflags=subprocess.CREATE_NO_WINDOW)
    timings["chkdsk"] = time.perf_counter() - t
    if chk.returncode != 0:
        tail = " ".join(chk.stdout.split()[-30:])
        raise FormatError(f"chkdsk reported problems on the new file system: {tail}")
    return drive


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--disk", type=int, required=True)
    ap.add_argument("--size", type=int, required=True)
    ap.add_argument("--label", default="KOTA")
    ap.add_argument("--result", required=True)
    a = ap.parse_args(argv)
    timings: dict[str, float] = {}
    result = {"ok": False, "message": "", "drive": "", "timings": timings}
    try:
        result["drive"] = format_card(a.disk, a.size, a.label, timings)
        result["ok"] = True
        result["message"] = "Formatted FAT32 (MBR), checked with chkdsk."
    except FormatError as exc:
        result["message"] = str(exc)
    except Exception as exc:                            # report everything to the tool
        result["message"] = f"Unexpected error: {exc!r}"
    with open(a.result, "w", encoding="utf-8") as fh:
        json.dump(result, fh)
    return 0 if result["ok"] else 1


if __name__ == "__main__":
    sys.exit(main())