# Ghidra script to extract .fwimage section from GSP firmware ELF
# Usage: analyzeHeadless ... -postScript extract_fwimage.py

import sys
import os

# Get the current program
program = getCurrentProgram()

# Get the .fwimage section
fwimage_section = program.getSection(".fwimage")
if not fwimage_section:
    print("No .fwimage section found")
    sys.exit(1)

print(f"Found .fwimage section: {fwimage_section.getStart()} - {fwimage_section.getEnd()}")
print(f"Size: {fwimage_section.getLength()} bytes")

# Get the data from the section
data = getBytes(fwimage_section.getStart(), fwimage_section.getLength())

# Save to file
output_dir = "/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP"
os.makedirs(output_dir, exist_ok=True)
fwimage_path = os.path.join(output_dir, "gsp-570.144.fwimage.bin")

with open(fwimage_path, "wb") as f:
    f.write(data)

print(f"Saved .fwimage to {fwimage_path}")

# Also extract .fwversion
fwversion_section = program.getSection(".fwversion")
if fwversion_section:
    version_data = getBytes(fwversion_section.getStart(), fwversion_section.getLength())
    version_str = version_data.decode('ascii').strip('\x00')
    print(f"GSP Firmware Version: {version_str}")
    
    with open(os.path.join(output_dir, "gsp_version.txt"), "w") as f:
        f.write(version_str)

# Extract signatures
for i in range(3):
    sig_section = program.getSection(f".fwsignature{i}")
    if sig_section:
        sig_data = getBytes(sig_section.getStart(), sig_section.getLength())
        sig_path = os.path.join(output_dir, f"gsp-570.144.fwsignature{i}.bin")
        with open(sig_path, "wb") as f:
            f.write(sig_data)
        print(f"Saved signature {i} to {sig_path} ({len(sig_data)} bytes)")

# Create a new Ghidra project for the .fwimage binary
print("Creating Ghidra project for .fwimage binary...")

# We'll do this via a separate script call
print("Done extracting sections.")