# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# Ghidra script to do basic disassembly analysis of GSP .fwimage binary
# Usage: analyzeHeadless ... -postScript basic_analysis.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 Basic 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. Check memory blocks
print("\n--- Memory Blocks ---")
mem_blocks = []
for block in memory.getBlocks():
    mem_blocks.append({
        "name": block.getName(),
        "start": str(block.getStart()),
        "end": str(block.getEnd()),
        "size": block.getSize(),
        "execute": block.isExecute(),
        "read": block.isRead(),
        "write": block.isWrite()
    })
    print(f"  {block.getName()}: {block.getStart()} - {block.getEnd()} ({block.getSize()} bytes) r={block.isRead()} w={block.isWrite()} x={block.isExecute()}")

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

# 2. Find all functions (but don't decompile yet)
print("\n--- Finding Functions ---")
functions = []
func_iter = function_manager.getFunctions(True)
count = 0
while func_iter.hasNext():
    func = func_iter.next()
    functions.append({
        "name": func.getName(),
        "entry": str(func.getEntryPoint()),
        "size": func.getBody().getNumAddresses()
    })
    count += 1
    if count % 1000 == 0:
        print(f"  Found {count} functions...")

print(f"Total functions: {len(functions)}")

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

# 3. Find strings
print("\n--- Finding Strings ---")
strings = []
string_iter = listing.getDefinedData(True)
while string_iter.hasNext():
    data = string_iter.next()
    if data.getDataType().getName() == "string":
        value = data.getValue()
        if value and len(str(value)) > 4:
            strings.append({
                "address": str(data.getAddress()),
                "value": str(value)
            })

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

# Filter for interesting strings
interesting_keywords = ["vgpu", "vmiop", "rpc", "hal", "bar", "migrate", "init", "boot", "kgsp", "wpr", "firmware", "version", "error", "fail", "assert", "panic", "nvidia", "rmapi", "gsp", "falcon", "mailbox", "interrupt", "dma", "iommu", "page", "mmu", "pde", "pte"]
interesting_strings = [s for s in strings if any(kw in s['value'].lower() for kw in interesting_keywords)]
print(f"Interesting strings: {len(interesting_strings)}")

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

# 4. Find entry point (offset 0)
print("\n--- Entry Point Analysis ---")
entry_addr = program.getMinAddress()
print(f"Program entry (min address): {entry_addr}")

# Check what's at address 0
func_at_zero = function_manager.getFunctionAt(entry_addr)
if func_at_zero:
    print(f"Function at entry: {func_at_zero.getName()}")
else:
    print("No function at entry point - creating one...")
    # Create function at entry point
    func = function_manager.createFunction("entry_point", entry_addr, None, None)
    if func:
        print(f"Created function: {func.getName()}")

# 5. Look for RISC-V standard function prologues
print("\n--- Looking for Function Prologues ---")
# RISC-V typical prologue: addi sp, sp, -X; sd ra, X-8(sp); sd s0, X-16(sp); addi s0, sp, X
# Let's search for common patterns

# 6. Find cross-references to interesting addresses
print("\n--- Analysis Complete ---")
print(f"Results saved to {output_dir}")