// Static Ghidra headless summary script for AgentRE.
// Usage from analyzeHeadless:
//   -postScript AgentRESummary.java /path/to/output.txt

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.DataIterator;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.InstructionIterator;
import ghidra.program.model.listing.Listing;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;

import java.io.File;
import java.io.FileWriter;
import java.io.PrintWriter;

public class AgentRESummary extends GhidraScript {
    private String clean(String s) {
        if (s == null) {
            return "";
        }
        s = s.replace('\n', ' ').replace('\r', ' ');
        if (s.length() > 240) {
            return s.substring(0, 240) + "...";
        }
        return s;
    }

    public void run() throws Exception {
        String[] args = getScriptArgs();
        if (args.length < 1) {
            println("AgentRESummary requires an output path argument");
            return;
        }

        File outFile = new File(args[0]);
        try (PrintWriter out = new PrintWriter(new FileWriter(outFile))) {
            Listing listing = currentProgram.getListing();

            out.println("== Program ==");
            out.println("name: " + currentProgram.getName());
            out.println("format: " + currentProgram.getExecutableFormat());
            out.println("language: " + currentProgram.getLanguageID());
            out.println("compiler: " + currentProgram.getCompilerSpec().getCompilerSpecID());
            out.println("image_base: " + currentProgram.getImageBase());

            out.println();
            out.println("== Memory Blocks ==");
            for (MemoryBlock block : currentProgram.getMemory().getBlocks()) {
                out.println(block.getName()
                        + " start=" + block.getStart()
                        + " end=" + block.getEnd()
                        + " size=" + block.getSize()
                        + " r=" + block.isRead()
                        + " w=" + block.isWrite()
                        + " x=" + block.isExecute());
            }

            out.println();
            out.println("== External Symbols / Imports ==");
            int count = 0;
            SymbolIterator symbols = currentProgram.getSymbolTable().getExternalSymbols();
            while (symbols.hasNext() && count < 600) {
                Symbol sym = symbols.next();
                out.println(sym.getName(true));
                count++;
            }
            if (count >= 600) {
                out.println("[truncated imports at 600]");
            }

            out.println();
            out.println("== Strings ==");
            count = 0;
            DataIterator data = listing.getDefinedData(true);
            while (data.hasNext() && count < 500) {
                Data d = data.next();
                Object value = null;
                try {
                    value = d.getValue();
                } catch (Throwable ignored) {
                    value = null;
                }
                if (value instanceof String) {
                    String text = clean((String) value);
                    if (text.length() >= 4) {
                        out.println(d.getAddress() + " " + text);
                        count++;
                    }
                }
            }
            if (count >= 500) {
                out.println("[truncated strings at 500]");
            }

            out.println();
            out.println("== Functions ==");
            count = 0;
            FunctionIterator funcs = listing.getFunctions(true);
            while (funcs.hasNext() && count < 220) {
                Function f = funcs.next();
                out.println(f.getEntryPoint() + " " + f.getName(true));
                count++;
            }
            if (count >= 220) {
                out.println("[truncated functions at 220]");
            }

            out.println();
            out.println("== Function Instruction Excerpts ==");
            int funcsPrinted = 0;
            funcs = listing.getFunctions(true);
            while (funcs.hasNext() && funcsPrinted < 40) {
                Function f = funcs.next();
                String name = f.getName(true).toLowerCase();
                boolean interesting = name.contains("main")
                        || name.contains("connect")
                        || name.contains("socket")
                        || name.contains("send")
                        || name.contains("recv")
                        || name.contains("exec")
                        || name.contains("crypt")
                        || name.contains("decrypt")
                        || name.contains("attack")
                        || name.contains("scan")
                        || name.contains("loader")
                        || name.contains("start");
                if (!interesting && funcsPrinted >= 12) {
                    continue;
                }
                out.println();
                out.println("-- " + f.getEntryPoint() + " " + f.getName(true) + " --");
                InstructionIterator instrs = listing.getInstructions(f.getBody(), true);
                int insnCount = 0;
                while (instrs.hasNext() && insnCount < 45) {
                    Instruction insn = instrs.next();
                    Address addr = insn.getAddress();
                    out.println("  " + addr + "  " + insn.toString());
                    insnCount++;
                }
                funcsPrinted++;
            }
        }

        println("AgentRE summary written to " + outFile.getAbsolutePath());
    }
}