import hashlib
import struct


def fix_esp32_image(file_path):
    with open(file_path, "rb") as f:
        data = bytearray(f.read())

    n_segments = data[1]

    pos = 0x18
    checksum = 0xEF
    for i in range(n_segments):
        seg_len = struct.unpack("<I", data[pos + 4 : pos + 8])[0]
        seg_data = data[pos + 8 : pos + 8 + seg_len]
        for b in seg_data:
            checksum ^= b
        pos += 8 + seg_len

    pad_len = 15 - (pos % 16)
    pos += pad_len

    print(f"Old Checksum: {hex(data[pos])}, New Checksum: {hex(checksum)}")
    data[pos] = checksum

    hash_data = data[: pos + 1]
    new_hash = hashlib.sha256(hash_data).digest()

    if len(data) >= pos + 1 + 32:
        print(f"Old Hash: {data[pos + 1 : pos + 33].hex()}")
        print(f"New Hash: {new_hash.hex()}")
        data[pos + 1 : pos + 33] = new_hash

    with open(file_path, "wb") as f:
        f.write(data)
    print(f"Fixed {file_path}")


if __name__ == "__main__":
    files = [
        "/home/sapient/Public/03_Projects_and_Dev/Meshcore_TDECK/Ultra-TDeck-v9.2.patched.bin",
        "/home/sapient/Public/03_Projects_and_Dev/Meshcore_TDECK/MeshOS-TDeck-1.1.8.patched.bin",
    ]
    for f in files:
        fix_esp32_image(f)