saltylab-firmware/scripts/flash_firmware.py

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

Flashes firmware via USB DFU using dfu-util.
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 firmware.bin --trigger-dfu /dev/ttyUSB0

Options:
    --vid HEX           USB vendor ID   (default: 0x0483 — STMicroelectronics)
    --pid HEX           USB product ID  (default: 0xDF11 — DFU mode)
    --alt N             DFU alt setting (default: 0 — internal flash)
    --rollback          Flash the previous firmware backup
    --trigger-dfu PORT  Send DFU_ENTER over JLink UART before flashing
    --dry-run           Print dfu-util command but do not execute

Requirements:
    pip install pyserial  (only if using --trigger-dfu)
    dfu-util >= 0.9 installed and in PATH

Dual-bank note:
    STM32F722 has single-bank 512 KB flash; hardware A/B rollback is not
    supported.  Rollback is implemented here by saving a backup of the
    previous binary (.firmware_backup.bin) before each flash.
"""

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

# ---- STM32F722 flash constants ----
FLASH_BASE  = 0x08000000
FLASH_SIZE  = 0x80000       # 512 KB

# ---- DFU device defaults (STM32 system bootloader) ----
DFU_VID     = 0x0483        # STMicroelectronics
DFU_PID     = 0xDF11        # DFU mode

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 stm32_crc32(data: bytes) -> int:
    """
    Compute CRC-32/MPEG-2 matching STM32F7 hardware CRC unit.

    STM32 algorithm:
        Polynomial : 0x04C11DB7
        Initial    : 0xFFFFFFFF
        Width      : 32 bits
        Reflection : none (MSB-first)
        Feed size  : 32-bit words from flash (little-endian CPU read)

    When the STM32 reads a flash word it gets a little-endian uint32;
    the hardware CRC unit processes bits[31:24] first, then [23:16],
    [15:8], [7:0].  This Python implementation replicates that behaviour.

    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, vid: int, pid: int, alt: int = 0,
          dry_run: bool = False) -> int:
    """
    Flash firmware using dfu-util.  Returns the process exit code.

    Uses --dfuse-address with :leave to reset into application after flash.
    """
    addr = f'0x{FLASH_BASE:08X}'
    cmd  = [
        'dfu-util',
        '--device',          f'{vid:04x}:{pid:04x}',
        '--alt',             str(alt),
        '--dfuse-address',   f'{addr}:leave',
        '--download',        bin_path,
    ]
    print('Running:', ' '.join(cmd))
    if dry_run:
        print('[dry-run] skipping dfu-util execution')
        return 0
    return subprocess.call(cmd)


# ---- Main ----

def main() -> int:
    parser = argparse.ArgumentParser(
        description='SaltyLab firmware OTA flash via USB DFU (Issue #124)'
    )
    parser.add_argument('firmware', nargs='?',
                        help='Firmware .bin file to flash')
    parser.add_argument('--vid',  type=lambda x: int(x, 0), default=DFU_VID,
                        help=f'USB vendor ID  (default: 0x{DFU_VID:04X})')
    parser.add_argument('--pid',  type=lambda x: int(x, 0), default=DFU_PID,
                        help=f'USB product ID (default: 0x{DFU_PID:04X})')
    parser.add_argument('--alt',  type=int, default=0,
                        help='DFU alt setting (default: 0 — internal flash)')
    parser.add_argument('--rollback', action='store_true',
                        help='Flash the previous firmware backup')
    parser.add_argument('--trigger-dfu', metavar='PORT',
                        help='Trigger DFU via JLink UART before flashing '
                             '(e.g. /dev/ttyUSB0 or COM3)')
    parser.add_argument('--dry-run', action='store_true',
                        help='Print dfu-util command without executing it')
    args = parser.parse_args()

    # Optionally trigger DFU mode over JLink serial
    if args.trigger_dfu:
        trigger_dfu_via_jlink(args.trigger_dfu)
        print('Waiting 3 s for USB DFU enumeration…')
        time.sleep(3)

    # 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

    # STM32 hardware CRC (for cross-checking with firmware telemetry)
    with open(target, 'rb') as fh:
        bin_data = fh.read()
    crc_hw = stm32_crc32(bin_data.ljust(FLASH_SIZE, b'\xff'))
    print(f'CRC-32 : 0x{crc_hw:08X}  (MPEG-2 / STM32 HW, 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
    rc = flash(target, args.vid, args.pid, args.alt, args.dry_run)
    if rc == 0:
        print('Flash complete — device should reset into application.')
    else:
        print(f'ERROR: dfu-util exited with code {rc}', file=sys.stderr)
    return rc


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