DoorCounter/tools/flash_device.py

#!/usr/bin/env python3
"""
flash_device.py — Write NVS config to TimerCamera-F over serial.

Requires: pip install esptool esp-idf-nvs-partition-gen
Usage:
  python flash_device.py \\
    --port /dev/ttyUSB0 \\
    --device-id dc-0042 \\
    --location-id retailer-123 \\
    --hmac-secret <32-byte-hex>   # omit to auto-generate \\
    [--wifi-ssid "StoreWiFi"] \\
    [--wifi-password "secret"] \\
    [--line-offset 50]
"""
import argparse
import os
import re
import secrets
import subprocess
import sys
import tempfile

HMAC_SECRET_RE = re.compile(r"^[0-9a-fA-F]{64}$")


NVS_NAMESPACE        = "doorcounter"
NVS_PARTITION_OFFSET = "0x9000"
NVS_PARTITION_SIZE   = "0x5000"   # matches firmware partition table (20KB)

# Characters that would change the field/row structure of the NVS-CSV format
# (key,type,encoding,value). A value containing any of these would either
# split into more fields or add rows, silently provisioning the wrong keys.
_CSV_FORBIDDEN = (",", '"', "\n", "\r")


def _reject_csv_metacharacters(name, value):
    """Exit with an error if value contains a character that would corrupt
    the NVS CSV. Used for operator-supplied strings (device id, location id,
    WiFi credentials)."""
    for c in _CSV_FORBIDDEN:
        if c in value:
            print(
                f"Error: --{name} contains forbidden character {c!r}; "
                f"this would corrupt the NVS partition CSV.",
                file=sys.stderr,
            )
            sys.exit(1)


def build_nvs_csv(device_id, location_id, hmac_secret,
                  wifi_ssid=None, wifi_pass=None, line_offset=50):
    rows = [
        "key,type,encoding,value",
        f"{NVS_NAMESPACE},namespace,,",
        f"device_id,data,string,{device_id}",
        f"location_id,data,string,{location_id}",
        f"hmac_secret,data,string,{hmac_secret}",
        f"line_offset,data,u32,{line_offset}",
    ]
    if wifi_ssid is not None:
        rows.append(f"wifi_ssid,data,string,{wifi_ssid}")
    if wifi_pass is not None:
        rows.append(f"wifi_pass,data,string,{wifi_pass}")
    return "\n".join(rows) + "\n"


def main():
    parser = argparse.ArgumentParser(
        description="Provision TimerCamera-F NVS config over serial")
    parser.add_argument("--port",          required=True,
                        help="Serial port, e.g. /dev/ttyUSB0 or COM3")
    parser.add_argument("--device-id",     required=True,
                        help="Unique device ID, e.g. dc-0042")
    parser.add_argument("--location-id",   required=True,
                        help="Retailer location ID, e.g. retailer-123")
    parser.add_argument("--hmac-secret",   default=None,
                        help="32-byte hex HMAC secret (auto-generated if omitted)")
    parser.add_argument("--wifi-ssid",     default=None,
                        help="WiFi SSID (optional — user can set via captive portal)")
    parser.add_argument("--wifi-password", default=None,
                        help="WiFi password (optional)")
    parser.add_argument("--line-offset",   type=int, default=50,
                        help="Virtual line position %% of frame height (default 50)")
    args = parser.parse_args()

    hmac_secret = args.hmac_secret or secrets.token_hex(32)
    if not HMAC_SECRET_RE.match(hmac_secret):
        print("Error: --hmac-secret must be exactly 64 hex characters (32 bytes)",
              file=sys.stderr)
        sys.exit(1)
    if args.hmac_secret is None:
        print(f"Generated HMAC secret: {hmac_secret}")
        print("  *** SAVE THIS — you need it to register the device on the server ***")

    if args.line_offset < 0 or args.line_offset > 100:
        print("Error: --line-offset must be 0-100", file=sys.stderr)
        sys.exit(1)

    _reject_csv_metacharacters("device-id",   args.device_id)
    _reject_csv_metacharacters("location-id", args.location_id)
    if args.wifi_ssid is not None:
        _reject_csv_metacharacters("wifi-ssid",     args.wifi_ssid)
    if args.wifi_password is not None:
        _reject_csv_metacharacters("wifi-password", args.wifi_password)

    with tempfile.TemporaryDirectory() as tmp:
        csv_path = os.path.join(tmp, "nvs.csv")
        bin_path = os.path.join(tmp, "nvs.bin")

        csv_content = build_nvs_csv(
            args.device_id, args.location_id, hmac_secret,
            args.wifi_ssid, args.wifi_password, args.line_offset
        )
        with open(csv_path, "w") as f:
            f.write(csv_content)

        # Generate NVS binary
        ret = subprocess.run(
            [sys.executable, "-m", "esp_idf_nvs_partition_gen", "generate",
             csv_path, bin_path, NVS_PARTITION_SIZE],
            capture_output=True, text=True
        )
        if ret.returncode != 0:
            print(f"nvs_partition_gen error:\n{ret.stderr}", file=sys.stderr)
            sys.exit(1)

        # Flash NVS partition
        ret = subprocess.run(
            ["esptool.py", "--port", args.port, "--chip", "esp32",
             "write_flash", NVS_PARTITION_OFFSET, bin_path]
        )
        sys.exit(ret.returncode)


if __name__ == "__main__":
    main()