# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# Ghidra script to analyze GSP .fwimage binary
# Usage: analyzeHeadless ... -postScript analyze_gsp_fwimage.py

import sys
import os
import json

# Get the current program
program = getCurrentProgram()
listing = program.getListing()
memory = program.getMemory()
function_manager = program.getFunctionManager()

print("=== GSP Firmware .fwimage Analysis ===")
print(f"Program: {program.getName()}")
print(f"Architecture: {program.getLanguage().toString()}")

# Create output directory
output_dir = "/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP"
os.makedirs(output_dir, exist_ok=True)

# 1. Find all functions
print("\n--- Finding Functions ---")
functions = []
func_iter = function_manager.getFunctions(True)
while func_iter.hasNext():
    func = func_iter.next()
    functions.append({
        "name": func.getName(),
        "entry": str(func.getEntryPoint()),
        "size": func.getBody().getNumAddresses(),
        "signature": func.getSignature().toString() if func.getSignature() else ""
    })

print(f"Found {len(functions)} functions")

# Save function list
with open(os.path.join(output_dir, "functions.json"), "w") as f:
    json.dump(functions, f, indent=2)

# 2. Find RPC-related functions (vRPC/pRPC handlers)
print("\n--- Finding RPC Handlers ---")
rpc_keywords = [
    "rpc", "Rpc", "RPC",
    "vgpu", "Vgpu", "VGPU",
    "vmiop", "Vmiop", "VMIOP",
    "hal", "Hal", "HAL",
    "bar", "Bar", "BAR",
    "migrate", "Migrate", "MIGRATE",
    "init", "Init", "INIT",
    "boot", "Boot", "BOOT",
    "kgsp", "Kgsp", "KGSP"
]

rpc_functions = []
for func in functions:
    name_lower = func["name"].lower()
    for kw in rpc_keywords:
        if kw.lower() in name_lower:
            rpc_functions.append(func)
            break

print(f"Found {len(rpc_functions)} potential RPC-related functions")
with open(os.path.join(output_dir, "rpc_functions.json"), "w") as f:
    json.dump(rpc_functions, f, indent=2)

# 3. Find strings related to vGPU initialization
print("\n--- Finding vGPU-related Strings ---")
vgpu_strings = []
string_iter = listing.getDefinedData(True)
while string_iter.hasNext():
    data = string_iter.next()
    if data.getDataType().getName() == "string":
        value = data.getValue()
        if value:
            val_str = str(value).lower()
            if any(kw in val_str for kw in ["vgpu", "vmiop", "rpc", "hal", "bar", "migrate", "init", "boot", "kgsp", "wpr", "firmware", "version"]):
                vgpu_strings.append({
                    "address": str(data.getAddress()),
                    "value": str(value)
                })

print(f"Found {len(vgpu_strings)} vGPU-related strings")
with open(os.path.join(output_dir, "vgpu_strings.json"), "w") as f:
    json.dump(vgpu_strings, f, indent=2)

# 4. Find entry points and initialization functions
print("\n--- Finding Entry Points ---")
entry_points = []
# Look for common entry point names
entry_names = [
    "kgspInitRm_IMPL", "kgspInitRm", "kgspBootGspRm", "kgspBootGspRm_IMPL",
    "gspBoot", "gspInit", "gspMain",
    "nv_vgpu_vfio_get_ops", "nv_vgpu_vfio_set_ops",
    "vgpu_init", "vgpu_init_IMPL",
    "vmiop_init", "vmiop_init_IMPL",
    "hal_init", "hal_init_IMPL"
]

for name in entry_names:
    func = function_manager.getFunction(name)
    if func:
        entry_points.append({
            "name": name,
            "entry": str(func.getEntryPoint()),
            "size": func.getBody().getNumAddresses()
        })
        print(f"  Found: {name} at {func.getEntryPoint()}")

# Also check for functions called from reset vector / entry
# Look at the first few instructions
entry_addr = program.getMinAddress()
print(f"\nProgram entry address: {entry_addr}")

# 5. Find WPR (Write Protected Region) related code
print("\n--- Finding WPR-related Code ---")
wpr_functions = []
for func in functions:
    name_lower = func["name"].lower()
    if "wpr" in name_lower or "write_protect" in name_lower or "protect" in name_lower:
        wpr_functions.append(func)

print(f"Found {len(wpr_functions)} WPR-related functions")
with open(os.path.join(output_dir, "wpr_functions.json"), "w") as f:
    json.dump(wpr_functions, f, indent=2)

# 6. Find memory management functions
print("\n--- Finding Memory Management Functions ---")
mem_keywords = ["alloc", "free", "map", "unmap", "pin", "unpin", "dma", "iommu", "mmio", "bar"]
mem_functions = []
for func in functions:
    name_lower = func["name"].lower()
    for kw in mem_keywords:
        if kw in name_lower:
            mem_functions.append(func)
            break

print(f"Found {len(mem_functions)} memory management functions")
with open(os.path.join(output_dir, "mem_functions.json"), "w") as f:
    json.dump(mem_functions, f, indent=2)

# 7. Find interrupt/MSI-X related functions
print("\n--- Finding Interrupt/MSI-X Functions ---")
irq_keywords = ["irq", "interrupt", "msi", "msix", "eventfd", "doorbell"]
irq_functions = []
for func in functions:
    name_lower = func["name"].lower()
    for kw in irq_keywords:
        if kw in name_lower:
            irq_functions.append(func)
            break

print(f"Found {len(irq_functions)} interrupt-related functions")
with open(os.path.join(output_dir, "irq_functions.json"), "w") as f:
    json.dump(irq_functions, f, indent=2)

# 8. Generate summary report
print("\n--- Generating Summary Report ---")
report = {
    "firmware": "gsp-570.144",
    "architecture": program.getLanguage().toString(),
    "total_functions": len(functions),
    "rpc_functions": len(rpc_functions),
    "vgpu_strings": len(vgpu_strings),
    "entry_points": entry_points,
    "wpr_functions": len(wpr_functions),
    "mem_functions": len(mem_functions),
    "irq_functions": len(irq_functions)
}

with open(os.path.join(output_dir, "analysis_summary.json"), "w") as f:
    json.dump(report, f, indent=2)

# Print summary
print(json.dumps(report, indent=2))

# 9. Decompile key functions and save
print("\n--- Decompiling Key Functions ---")
from ghidra.app.decompiler import DecompInterface
from ghidra.util.task import ConsoleTaskMonitor

decompiler = DecompInterface()
decompiler.openProgram(program)
monitor = ConsoleTaskMonitor()

key_functions = entry_points + rpc_functions[:20] + wpr_functions + mem_functions[:10] + irq_functions[:10]
decompiled = {}

for func_info in key_functions:
    func = function_manager.getFunctionAt(func_info["entry"])
    if func:
        result = decompiler.decompileFunction(func, 60, monitor)
        if result.decompileCompleted():
            decompiled[func_info["name"]] = result.getDecompiledFunction().getC()
            print(f"  Decompiled: {func_info['name']}")
        else:
            print(f"  Failed to decompile: {func_info['name']}")

# Save decompiled code
with open(os.path.join(output_dir, "decompiled_key_functions.c"), "w") as f:
    for name, code in decompiled.items():
        f.write(f"\n// ===== {name} =====\n")
        f.write(code)
        f.write("\n")

print(f"\nDecompiled {len(decompiled)} key functions")
print("Analysis complete!")