saltylab-firmware/legacy/stm32/scripts/flash_firmware.py

#!/usr/bin/env python3
"""SaltyLab Firmware OTA Flash Script — Issue #124

Flashes ESP32-S3 firmware via PlatformIO (pio run -t upload).
Supports CRC32 integrity verification and host-side backup/rollback.

Usage:
    python flash_firmware.py firmware.bin [options]
    python flash_firmware.py --rollback
    python flash_firmware.py --board balance   # flash esp32/balance/ via PlatformIO
    python flash_firmware.py --board io        # flash esp32/io/ via PlatformIO

Options:
    --board NAME        Board to flash: 'balance' (ESP32-S3 BALANCE) or 'io' (ESP32-S3 IO)
    --rollback          Flash the previous firmware backup
    --dry-run           Print pio command but do not execute

Requirements:
    pip install platformio   (PlatformIO CLI)
    pio >= 6.x installed and in PATH

ESP32-S3 note:
    ESP32-S3 BALANCE board uses CH343G USB bridge — flashing via USB UART.
    ESP32-S3 IO board uses built-in JTAG/USB-CDC.
    Both flashed with: pio run -t upload  in the respective esp32/ subdirectory.
"""

import argparse
import binascii
import os
import shutil
import struct
import subprocess
import sys
import time

# ---- ESP32-S3 flash constants ----
# ESP32-S3 flash is managed by PlatformIO/esptool; these values are kept
# for reference only (CRC utility functions below are still valid for
# cross-checking firmware images).
FLASH_BASE  = 0x00000000
FLASH_SIZE  = 0x800000      # 8 MB

# ---- PlatformIO board directories ----
BOARD_DIRS = {
    "balance": "esp32/balance",
    "io":      "esp32/io",
}

BACKUP_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)),
                            '.firmware_backup.bin')


# ---- CRC utilities ----

def crc32_file(path: str) -> int:
    """
    Compute CRC-32/ISO-HDLC (standard Python binascii.crc32) of a file.
    Used for pre-flash integrity verification; consistent across runs.
    """
    with open(path, 'rb') as fh:
        data = fh.read()
    return binascii.crc32(data) & 0xFFFFFFFF


def esp32_crc32(data: bytes) -> int:
    """
    Compute CRC-32/MPEG-2 for firmware image integrity verification.

    Note: ESP32-S3 uses esptool for flashing; this CRC is for host-side
    integrity checking only (not the ESP32 hardware CRC unit).

    data should be padded to a 4-byte boundary with 0xFF before calling.
    """
    if len(data) % 4:
        data += b'\xff' * (4 - len(data) % 4)

    crc = 0xFFFFFFFF
    for i in range(0, len(data), 4):
        # Little-endian: byte0 is LSB, byte3 is MSB of the 32-bit word
        word = struct.unpack_from('<I', data, i)[0]
        crc ^= word
        for _ in range(32):
            if crc & 0x80000000:
                crc = ((crc << 1) ^ 0x04C11DB7) & 0xFFFFFFFF
            else:
                crc = (crc << 1) & 0xFFFFFFFF
    return crc


# ---- JLink protocol helpers ----

def _crc16_xmodem(data: bytes) -> int:
    """CRC-16/XMODEM (poly 0x1021, init 0x0000) — JLink frame CRC."""
    crc = 0x0000
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = ((crc << 1) ^ 0x1021) & 0xFFFF
            else:
                crc = (crc << 1) & 0xFFFF
    return crc


def _build_jlink_frame(cmd: int, payload: bytes = b'') -> bytes:
    """Build a JLink binary frame: [STX][LEN][CMD][PAYLOAD][CRC_hi][CRC_lo][ETX]."""
    STX, ETX = 0x02, 0x03
    body   = bytes([cmd]) + payload
    length = len(body)
    crc    = _crc16_xmodem(body)
    return bytes([STX, length, cmd]) + payload + bytes([crc >> 8, crc & 0xFF, ETX])


def trigger_dfu_via_jlink(port: str, baud: int = 921600) -> None:
    """Send JLINK_CMD_DFU_ENTER (0x06) over USART1 to put device in DFU mode."""
    try:
        import serial
    except ImportError:
        print("ERROR: pyserial not installed.  Run: pip install pyserial",
              file=sys.stderr)
        sys.exit(1)

    frame = _build_jlink_frame(0x06)  # JLINK_CMD_DFU_ENTER, no payload
    with serial.Serial(port, baud, timeout=2) as ser:
        ser.write(frame)
        time.sleep(0.1)
    print(f"DFU_ENTER sent to {port} ({len(frame)} bytes)")


# ---- Flash ----

def flash(bin_path: str, board: str = "balance",
          dry_run: bool = False) -> int:
    """
    Flash firmware using PlatformIO (pio run -t upload).
    Returns the process exit code.

    board: 'balance' → esp32/balance/, 'io' → esp32/io/
    """
    board_dir = BOARD_DIRS.get(board, BOARD_DIRS["balance"])
    cmd = ['pio', 'run', '-t', 'upload', '--project-dir', board_dir]
    print('Running:', ' '.join(cmd))
    if dry_run:
        print('[dry-run] skipping pio upload execution')
        return 0
    return subprocess.call(cmd)


# ---- Main ----

def main() -> int:
    parser = argparse.ArgumentParser(
        description='SaltyLab ESP32-S3 firmware flash via PlatformIO (Issue #124)'
    )
    parser.add_argument('firmware', nargs='?',
                        help='Firmware .bin file (for CRC check only; PlatformIO handles actual flash)')
    parser.add_argument('--board', default='balance',
                        choices=['balance', 'io'],
                        help='Board to flash: balance (ESP32-S3 BALANCE) or io (ESP32-S3 IO)')
    parser.add_argument('--rollback', action='store_true',
                        help='Flash the previous firmware backup')
    parser.add_argument('--dry-run', action='store_true',
                        help='Print pio command without executing it')
    args = parser.parse_args()

    # Determine target binary
    if args.rollback:
        if not os.path.exists(BACKUP_PATH):
            print(f'ERROR: No backup found at {BACKUP_PATH}', file=sys.stderr)
            return 1
        target = BACKUP_PATH
        print(f'Rolling back to {BACKUP_PATH}')
    elif args.firmware:
        target = args.firmware
    else:
        parser.print_help()
        return 1

    if not os.path.exists(target):
        print(f'ERROR: File not found: {target}', file=sys.stderr)
        return 1

    # CRC32 integrity check
    crc_std = crc32_file(target)
    size    = os.path.getsize(target)
    print(f'Binary : {target}  ({size} bytes)')
    print(f'CRC-32 : 0x{crc_std:08X}  (ISO-HDLC)')

    if size > FLASH_SIZE:
        print(f'ERROR: Binary ({size} bytes) exceeds flash size '
              f'({FLASH_SIZE} bytes)', file=sys.stderr)
        return 1

    # CRC for cross-checking firmware integrity
    with open(target, 'rb') as fh:
        bin_data = fh.read()
    crc_hw = esp32_crc32(bin_data.ljust(FLASH_SIZE, b'\xff'))
    print(f'CRC-32 : 0x{crc_hw:08X}  (MPEG-2, padded to {FLASH_SIZE // 1024} KB)')

    # Save backup before flashing (skip when rolling back)
    if not args.rollback:
        shutil.copy2(target, BACKUP_PATH)
        print(f'Backup : {BACKUP_PATH}')

    # Flash via PlatformIO
    rc = flash(target, args.board, args.dry_run)
    if rc == 0:
        print('Flash complete — ESP32-S3 should reset into application.')
    else:
        print(f'ERROR: pio run -t upload exited with code {rc}', file=sys.stderr)
    return rc


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