// Exports a Ghidra analysis for CodeCraft: program info, function index,
// pseudo-C per function, strings, imports and exports.
// Usage: analyzeHeadless ... -postScript CodeCraftExport.java <output dir> [max functions]
//@category CodeCraft

import ghidra.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompileResults;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.address.AddressIterator;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Program;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;
import ghidra.program.util.DefinedStringIterator;

import java.io.File;
import java.io.PrintWriter;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;

public class CodeCraftExport extends GhidraScript {

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

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        File out = new File(args[0]);
        int maxFunctions = args.length > 1 ? Integer.parseInt(args[1]) : Integer.MAX_VALUE;
        File funcDir = new File(out, "functions");
        funcDir.mkdirs();
        Program p = currentProgram;

        try (PrintWriter info = new PrintWriter(new File(out, "info.txt"), StandardCharsets.UTF_8)) {
            info.println("name: " + p.getName());
            info.println("format: " + p.getExecutableFormat());
            info.println("language: " + p.getLanguageID());
            info.println("compiler: " + p.getCompilerSpec().getCompilerSpecID());
            info.println("image_base: " + p.getImageBase());
            info.println("functions: " + p.getFunctionManager().getFunctionCount());
        }

        DecompInterface di = new DecompInterface();
        di.openProgram(p);
        int n = 0;
        try (PrintWriter idx = new PrintWriter(new File(out, "functions.tsv"), StandardCharsets.UTF_8)) {
            idx.println("address\tname\tsize\tfile\tsignature");
            FunctionIterator it = p.getFunctionManager().getFunctions(true);
            while (it.hasNext() && !monitor.isCancelled() && n < maxFunctions) {
                Function f = it.next();
                if (f.isThunk() || f.isExternal()) {
                    continue;
                }
                String base = f.getEntryPoint().toString() + "_" + f.getName().replaceAll("[^A-Za-z0-9_]", "_");
                if (base.length() > 120) {
                    base = base.substring(0, 120);
                }
                DecompileResults r = di.decompileFunction(f, 60, monitor);
                String code = (r != null && r.decompileCompleted())
                        ? r.getDecompiledFunction().getC()
                        : "/* decompilation failed: " + (r == null ? "no result" : r.getErrorMessage()) + " */\n";
                Files.writeString(new File(funcDir, base + ".c").toPath(),
                        "// " + f.getName() + " @ " + f.getEntryPoint() + "\n" + code, StandardCharsets.UTF_8);
                idx.println(f.getEntryPoint() + "\t" + f.getName() + "\t" + f.getBody().getNumAddresses()
                        + "\tfunctions/" + base + ".c\t" + clean(f.getSignature().getPrototypeString()));
                n++;
                if (n % 200 == 0) {
                    println("CodeCraft: decompiled " + n + " functions");
                }
            }
        }
        di.dispose();

        try (PrintWriter sw = new PrintWriter(new File(out, "strings.tsv"), StandardCharsets.UTF_8)) {
            sw.println("address\tstring");
            for (Data d : DefinedStringIterator.forProgram(p)) {
                Object v = d.getValue();
                if (v != null) {
                    sw.println(d.getAddress() + "\t" + clean(v.toString()));
                }
            }
        }

        try (PrintWriter iw = new PrintWriter(new File(out, "imports.tsv"), StandardCharsets.UTF_8)) {
            iw.println("library\tname");
            SymbolIterator si = p.getSymbolTable().getExternalSymbols();
            while (si.hasNext()) {
                Symbol s = si.next();
                iw.println(s.getParentNamespace().getName() + "\t" + s.getName());
            }
        }

        try (PrintWriter ew = new PrintWriter(new File(out, "exports.tsv"), StandardCharsets.UTF_8)) {
            ew.println("name\taddress");
            AddressIterator ai = p.getSymbolTable().getExternalEntryPointIterator();
            while (ai.hasNext()) {
                Address a = ai.next();
                Symbol s = p.getSymbolTable().getPrimarySymbol(a);
                ew.println((s == null ? "?" : s.getName()) + "\t" + a);
            }
        }

        Files.writeString(new File(out, ".done").toPath(), String.valueOf(n), StandardCharsets.UTF_8);
        println("CodeCraft: export complete, " + n + " functions");
    }
}