#!/usr/bin/env python3

import argparse
import sys
from dataclasses import dataclass
from pathlib import Path


FLASH_SIZE = 0x8000
SLOT_METADATA_BASE = 0x7FC0
SLOT_STRIDE = 0x40
SLOT_COUNT = 2

CHECKSUM_OFFSET = 0x00
ENTRY_POINT_OFFSET = 0x01
START_SECTOR_OFFSET = 0x05
SIZE_OFFSET = 0x09
PAGE_CKSUM_LO_OFFSET = 0x10
PAGE_CKSUM_HI_OFFSET = 0x11


def read_u16_le(data: bytes, offset: int) -> int:
    return int.from_bytes(data[offset : offset + 2], "little")


def read_u32_le(data: bytes, offset: int) -> int:
    return int.from_bytes(data[offset : offset + 4], "little")


@dataclass
class SlotMetadata:
    slot: int
    metadata_offset: int
    checksum: int
    entry_point: int
    start_sector: int
    size: int
    page_checksum: int

    @property
    def image_start(self) -> int:
        return (self.start_sector + 1) << 7

    @property
    def image_end(self) -> int:
        return self.image_start + self.size


def parse_slot_metadata(image: bytes, slot: int) -> SlotMetadata:
    base = SLOT_METADATA_BASE - (slot * SLOT_STRIDE)
    return SlotMetadata(
        slot=slot,
        metadata_offset=base,
        checksum=image[base + CHECKSUM_OFFSET],
        entry_point=read_u32_le(image, base + ENTRY_POINT_OFFSET),
        start_sector=read_u16_le(image, base + START_SECTOR_OFFSET),
        size=read_u32_le(image, base + SIZE_OFFSET),
        page_checksum=(
            image[base + PAGE_CKSUM_LO_OFFSET]
            | (image[base + PAGE_CKSUM_HI_OFFSET] << 8)
        ),
    )


def compute_slot_checksum(image: bytes, meta: SlotMetadata) -> tuple[int, bool]:
    if meta.image_end > len(image):
        raise ValueError(
            f"slot {meta.slot}: image range 0x{meta.image_start:04x}-0x{meta.image_end:04x} "
            f"extends past dump size 0x{len(image):04x}"
        )

    region = image[meta.image_start : meta.image_end]
    checksum = (-sum(region)) & 0xFF
    has_data = any(byte not in (0x00, 0xFF) for byte in region)
    return checksum, has_data


def validate_slot(image: bytes, slot: int) -> str:
    meta = parse_slot_metadata(image, slot)
    computed_checksum, has_data = compute_slot_checksum(image, meta)
    valid = computed_checksum == meta.checksum and has_data

    lines = [
        f"slot {slot}: {'VALID' if valid else 'INVALID'}",
        f"  metadata_offset: 0x{meta.metadata_offset:04x}",
        f"  entry_point:     0x{meta.entry_point:08x}",
        f"  start_sector:    0x{meta.start_sector:04x}",
        f"  size:            0x{meta.size:08x} ({meta.size})",
        f"  image_range:     0x{meta.image_start:04x}-0x{meta.image_end:04x}",
        f"  stored_checksum: 0x{meta.checksum:02x}",
        f"  calc_checksum:   0x{computed_checksum:02x}",
        f"  has_data:        {'yes' if has_data else 'no'}",
        f"  page_checksum:   0x{meta.page_checksum:04x}",
    ]
    return "\n".join(lines)


def fixup_slot_checksum(image: bytearray, slot: int) -> tuple[SlotMetadata, int]:
    meta = parse_slot_metadata(image, slot)
    computed_checksum, _ = compute_slot_checksum(image, meta)
    image[meta.metadata_offset + CHECKSUM_OFFSET] = computed_checksum
    return meta, computed_checksum


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description=(
            "Validate S-MCU slot checksums in a firmware dump using the same "
            "metadata checksum scheme as the bootloader."
        )
    )
    parser.add_argument("dump", type=Path, help="Path to the firmware dump file.")
    parser.add_argument(
        "--fixup-slot",
        type=int,
        choices=range(SLOT_COUNT),
        help="Recompute and rewrite the metadata checksum byte for the specified slot.",
    )
    parser.add_argument(
        "--output",
        type=Path,
        help="Output path for a modified dump. Required with --fixup-slot unless --in-place is used.",
    )
    parser.add_argument(
        "--in-place",
        action="store_true",
        help="Rewrite the input dump in place when used with --fixup-slot.",
    )
    return parser


def main() -> int:
    args = build_parser().parse_args()
    image = bytearray(args.dump.read_bytes())

    if args.fixup_slot is not None and not args.in_place and args.output is None:
        print(
            "--output or --in-place is required when using --fixup-slot",
            file=sys.stderr,
        )
        return 1

    if args.in_place and args.output is not None:
        print("--output and --in-place are mutually exclusive", file=sys.stderr)
        return 1

    if len(image) < SLOT_METADATA_BASE + 1:
        print(
            f"Dump is too small to contain slot metadata at 0x{SLOT_METADATA_BASE:04x}",
            file=sys.stderr,
        )
        return 1

    if len(image) != FLASH_SIZE:
        print(
            f"Warning: dump size is 0x{len(image):x}, expected 0x{FLASH_SIZE:x}",
            file=sys.stderr,
        )

    try:
        if args.fixup_slot is not None:
            meta, computed_checksum = fixup_slot_checksum(image, args.fixup_slot)
            destination = args.dump if args.in_place else args.output
            destination.write_bytes(image)
            print(
                f"Updated slot {args.fixup_slot} checksum at 0x{meta.metadata_offset + CHECKSUM_OFFSET:04x} "
                f"to 0x{computed_checksum:02x}"
            )
            print(f"Wrote modified dump to {destination}")
            print()

        for slot in range(SLOT_COUNT):
            print(validate_slot(image, slot))
            if slot != SLOT_COUNT - 1:
                print()
    except ValueError as exc:
        print(str(exc), file=sys.stderr)
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())

> ./validate_smcu_dump.py ../../Flash\ Mod/S-MCU_redirect.bin --fixup-slot 0 --in-place Updated slot 0 checksum at 0x7fc0 to 0x71 
Wrote modified dump to ../../Flash Mod/S-MCU_redirect.bin 
slot 0: VALID 
metadata_offset: 0x7fc0 
entry_point: 0x00001491 
start_sector: 0x0028 
size: 0x00001880 (6272) 
image_range: 0x1480-0x2d00 
stored_checksum: 0x71 
calc_checksum: 0x71 
has_data: yes 
page_checksum: 0x0000 

slot 1: INVALID 
metadata_offset: 0x7f80 
entry_point: 0x00000000 
start_sector: 0x0000 
size: 0x00000000 (0) 
image_range: 0x0080-0x0080 
stored_checksum: 0x00 
calc_checksum: 0x00 
has_data: no 
page_checksum: 0x0000 

# Call this script from glasglow with the following command:
#   glasgow script i2c_controller -V 3.3 --scl A1 --sda A0 -f 800
#
# NOTE: Sometimes the writes fail. I just pick up again at the last 
#       succesfully programmed offset.

import ctypes
import enum
import asyncio
import argparse
import time
import traceback

import glasgow

def hexint(value: str) -> int:
    return int(value, 0)

parser = argparse.ArgumentParser(description="S-MCU bootloader flashing script")
parser.add_argument("--no-bootloader", default=False,action="store_true", help="Don't attempt to enter the bootloader, assume already in bootloader mode" )
subparsers = parser.add_subparsers(dest="command")

subparser = subparsers.add_parser("flash", help="Flash firmware to the device")
subparser.add_argument("firmware", type=argparse.FileType("rb"), help="Path to the firmware binary to flash")
subparser.add_argument("offset", type=hexint, help="Offset to flash to")
subparser.add_argument("--seek", type=hexint, help="Seek to the specified offset in the firmware file before reading (default: 0)", default=0)

subparser = subparsers.add_parser("validate", help="Validate the currently flashed firmware")
subparser = subparsers.add_parser("boot", help="Boot the device")

args = parser.parse_args(args.script_args)


BRIDGE_ADDR = 0x5b
BRIDGE_MAGIC = 0xeb

class BridgeOpcodes(enum.IntEnum):
    RECOVERY = 0x00
    IDENTIFY = 0x04

class BridgeCMD(ctypes.Structure):
    _fields_ = [
        ("cmd", ctypes.c_uint8),
        ("addr_hi", ctypes.c_uint8),
        ("addr_mid", ctypes.c_uint8),
        ("addr_lo", ctypes.c_uint8),
        ("footer", ctypes.c_uint8)
    ]
    
BL_ADDR = 0x38
BL_SOP = 0x01
BL_EOP = 0x17

class BLError(Exception):
    pass

class BLOpcodes(enum.IntEnum):
    GET_VERSION = 0x38
    
    VALIDATE_SLOT = 0x31
    WRITE_FLASH_ROW = 0x39
    
    SET_ROW_BUFFER = 0x35
    APPEND_ROW_BUFFER = 0x37
    
    BOOT = 0x3b

class BLCmdHdr(ctypes.Structure):
    _fields_ = [
        ("sop", ctypes.c_uint8),
        ("cmd", ctypes.c_uint8),
        ("payload_len", ctypes.c_uint16),
    ]
class BLCmdFooter(ctypes.Structure):
    _fields_ = [
        ("checksum", ctypes.c_uint16),
        ("eop", ctypes.c_uint8)
    ]

async def bridge_identify():
    cmd = BridgeCMD()
    cmd.cmd = BridgeOpcodes.IDENTIFY
    cmd.footer = BRIDGE_MAGIC
    
    async with i2c_iface.transaction():
        await i2c_iface.write(BRIDGE_ADDR, bytes(cmd))
        response = await i2c_iface.read(BRIDGE_ADDR, 2)
    
    return response

async def enter_bootloader():
    cmd = BridgeCMD()
    cmd.cmd = BridgeOpcodes.RECOVERY
    cmd.addr_hi = 0x38 # recovery magic trigger value
    cmd.footer = BRIDGE_MAGIC
    
    await i2c_iface.write(BRIDGE_ADDR, bytes(cmd))

async def do_bl_cmd(cmd: int | BLOpcodes, payload: bytes = b"", retry : bool=False) -> bytes:
    hdr = BLCmdHdr()
    hdr.sop = BL_SOP
    hdr.cmd = cmd
    hdr.payload_len = len(payload)
    
    pkt = bytes(hdr) + payload
    checksum = (-sum(pkt)) & 0xFFFF
    footer = BLCmdFooter()
    footer.checksum = checksum
    footer.eop = BL_EOP
    pkt += bytes(footer)
    
    if retry:
        while True:
            try:
                await i2c_iface.write(BL_ADDR, pkt)
                break
            except glasgow.applet.interface.i2c_controller.I2CNotAcknowledged as e:
                print(f"Failed to write: {e}")
                print(f"Retrying...")
                await asyncio.sleep(0.5)
    else:
        await i2c_iface.write(BL_ADDR, pkt)
            

    time.sleep(0.2)
    async with i2c_iface.transaction():
        # just fixing the return length for now... guess we could read the length
        # then the rest...
        response = await i2c_iface.read(BL_ADDR, 0x32)
    
    rsp_hdr = BLCmdHdr.from_buffer_copy(response[:ctypes.sizeof(BLCmdHdr)])
    if rsp_hdr.sop != BL_SOP:
        raise BLError("Invalid start of packet")
    if rsp_hdr.cmd != 0x00: # response packets have cmd=0
        raise BLError(f"Unexpected response cmd: {rsp_hdr.cmd}")
    if rsp_hdr.payload_len > len(response) - ctypes.sizeof(BLCmdHdr) - ctypes.sizeof(BLCmdFooter):
        raise BLError("Response payload length mismatch")
    
    payload_start = ctypes.sizeof(BLCmdHdr)
    payload_end = payload_start + rsp_hdr.payload_len
    
    rsp_data = response[payload_start : payload_end]
    rsp_footer = BLCmdFooter.from_buffer_copy(response[payload_end : payload_end + ctypes.sizeof(BLCmdFooter)])
    if rsp_footer.eop != BL_EOP:
        raise BLError("Invalid end of packet")
    expected_checksum = (-sum(response[:payload_end])) & 0xFFFF
    if rsp_footer.checksum != expected_checksum:
        raise BLError("Checksum mismatch")
    
    return rsp_data
    
async def bl_version():
    return await do_bl_cmd(BLOpcodes.GET_VERSION)

async def validate():
    return await do_bl_cmd(BLOpcodes.VALIDATE_SLOT)

async def boot():
    try:
        return await do_bl_cmd(BLOpcodes.BOOT)
    except glasgow.applet.interface.i2c_controller.I2CNotAcknowledged as e:
        print(f"Boot command sent, error is likely as the device is rebooting: {e}")

async def program_row(row_id: int, data: bytes, chunk_size: int = 0x20):
    if len(data) != 0x80:
        raise ValueError("Row data not 0x80 bytes in length")
    
    if chunk_size != 0x80:
        # we are going to chunk the row data into smaller pieces to avoid overwhelming the bootloader, which seems to have some internal buffering limits. The bootloader will reassemble the chunks into the full row before programming.
        while len(data) > 0:
            chunk = data[:chunk_size]
            await do_bl_cmd(BLOpcodes.APPEND_ROW_BUFFER, chunk, retry=True)
            data = data[chunk_size:]
        
        # finally send the command to program the row from the buffer
        payload = b"\x00" + row_id.to_bytes(2, "little")
        await do_bl_cmd(BLOpcodes.WRITE_FLASH_ROW, payload, retry=True)
    else:
        # if the chunk size is the full row size, we can just send it all in one command
        payload = b"\x00" + row_id.to_bytes(2, "little") + data
        await do_bl_cmd(BLOpcodes.WRITE_FLASH_ROW, payload, retry=True)

async def program(offset: int, data: bytes):
    row_id = offset // 0x80
    while len(data) > 0:
        row_data = data[:0x80] # max row size is 128 bytes
        if len(row_data) < 0x80:
            row_data += b"\x00" * (0x80 - len(row_data))
            
        await program_row(row_id, row_data)
        print(f"Programmed row {row_id} at offset 0x{row_id * 0x80:04x}")
        
        row_id += 1
        data = data[0x80:]

async def main():
    # first try to enter the bootloader, in case we're not already there
    try:
        rsp = await bridge_identify()
        print(f"Response: {rsp.hex()}")
        if args.no_bootloader:
            print("Skipping bootloader entry due to --no-bootloader flag")
        else:
            print("Device is responsive, attempting to enter bootloader...")
            await enter_bootloader()
            await asyncio.sleep(1.0)
    except Exception as e:
        print(f"Error: {e}")
        print("Assuming in recovery mode...")
    
    
    for i in range(1):  
        version = await bl_version()
        print(f"Bootloader version response: {version.hex()}")
        
    if args.command == "flash":
        fd = args.firmware
        fd.seek(args.seek)
        firmware_data = fd.read()
        await program(args.offset, firmware_data)
    elif args.command == "validate":
        validation_result = await validate()
        print(f"Validation result: {validation_result.hex()}")
    elif args.command == "boot":
        await boot()
        print("Boot command sent. Device should reboot now.")

try:
    await main()
except KeyboardInterrupt:
    print("Interrupted by user")
except Exception as e:
    traceback.print_exc()

[0000000000]MEP_ipolling() 
irq detected!!!! 
FE raid id 0x0 raid 0x0 nr_disk 0x0 status 0xf stripe_size 0x0 
Hello, ANVIL Modified Loader!!!!
T: 0,0 
2x4 !! 
single host !! 
falling edge detect 0x0 
REG 0xE1248158 value