import hashlib
import struct
import os

def fix_esp32_image(file_path):
    print(f"Fixing {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):
        header = data[pos:pos+8]
        if len(header) < 8: break
        seg_len = struct.unpack("<I", header[4: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
    
    if pos < len(data):
        data[pos] = checksum
        hash_data = data[:pos+1]
        new_hash = hashlib.sha256(hash_data).digest()
        if len(data) >= pos + 1 + 32:
            data[pos+1:pos+33] = new_hash
            
    with open(file_path, "wb") as f:
        f.write(data)
    print(f"Fixed footer for {file_path}")

if __name__ == "__main__":
    PATCHED_DIR = "/home/sapient/Public/esp32-analysis/patched_binaries"
    files = [
        os.path.join(PATCHED_DIR, "Ultra-TDeck-v9.2.patched.bin"),
        os.path.join(PATCHED_DIR, "MeshOS-TDeck-1.1.8.patched.bin")
    ]
    for f in files:
        fix_esp32_image(f)