// Ghidra script to do basic disassembly analysis of GSP .fwimage binary
// Usage: analyzeHeadless ... -postScript BasicAnalysis.java

import ghidra.app.script.GhidraScript;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.*;
import ghidra.program.model.symbol.*;
import ghidra.program.model.address.*;
import ghidra.util.task.TaskMonitor;
import java.io.*;
import java.util.*;
import java.nio.file.*;

public class BasicAnalysis extends GhidraScript {

    @Override
    protected void run() throws Exception {
        Program program = getCurrentProgram();
        Listing listing = program.getListing();
        Memory memory = program.getMemory();
        FunctionManager functionManager = program.getFunctionManager();

        println("=== GSP Firmware .fwimage Basic Analysis ===");
        println("Program: " + program.getName());
        println("Architecture: " + program.getLanguage().toString());

        // Create output directory
        String outputDir = "/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP";
        Files.createDirectories(Paths.get(outputDir));

        // 1. Check memory blocks
        println("\n--- Memory Blocks ---");
        List<Map<String, Object>> memBlocks = new ArrayList<>();
        for (MemoryBlock block : memory.getBlocks()) {
            Map<String, Object> blockInfo = new HashMap<>();
            blockInfo.put("name", block.getName());
            blockInfo.put("start", block.getStart().toString());
            blockInfo.put("end", block.getEnd().toString());
            blockInfo.put("size", block.getSize());
            blockInfo.put("execute", block.isExecute());
            blockInfo.put("read", block.isRead());
            blockInfo.put("write", block.isWrite());
            memBlocks.add(blockInfo);
            println("  " + block.getName() + ": " + block.getStart() + " - " + block.getEnd() + 
                    " (" + block.getSize() + " bytes) r=" + block.isRead() + " w=" + block.isWrite() + " x=" + block.isExecute());
        }
        writeJson(outputDir + "/memory_blocks.json", memBlocks);

        // 2. Find all functions
        println("\n--- Finding Functions ---");
        List<Map<String, Object>> functions = new ArrayList<>();
        FunctionIterator funcIter = functionManager.getFunctions(true);
        int count = 0;
        while (funcIter.hasNext()) {
            Function func = funcIter.next();
            Map<String, Object> funcInfo = new HashMap<>();
            funcInfo.put("name", func.getName());
            funcInfo.put("entry", func.getEntryPoint().toString());
            funcInfo.put("size", func.getBody().getNumAddresses());
            functions.add(funcInfo);
            count++;
            if (count % 1000 == 0) {
                println("  Found " + count + " functions...");
            }
        }
        println("Total functions: " + functions.size());
        writeJson(outputDir + "/functions.json", functions);

        // 3. Find strings
        println("\n--- Finding Strings ---");
        List<Map<String, String>> strings = new ArrayList<>();
        DataIterator dataIter = listing.getDefinedData(true);
        while (dataIter.hasNext()) {
            Data data = dataIter.next();
            if (data.getDataType().getName().equals("string")) {
                Object value = data.getValue();
                if (value != null && value.toString().length() > 4) {
                    Map<String, String> strInfo = new HashMap<>();
                    strInfo.put("address", data.getAddress().toString());
                    strInfo.put("value", value.toString());
                    strings.add(strInfo);
                }
            }
        }
        println("Found " + strings.size() + " strings");

        // Filter for interesting strings
        String[] interestingKeywords = {"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"};
        List<Map<String, String>> interestingStrings = new ArrayList<>();
        for (Map<String, String> s : strings) {
            String val = s.get("value").toLowerCase();
            for (String kw : interestingKeywords) {
                if (val.contains(kw)) {
                    interestingStrings.add(s);
                    break;
                }
            }
        }
        println("Interesting strings: " + interestingStrings.size());
        writeJson(outputDir + "/interesting_strings.json", interestingStrings);

        // 4. Find entry point
        println("\n--- Entry Point Analysis ---");
        Address entryAddr = program.getMinAddress();
        println("Program entry (min address): " + entryAddr);
        Function funcAtZero = functionManager.getFunctionAt(entryAddr);
        if (funcAtZero != null) {
            println("Function at entry: " + funcAtZero.getName());
        } else {
            println("No function at entry point - creating one...");
            Function func = functionManager.createFunction("entry_point", entryAddr, null, null, monitor);
            if (func != null) {
                println("Created function: " + func.getName());
            }
        }

        println("\n--- Analysis Complete ---");
        println("Results saved to " + outputDir);
    }

    private void writeJson(String path, Object data) throws IOException {
        StringBuilder json = new StringBuilder();
        json.append("[\n");
        if (data instanceof List) {
            List<?> list = (List<?>) data;
            for (int i = 0; i < list.size(); i++) {
                Object item = list.get(i);
                if (item instanceof Map) {
                    Map<?, ?> map = (Map<?, ?>) item;
                    json.append("  {");
                    int j = 0;
                    for (Map.Entry<?, ?> entry : map.entrySet()) {
                        json.append("\"").append(entry.getKey()).append("\": ");
                        Object val = entry.getValue();
                        if (val instanceof String) {
                            json.append("\"").append(escapeJson(val.toString())).append("\"");
                        } else {
                            json.append(val);
                        }
                        if (j < map.size() - 1) json.append(", ");
                        j++;
                    }
                    json.append("}");
                }
                if (i < list.size() - 1) json.append(",");
                json.append("\n");
            }
        }
        json.append("]");
        Files.writeString(Paths.get(path), json.toString());
    }

    private String escapeJson(String s) {
        return s.replace("\\", "\\\\").replace("\"", "\\\"").replace("\n", "\\n").replace("\r", "\\r").replace("\t", "\\t");
    }
}