WsolverPreAnalysis.java

// WsolverPreAnalysis.java — Ghidra headless pre-analysis script (Java)
//
// Purpose:
//   For each dangerous sink call site in the binary, determine whether the
//   critical argument is reachable from user-controlled input.
//   Emits ghidra_pre.json which wsolver reads via --ghidra-pre to skip
//   BOUNDED targets before running KLEE.
//
// Invocation (via wghidra):
//   analyzeHeadless /tmp/proj ProjName -import <binary>  \
//     -postScript WsolverPreAnalysis.java <sinks_json> <out_json> \
//     -scriptPath <dir> -deleteProject
//
// Taint classifications:
//   USER_TAINTED — dangerous arg reachable from read/fgets/argv/recv/getenv/...
//   BOUNDED      — dangerous arg derived from strlen/stat/constant/ptr subtraction
//   UNKNOWN      — could not determine; KLEE will still run on these
//
//@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 WsolverPreAnalysis extends GhidraScript {

    // ── User-input source functions ───────────────────────────────────────────
    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"
    ));

    // ── Bounded (non-user-controlled) source functions ────────────────────────
    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"
    ));

    // ── Sink → dangerous argument index (0-based) ─────────────────────────────
    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("execl",    1);
        SINK_ARG_INDEX.put("execlp",   1);
        SINK_ARG_INDEX.put("execvp",   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);
    }

    // ── Main entry point ──────────────────────────────────────────────────────

    @Override
    public void run() throws Exception {
        String[] args       = getScriptArgs();
        String sinksPath    = args.length > 0 ? args[0] : "/usr/local/bin/sinks.json";
        String outPath      = args.length > 1 ? args[1] : "/tmp/ghidra_pre.json";

        ConsoleTaskMonitor monitor = new ConsoleTaskMonitor();

        // Load sink names from sinks.json
        Set<String> sinkNames = loadSinkNames(sinksPath);
        println("[wghidra-pre] Loaded " + sinkNames.size() + " sink names from " + sinksPath);

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

        FunctionManager funcMgr    = currentProgram.getFunctionManager();
        String          binaryPath = currentProgram.getExecutablePath();
        int             funcCount  = funcMgr.getFunctionCount();

        println("[wghidra-pre] Scanning " + funcCount + " functions for sink call sites...");

        // results: hex_addr → JSON object string
        Map<String, String> results = new LinkedHashMap<>();

        FunctionIterator funcIter = funcMgr.getFunctions(true);
        while (funcIter.hasNext()) {
            Function func = funcIter.next();
            String   fname    = func.getName();
            long     addrInt  = func.getEntryPoint().getOffset();
            String   addrHex  = "0x" + Long.toHexString(addrInt);

            // Collect all callee names for this function
            Set<String> callees = getFunctionCallees(func);

            // Check whether this function calls any dangerous sink
            Set<String> sinksHit = new HashSet<>(callees);
            sinksHit.retainAll(sinkNames);
            if (sinksHit.isEmpty()) continue;

            // Pick the first matching sink (highest priority handled via sinkNames order)
            String sinkName = pickHighestPrioritySink(sinksHit);
            int    argIndex = SINK_ARG_INDEX.getOrDefault(sinkName, 0);

            // Skip PLT stubs and external thunks — they have no real body to
            // analyse; their "callees" are just the function itself.  Ghidra
            // decompiles them as a single-expression stub showing the libc
            // signature (e.g. "size_t __n"), which produces meaningless results.
            if (func.isThunk() || func.isExternal()) {
                println("[wghidra-pre]   " + fname + " @ " + addrHex
                        + " — skipping PLT/external thunk");
                continue;
            }

            println("[wghidra-pre]   " + fname + " @ " + addrHex + " calls " + sinkName);

            // Decompile
            String cCode = decompileFunction(decomp, func, monitor);
            if (cCode == null) {
                results.put(addrHex, makeEntry(fname, sinkName, argIndex,
                    "UNKNOWN", "", "decompile_failed"));
                continue;
            }

            // Extract the dangerous argument expression
            String argExpr = extractSinkArg(cCode, sinkName, argIndex);

            // Classify
            String[]  taintResult = classify(argExpr, callees);
            String    taint       = taintResult[0];
            String    sources     = taintResult[1];

            // If UNKNOWN, walk one level of callers to check for user input
            if ("UNKNOWN".equals(taint)) {
                outer:
                for (Reference ref : getReferencesTo(func.getEntryPoint())) {
                    if (!ref.getReferenceType().isCall()) continue;
                    Function caller = getFunctionContaining(ref.getFromAddress());
                    if (caller == null) continue;
                    Set<String> callerCallees = getFunctionCallees(caller);
                    String[] r2 = classify("", callerCallees);
                    if ("USER_TAINTED".equals(r2[0])) {
                        taint   = "USER_TAINTED";
                        sources = r2[1];
                        break outer;
                    }
                }
            }

            // Truncate argExpr
            if (argExpr.length() > 200) argExpr = argExpr.substring(0, 200);

            results.put(addrHex, makeEntry(fname, sinkName, argIndex, taint, argExpr, sources));
        }

        decomp.dispose();
        writeOutput(outPath, binaryPath, results);
        println("[wghidra-pre] Wrote " + results.size() + " entries to " + outPath);
    }

    // ── Helpers ───────────────────────────────────────────────────────────────

    private Set<String> loadSinkNames(String path) {
        Set<String> names = new HashSet<>();
        try {
            String json = new String(Files.readAllBytes(Paths.get(path)));
            // Extract all "name": "..." and "aliases": [...] values
            // Simple regex — no external JSON lib needed
            Pattern nameP  = Pattern.compile("\"name\"\\s*:\\s*\"([^\"]+)\"");
            Pattern aliasP = Pattern.compile("\"([^\"]+)\"");
            Matcher m = nameP.matcher(json);
            while (m.find()) names.add(m.group(1));
            // Also grab aliases arrays
            Pattern aliasBlock = Pattern.compile("\"aliases\"\\s*:\\s*\\[([^\\]]*)]");
            Matcher mb = aliasBlock.matcher(json);
            while (mb.find()) {
                Matcher ma = aliasP.matcher(mb.group(1));
                while (ma.find()) names.add(ma.group(1));
            }
        } catch (Exception e) {
            println("[wghidra-pre] WARNING: could not load sinks.json: " + e.getMessage());
            names.addAll(SINK_ARG_INDEX.keySet());
        }
        return names;
    }

    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 pickHighestPrioritySink(Set<String> sinksHit) {
        // Prefer critical > high sinks by checking known ordering
        String[] preferred = {"gets","strcpy","strcat","sprintf","memcpy",
                              "memmove","snprintf","printf","system","execve",
                              "malloc","calloc"};
        for (String p : preferred) {
            if (sinksHit.contains(p)) return p;
        }
        return sinksHit.iterator().next();
    }

    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) {
        // Find sinkName(...) and extract argument at argIndex
        Pattern p = Pattern.compile("\\b" + Pattern.quote(sinkName) + "\\s*\\(");
        Matcher m = p.matcher(cCode);
        if (!m.find()) return "";

        int start = m.end();
        // Walk the argument list respecting nested parentheses
        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) : "";
    }

    /** Returns [taint, sources_string] */
    private String[] classify(String argExpr, Set<String> callees) {
        String lower = argExpr.toLowerCase();

        // argv / argc pattern in expression
        if (lower.contains("argv") || (lower.contains("param_1") && lower.contains("argc")))
            return new String[]{"USER_TAINTED", "argv"};

        // User-input source in callee set
        Set<String> taintHits = new HashSet<>(callees);
        taintHits.retainAll(USER_INPUT_SOURCES);
        if (!taintHits.isEmpty())
            return new String[]{"USER_TAINTED", join(taintHits)};

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

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

        // sizeof in expression
        if (lower.contains("sizeof"))
            return new String[]{"BOUNDED", "sizeof"};

        // Pointer / array subtraction, with optional element-size multiplication:
        //   (local_70 - uVar10) * 8   or   end_ptr - start_ptr
        // Key signal: " - " between local/param/uvar variables
        boolean hasSubtraction = argExpr.contains(" - ");
        boolean hasLocalVar    = lower.contains("local_") || lower.contains("uvar")
                              || lower.contains("param")  || lower.contains("ptrvar");
        if (hasSubtraction && hasLocalVar)
            return new String[]{"BOUNDED", "pointer_subtraction"};

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

        // Bounded source in callee set (and no taint sources)
        Set<String> boundedHits = new HashSet<>(callees);
        boundedHits.retainAll(BOUNDED_SOURCES);
        if (!boundedHits.isEmpty())
            return new String[]{"BOUNDED", join(boundedHits)};

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

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

    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();
    }

    private String makeEntry(String name, String sink, int argIndex,
                             String taint, String argExpr, String sources) {
        return "{"
            + "\"name\":" + jsonStr(name) + ","
            + "\"sink\":" + jsonStr(sink) + ","
            + "\"arg_index\":" + argIndex + ","
            + "\"taint\":" + jsonStr(taint) + ","
            + "\"arg_expr\":" + jsonStr(argExpr) + ","
            + "\"taint_sources\":[" + jsonStr(sources) + "]"
            + "}";
    }

    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("  \"functions\": {\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());
    }
}