#!/usr/bin/env python3
import os
import sys

# TDD Verification Harness for Patched ESP32-S3 Binaries
# Extensible framework for verifying known 'always-true' patches.

# --- Patch Definitions ---
# The keys represent the target identifier (substring match against filename).
# 'offset': Relative to the start of the app partition (usually 0x10000 in merged binaries).
# 'original': The unpatched bytecode.
# 'patched': The expected bytecode after patching.
PATCH_DEFS = {
    "Ultra-TDeck": {
        "patches": [
            {
                "name": "Always-True License Check",
                "app_offset": 0xBA62D,
                "original": bytes([0x26, 0x15, 0x3C]),
                "patched": bytes([0x06, 0x0F, 0x00]),
            }
        ]
    },
    "MeshOS": {
        "patches": [
            {
                "name": "Always-True License Check",
                "app_offset": 0xB99ED,
                "original": bytes([0x26, 0x19, 0x3C]),
                "patched": bytes([0x06, 0x0F, 0x00]),
            }
        ]
    },
}

# --- Known Binaries to Verify ---
BINARIES_TO_CHECK = [
    "patched_binaries/Ultra-TDeck-v9.2.patched.bin",
    "patched_binaries/Ultra-TDeck-v9.2-merged.patched.bin",
    "patched_binaries/MeshOS-TDeck-1.1.8.patched.bin",
    "patched_binaries/MeshOS-TDeck-1.1.8-merged.patched.bin",
]


def verify_binary(filepath):
    print(f"[*] Verifying binary: {filepath}")

    if not os.path.exists(filepath):
        print(f"  [!] ERROR: File not found: {filepath}")
        return False

    filename = os.path.basename(filepath)

    # Determine the target type based on filename
    target_def = None
    for key, defs in PATCH_DEFS.items():
        if key in filename:
            target_def = defs
            break

    if not target_def:
        print(f"  [!] ERROR: No patch definitions found for {filename}")
        return False

    # Adjust offset if the binary is merged (app partition starts at 0x10000)
    is_merged = "merged" in filename.lower()
    base_offset = 0x10000 if is_merged else 0x0

    all_passed = True

    try:
        with open(filepath, "rb") as f:
            for patch in target_def["patches"]:
                absolute_offset = base_offset + patch["app_offset"]

                f.seek(absolute_offset)
                read_bytes = f.read(len(patch["patched"]))

                print(
                    f"  -> Checking '{patch['name']}' at offset 0x{absolute_offset:X}..."
                )

                if read_bytes == patch["patched"]:
                    print(
                        f"    [OK] Patch verified! Expected: {patch['patched'].hex()}, Found: {read_bytes.hex()}"
                    )
                elif read_bytes == patch["original"]:
                    print(
                        f"    [FAIL] Binary is UNPATCHED. Expected: {patch['patched'].hex()}, Found: {read_bytes.hex()} (Original)"
                    )
                    all_passed = False
                else:
                    print(
                        f"    [FAIL] Unknown bytes found. Expected: {patch['patched'].hex()}, Found: {read_bytes.hex()}"
                    )
                    all_passed = False

    except Exception as e:
        print(f"  [!] ERROR: Failed to read binary: {e}")
        return False

    return all_passed


def main():
    base_dir = os.path.dirname(os.path.abspath(__file__))

    print("=" * 60)
    print(" ESP32-S3 Firmware Patch Verification Harness")
    print("=" * 60)

    success_count = 0
    total_count = len(BINARIES_TO_CHECK)

    for relative_path in BINARIES_TO_CHECK:
        full_path = os.path.join(base_dir, relative_path)
        if verify_binary(full_path):
            success_count += 1
        print("-" * 60)

    print(f"Verification Summary: {success_count}/{total_count} binaries passed.")

    if success_count == total_count:
        print("RESULT: ALL PATCHES VERIFIED SUCCESSFULLY.")
        sys.exit(0)
    else:
        print("RESULT: VERIFICATION FAILED FOR ONE OR MORE BINARIES.")
        sys.exit(1)


if __name__ == "__main__":
    main()