WsolverPostAnalysis.java

// WsolverPostAnalysis.java — Ghidra headless post-analysis verification script (Java)
//
// Purpose:
//   For each KLEE hit with violations > 0 in wsolver_report.json, decompile
//   the function in Ghidra and verify whether the violation is plausible.
//   Emits ghidra_post.json; wsolve merges the verdicts back into the report.
//
// Invocation (via wghidra):
//   analyzeHeadless /tmp/proj ProjName -import <binary>  \
//     -postScript WsolverPostAnalysis.java <report_json> <out_json> \
//     -scriptPath <dir> -deleteProject
//
// Verdicts per hit:
//   CONFIRMED      — dangerous arg is user-tainted; KLEE finding is plausible
//   FALSE_POSITIVE — dangerous arg is bounded/constant; KLEE exploited harness
//   UNKNOWN        — could not determine; leave for human triage
//
//@category WSolver

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.listing.*;
import ghidra.program.model.symbol.*;
import ghidra.util.task.ConsoleTaskMonitor;

import java.io.*;
import java.nio.file.*;
import java.util.*;
import java.util.regex.*;

public class WsolverPostAnalysis extends GhidraScript {

    private static final Set<String> USER_INPUT_SOURCES = new HashSet<>(Arrays.asList(
        "read", "fread", "fgets", "getc", "getchar", "gets",
        "recv", "recvfrom", "recvmsg", "mq_receive", "msgrcv",
        "getenv", "getopt", "getopt_long",
        "atoi", "atol", "atoll", "strtol", "strtoul", "strtoll", "strtoull",
        "sscanf", "scanf", "fscanf",
        "accept", "connect"
    ));

    private static final Set<String> BOUNDED_SOURCES = new HashSet<>(Arrays.asList(
        "strlen", "strnlen", "wcslen",
        "stat", "fstat", "lstat",
        "ftell", "lseek", "fseeko",
        "getpid", "getuid", "geteuid", "getgid",
        "time", "clock", "clock_gettime",
        "malloc", "calloc", "realloc"
    ));

    private static final Map<String, Integer> SINK_ARG_INDEX = new HashMap<>();
    static {
        SINK_ARG_INDEX.put("strcpy",   1);  SINK_ARG_INDEX.put("strcat",   1);
        SINK_ARG_INDEX.put("gets",     0);  SINK_ARG_INDEX.put("sprintf",  1);
        SINK_ARG_INDEX.put("vsprintf", 1);  SINK_ARG_INDEX.put("snprintf", 2);
        SINK_ARG_INDEX.put("vsnprintf",2);  SINK_ARG_INDEX.put("memcpy",   2);
        SINK_ARG_INDEX.put("memmove",  2);  SINK_ARG_INDEX.put("bcopy",    2);
        SINK_ARG_INDEX.put("strncpy",  2);  SINK_ARG_INDEX.put("strncat",  2);
        SINK_ARG_INDEX.put("malloc",   0);  SINK_ARG_INDEX.put("calloc",   0);
        SINK_ARG_INDEX.put("realloc",  1);  SINK_ARG_INDEX.put("alloca",   0);
        SINK_ARG_INDEX.put("system",   0);  SINK_ARG_INDEX.put("popen",    0);
        SINK_ARG_INDEX.put("execve",   1);  SINK_ARG_INDEX.put("printf",   0);
        SINK_ARG_INDEX.put("fprintf",  1);  SINK_ARG_INDEX.put("vprintf",  0);
        SINK_ARG_INDEX.put("vfprintf", 1);
    }

    // ── Hit record ────────────────────────────────────────────────────────────
    private static class KleeHit {
        String funcName;   // e.g. "function_16f4e"
        String sinkName;   // e.g. "memcpy"
    }

    @Override
    public void run() throws Exception {
        String[] args      = getScriptArgs();
        String reportPath  = args.length > 0 ? args[0] : "/tmp/wsolver_report.json";
        String outPath     = args.length > 1 ? args[1] : "/tmp/ghidra_post.json";

        ConsoleTaskMonitor monitor = new ConsoleTaskMonitor();

        // Parse KLEE hits from wsolver_report.json
        List<KleeHit> hits = parseKleeHits(reportPath);
        if (hits.isEmpty()) {
            println("[wghidra-post] No KLEE violations found in " + reportPath);
            writeOutput(outPath, currentProgram.getExecutablePath(),
                        new LinkedHashMap<>());
            return;
        }
        println("[wghidra-post] Verifying " + hits.size() + " KLEE hit(s)...");

        // Set up decompiler
        DecompInterface decomp = new DecompInterface();
        decomp.openProgram(currentProgram);

        Map<String, String> results = new LinkedHashMap<>();

        for (KleeHit hit : hits) {
            Long addrLong = parseHexFromFuncName(hit.funcName);
            if (addrLong == null) {
                println("[wghidra-post]   " + hit.funcName + " — cannot parse address");
                results.put(hit.funcName,
                    makeEntry(hit.funcName, null, hit.sinkName,
                              "UNKNOWN", "ADDR_PARSE_FAILED", "", "", ""));
                continue;
            }

            // Try to find function, with common base-address deltas as fallback
            Function func = findFunction(addrLong);
            if (func == null) {
                println("[wghidra-post]   " + hit.funcName + " @ 0x" +
                        Long.toHexString(addrLong) + " — function not found");
                results.put(hit.funcName,
                    makeEntry(hit.funcName, "0x" + Long.toHexString(addrLong),
                              hit.sinkName, "UNKNOWN", "FUNCTION_NOT_FOUND", "", "", ""));
                continue;
            }

            println("[wghidra-post]   " + hit.funcName + " → " + func.getName() +
                    " (" + hit.sinkName + ")");

            // Decompile
            String cCode = decompileFunction(decomp, func, monitor);
            if (cCode == null) {
                results.put(hit.funcName,
                    makeEntry(hit.funcName, "0x" + Long.toHexString(addrLong),
                              hit.sinkName, "UNKNOWN", "DECOMPILE_FAILED", "", "", ""));
                continue;
            }

            // Collect callees
            Set<String> callees = getFunctionCallees(func);

            // Determine argument index
            int argIndex = SINK_ARG_INDEX.getOrDefault(hit.sinkName, 0);

            // Extract argument expression
            String argExpr  = extractSinkArg(cCode, hit.sinkName, argIndex);
            String snippet  = extractSinkLine(cCode, hit.sinkName);

            // Verify
            String[] verdict = verify(argExpr, callees, cCode);
            String   v       = verdict[0]; // CONFIRMED | FALSE_POSITIVE | UNKNOWN
            String   reason  = verdict[1];
            String   sources = verdict[2];

            println("[wghidra-post]     " + hit.sinkName + " → " + v +
                    " (" + reason + ")  arg=" +
                    (argExpr.length() > 60 ? argExpr.substring(0,60)+"..." : argExpr));

            results.put(hit.funcName,
                makeEntry(hit.funcName, "0x" + Long.toHexString(addrLong),
                          hit.sinkName, v, reason,
                          argExpr.length() > 300 ? argExpr.substring(0,300) : argExpr,
                          sources,
                          snippet.length() > 200 ? snippet.substring(0,200) : snippet));
        }

        decomp.dispose();

        long confirmed = results.values().stream()
            .filter(s -> s.contains("\"verdict\":\"CONFIRMED\"")).count();
        long fp = results.values().stream()
            .filter(s -> s.contains("\"verdict\":\"FALSE_POSITIVE\"")).count();
        long unk = results.values().stream()
            .filter(s -> s.contains("\"verdict\":\"UNKNOWN\"")).count();

        println("[wghidra-post] Results: " + confirmed + " confirmed, " +
                fp + " false positive(s), " + unk + " unknown");

        writeOutput(outPath, currentProgram.getExecutablePath(), results);
        println("[wghidra-post] Wrote " + outPath);
    }

    // ── Parse KLEE hits from wsolver_report.json ──────────────────────────────

    private List<KleeHit> parseKleeHits(String reportPath) throws Exception {
        List<KleeHit> hits = new ArrayList<>();
        String json = new String(Files.readAllBytes(Paths.get(reportPath)));

        // We only want hits inside klee jobs with violations > 0.
        // Strategy: find each "solver":"klee" job block, then within it find
        // "hits" arrays, then entries with "violations": N > 0.
        //
        // Simple state-machine parse — no external JSON lib.

        // Find all "func":"..." + "sink":"..." pairs inside hits with violations > 0
        // We look for blocks: { ... "violations": N ... } where N > 0
        Pattern hitBlock = Pattern.compile(
            "\\{[^{}]*\"func\"\\s*:\\s*\"([^\"]+)\"[^{}]*\"sink\"\\s*:\\s*\"([^\"]+)\"" +
            "[^{}]*\"violations\"\\s*:\\s*(\\d+)[^{}]*\\}");
        Matcher m = hitBlock.matcher(json);
        while (m.find()) {
            int violations = Integer.parseInt(m.group(3));
            if (violations > 0) {
                KleeHit h = new KleeHit();
                h.funcName = m.group(1);
                h.sinkName = m.group(2);
                hits.add(h);
            }
        }
        return hits;
    }

    // ── Address resolution ────────────────────────────────────────────────────

    private Long parseHexFromFuncName(String name) {
        // Handles "function_16f4e" and "FUN_00016f4e"
        Matcher m = Pattern.compile("[_]([0-9a-fA-F]{4,})$").matcher(name);
        if (m.find()) {
            try { return Long.parseLong(m.group(1), 16); }
            catch (NumberFormatException e) { /* fall through */ }
        }
        return null;
    }

    private Function findFunction(long addrInt) {
        long[] candidates = {
            addrInt,
            addrInt + 0x400000L,
            addrInt - 0x400000L,
            addrInt + 0x10000L,
            addrInt - 0x10000L,
        };
        for (long candidate : candidates) {
            try {
                Address addr = currentProgram.getAddressFactory()
                    .getDefaultAddressSpace().getAddress(candidate);
                Function f = getFunctionAt(addr);
                if (f != null) return f;
                f = getFunctionContaining(addr);
                if (f != null) return f;
            } catch (Exception e) { /* try next */ }
        }
        return null;
    }

    // ── Analysis helpers ──────────────────────────────────────────────────────

    private Set<String> getFunctionCallees(Function func) {
        Set<String> callees = new HashSet<>();
        Listing listing = currentProgram.getListing();
        InstructionIterator it = listing.getInstructions(func.getBody(), true);
        while (it.hasNext()) {
            Instruction instr = it.next();
            for (Reference ref : getReferencesFrom(instr.getAddress())) {
                if (!ref.getReferenceType().isCall()) continue;
                Address dest = ref.getToAddress();
                Function callee = getFunctionAt(dest);
                if (callee != null) {
                    callees.add(stripPlt(callee.getName()));
                } else {
                    Symbol sym = getSymbolAt(dest);
                    if (sym != null) callees.add(stripPlt(sym.getName()));
                }
            }
        }
        return callees;
    }

    private String stripPlt(String name) {
        int at = name.indexOf('@');
        return at >= 0 ? name.substring(0, at) : name;
    }

    private String decompileFunction(DecompInterface decomp, Function func,
                                     ConsoleTaskMonitor monitor) {
        try {
            DecompileResults res = decomp.decompileFunction(func, 30, monitor);
            if (res != null && res.decompileCompleted() && res.getDecompiledFunction() != null)
                return res.getDecompiledFunction().getC();
        } catch (Exception e) { /* fall through */ }
        return null;
    }

    private String extractSinkArg(String cCode, String sinkName, int argIndex) {
        Pattern p = Pattern.compile("\\b" + Pattern.quote(sinkName) + "\\s*\\(");
        Matcher m = p.matcher(cCode);
        if (!m.find()) return "";

        int start = m.end();
        List<String> args = new ArrayList<>();
        int depth = 1;
        StringBuilder current = new StringBuilder();
        for (int i = start; i < cCode.length() && depth > 0; i++) {
            char ch = cCode.charAt(i);
            if      (ch == '(')                   { depth++; current.append(ch); }
            else if (ch == ')')                   { depth--; if (depth > 0) current.append(ch); }
            else if (ch == ',' && depth == 1)     { args.add(current.toString().trim());
                                                    current = new StringBuilder(); }
            else                                  { current.append(ch); }
        }
        if (current.length() > 0) args.add(current.toString().trim());
        return argIndex < args.size() ? args.get(argIndex) : "";
    }

    private String extractSinkLine(String cCode, String sinkName) {
        for (String line : cCode.split("\n")) {
            if (line.contains(sinkName + "(")) return line.trim();
        }
        return "";
    }

    /** Returns [verdict, reason, sources] */
    private String[] verify(String argExpr, Set<String> callees, String cCode) {
        String lower = argExpr.toLowerCase();

        // ── Hard CONFIRMED indicators ────────────────────────────────────────
        Set<String> taintHits = new HashSet<>(callees);
        taintHits.retainAll(USER_INPUT_SOURCES);
        if (!taintHits.isEmpty())
            return new String[]{"CONFIRMED", "USER_TAINTED", join(taintHits)};

        if (lower.contains("argv"))
            return new String[]{"CONFIRMED", "ARGV_DERIVED", "argv"};

        // ── Hard FALSE_POSITIVE indicators ───────────────────────────────────

        // Constant numeric literal
        if (argExpr.matches("\\s*0x[0-9a-fA-F]+\\s*") || argExpr.matches("\\s*\\d+\\s*"))
            return new String[]{"FALSE_POSITIVE", "CONSTANT_SIZE", "constant"};

        // String literal constant: "TZ", "PATH", etc.
        if (argExpr.matches("\\s*\"[^\"]*\"\\s*"))
            return new String[]{"FALSE_POSITIVE", "STRING_LITERAL", "string_literal"};

        // sizeof
        if (lower.contains("sizeof"))
            return new String[]{"FALSE_POSITIVE", "SIZEOF_BOUNDED", "sizeof"};

        // Pointer / array subtraction, optionally scaled: (end - start) * 8
        boolean hasSubtraction = argExpr.contains(" - ");
        boolean hasLocalVar    = lower.contains("local_") || lower.contains("uvar")
                              || lower.contains("param")  || lower.contains("ptrvar");
        if (hasSubtraction && hasLocalVar)
            return new String[]{"FALSE_POSITIVE", "POINTER_SUBTRACTION", "pointer_subtraction"};

        // Bounded source in expression text
        for (String b : BOUNDED_SOURCES) {
            if (lower.contains(b))
                return new String[]{"FALSE_POSITIVE", "BOUNDED_SOURCE", b};
        }

        // Struct field access: *(type*)(param + n)
        if (argExpr.matches(".*\\*\\s*\\(\\w+\\s*\\*\\s*\\)\\s*\\(\\w+.*[+\\-].*"))
            return new String[]{"FALSE_POSITIVE", "STRUCT_FIELD", "struct_field"};

        // Bounded callees in function (but no taint sources)
        Set<String> boundedHits = new HashSet<>(callees);
        boundedHits.retainAll(BOUNDED_SOURCES);
        if (!boundedHits.isEmpty())
            return new String[]{"FALSE_POSITIVE", "BOUNDED_CALLEE", join(boundedHits)};

        return new String[]{"UNKNOWN", "UNDETERMINED", ""};
    }

    private String join(Set<String> s) {
        StringBuilder sb = new StringBuilder();
        for (String v : s) { if (sb.length() > 0) sb.append(","); sb.append(v); }
        return sb.toString();
    }

    // ── JSON output ───────────────────────────────────────────────────────────

    private String makeEntry(String funcName, String address, String sink,
                             String verdict, String reason,
                             String argExpr, String sources, String snippet) {
        return "{"
            + "\"address\":"           + jsonStr(address)  + ","
            + "\"sink\":"              + jsonStr(sink)      + ","
            + "\"verdict\":"           + jsonStr(verdict)   + ","
            + "\"reason\":"            + jsonStr(reason)    + ","
            + "\"arg_expr\":"          + jsonStr(argExpr)   + ","
            + "\"taint_sources\":"     + jsonStr(sources)   + ","
            + "\"decompiled_snippet\":" + jsonStr(snippet)  + "}";
    }

    private String jsonStr(String s) {
        if (s == null) return "\"\"";
        return "\"" + s.replace("\\", "\\\\")
                       .replace("\"", "\\\"")
                       .replace("\n", "\\n")
                       .replace("\r", "") + "\"";
    }

    private void writeOutput(String outPath, String binaryPath,
                             Map<String, String> results) throws Exception {
        StringBuilder sb = new StringBuilder();
        sb.append("{\n");
        sb.append("  \"binary\": ").append(jsonStr(binaryPath)).append(",\n");
        sb.append("  \"hits\": {\n");
        int i = 0;
        for (Map.Entry<String, String> e : results.entrySet()) {
            sb.append("    ").append(jsonStr(e.getKey()))
              .append(": ").append(e.getValue());
            if (++i < results.size()) sb.append(",");
            sb.append("\n");
        }
        sb.append("  }\n}\n");
        Files.write(Paths.get(outPath), sb.toString().getBytes());
    }
}