import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.*;
import ghidra.program.model.mem.*;
import ghidra.program.model.listing.*;
import ghidra.program.model.symbol.*;
import ghidra.program.model.scalar.*;
import ghidra.app.decompiler.*;
import java.util.*;
import java.util.regex.*;
import java.io.*;

// Pass 4 (read-only): network message registry table, ProcessTurn callee strings, +0x170 callback writers, misc decompiles.
public class SpineRecon4 extends GhidraScript {
    DecompInterface decomp; Memory mem; PrintWriter out; ReferenceManager rm; SymbolTable st;
    Set<Function> toDump = new LinkedHashSet<Function>();
    static final long[] DECOMP = {0x00891340L,0x007598e0L,0x00898800L,0x0089a640L,0x00899210L,0x0090c700L,0x008e5ac0L,0x00783be0L,0x00783980L,0x007856f0L,0x007cd2a0L,0x007d4400L,0x00752500L,0x0078a7c0L,0x007ad100L,0x007d7f70L,0x00814ea0L,0x0086b300L,0x007adc80L,0x007a0e20L,0x007999a0L,0x007ae480L,0x007af0b0L,0x008986a0L,0x0089cc70L,0x00898690L,0x00723df0L};
    static final long[] CALLEE_STRS = {0x007dc6c0L,0x007598e0L,0x00891340L,0x007da9a0L,0x007d4400L,0x007cda40L,0x007842b0L,0x00741cd0L,0x00732ab0L};

    String cstr(Address a, int max) {
        try { byte[] b = new byte[max]; int got = mem.getBytes(a, b); int i = 0;
            while (i < got && b[i] != 0 && (b[i]&0xff) >= 0x20 && (b[i]&0xff) < 0x7f) i++;
            if (i >= 1 && i < got && b[i] == 0) return new String(b, 0, i, "ISO-8859-1"); } catch (Exception e) {}
        return null;
    }
    String decompRes(Function f) {
        try { DecompileResults res = decomp.decompileFunction(f, 240, monitor);
            if (res == null || !res.decompileCompleted()) return "[decompile failed]";
            String c = res.getDecompiledFunction().getC();
            Matcher m = Pattern.compile("&?(DAT|PTR_s_|s_[A-Za-z0-9_]*|PTR_DAT|u_[A-Za-z0-9_]*)_([0-9a-f]{8})").matcher(c);
            StringBuffer sb = new StringBuffer();
            while (m.find()) { Address a = toAddr(Long.parseLong(m.group(2), 16)); String s = cstr(a, 100); String rep = m.group(0);
                if (s != null && s.length() <= 96) rep = "\"" + s + "\""; m.appendReplacement(sb, Matcher.quoteReplacement(rep)); }
            m.appendTail(sb); return sb.toString();
        } catch (Exception e) { return "[exception " + e.getMessage() + "]"; }
    }
    void dumpFunc(Function f) throws IOException {
        String n = String.format("%08x", f.getEntryPoint().getOffset());
        File fl = new File("/tmp/spine/" + n + ".c"); if (fl.exists()) return;
        PrintWriter w = new PrintWriter(new FileWriter(fl));
        w.println("// " + f.getName() + " @ " + f.getEntryPoint() + " size=" + f.getBody().getNumAddresses());
        StringBuilder cs = new StringBuilder(); for (Function c : f.getCallingFunctions(monitor)) cs.append(c.getName() + "@" + c.getEntryPoint() + " ");
        w.println("// CALLERS: " + cs);
        StringBuilder ce = new StringBuilder(); for (Function c : f.getCalledFunctions(monitor)) ce.append(c.getName() + " ");
        w.println("// CALLEES: " + ce);
        w.println(decompRes(f)); w.close();
    }
    String strsOf(Function f) {
        TreeMap<Address,String> refs = new TreeMap<Address,String>();
        AddressIterator ai = rm.getReferenceSourceIterator(f.getBody(), true);
        while (ai.hasNext()) { Address from = ai.next(); for (Reference r : rm.getReferencesFrom(from)) { Address to = r.getToAddress(); if (!mem.contains(to)) continue;
            String sv = cstr(to, 90); if (sv != null && sv.length() >= 4) refs.put(from, sv);
            Symbol s = st.getPrimarySymbol(to); if (s != null && s.getName(true).endsWith("::vftable") && !s.getName(true).contains("Helper")) refs.put(from, "<" + s.getName(true) + ">"); } }
        StringBuilder sb = new StringBuilder(); for (String s : refs.values()) sb.append(s.startsWith("<") ? s + " " : "\"" + s + "\" ");
        return sb.toString();
    }

    @Override
    public void run() throws Exception {
        mem = currentProgram.getMemory(); rm = currentProgram.getReferenceManager(); st = currentProgram.getSymbolTable();
        decomp = new DecompInterface(); decomp.openProgram(currentProgram);
        new File("/tmp/spine").mkdirs();
        out = new PrintWriter(new FileWriter("/tmp/spine/recon4.txt"));

        out.println("==== 1. NETWORK MESSAGE REGISTRY (static initializers calling FUN_008d2290) ====");
        // scan .text for  MOV ECX,imm32 ; CALL 0x008d2290  and read the two preceding pushes
        Address lo = toAddr(0x009b0000L), hi = toAddr(0x009c8000L);
        Address a = lo; int found = 0;
        byte[] buf = new byte[(int)(hi.subtract(lo))]; mem.getBytes(lo, buf);
        TreeMap<Integer,String> byId = new TreeMap<Integer,String>();
        for (int i = 0; i + 10 < buf.length; i++) {
            if ((buf[i]&0xff) == 0xB9 && (buf[i+5]&0xff) == 0xE8) {
                long rel = (buf[i+6]&0xffL) | ((buf[i+7]&0xffL)<<8) | ((buf[i+8]&0xffL)<<16) | ((buf[i+9]&0xffL)<<24);
                long target = lo.getOffset() + i + 10 + (int)rel;
                if (target != 0x008d2290L) continue;
                long entry = (buf[i+1]&0xffL) | ((buf[i+2]&0xffL)<<8) | ((buf[i+3]&0xffL)<<16) | ((buf[i+4]&0xffL)<<24);
                // walk back: PUSH name (68 imm32), PUSH id (6a imm8 | 68 imm32), PUSH factory (68 imm32)
                int p = i - 5; long name = 0, id = -1, factory = 0;
                if ((buf[p]&0xff) == 0x68) { name = (buf[p+1]&0xffL) | ((buf[p+2]&0xffL)<<8) | ((buf[p+3]&0xffL)<<16) | ((buf[p+4]&0xffL)<<24); }
                else continue;
                int q;
                if ((buf[p-2]&0xff) == 0x6a) { id = buf[p-1]&0xff; q = p - 2; }
                else if ((buf[p-5]&0xff) == 0x68) { id = (buf[p-4]&0xffL) | ((buf[p-3]&0xffL)<<8) | ((buf[p-2]&0xffL)<<16) | ((buf[p-1]&0xffL)<<24); q = p - 5; }
                else continue;
                if ((buf[q-5]&0xff) == 0x68) factory = (buf[q-4]&0xffL) | ((buf[q-3]&0xffL)<<8) | ((buf[q-2]&0xffL)<<16) | ((buf[q-1]&0xffL)<<24);
                String nm = cstr(toAddr(name), 64);
                Function ff = getFunctionAt(toAddr(factory));
                String line = String.format("id=0x%02x  entry=%08x  name=%-32s factory=%08x %s  init@%08x", id, entry, nm, factory, ff == null ? "(nofunc)" : ff.getName() + "(sz " + ff.getBody().getNumAddresses() + ")", lo.getOffset() + i);
                byId.put((int)id, line); found++;
            }
        }
        for (String l : byId.values()) out.println(l);
        out.println("total " + found);

        out.println("\n==== 2. WRITERS of StrategyServer+0x170 (callback) : scan for 'mov [reg+0x170], ' in functions referencing StrategyServer ====");
        // brute: instructions with scalar 0x170 operand as memory displacement, being a store
        for (long fa : new long[]{0x007d78d0L, 0x00888e80L, 0x007d20d0L}) {
            Function f = getFunctionAt(toAddr(fa)); if (f == null) continue;
            InstructionIterator ii = currentProgram.getListing().getInstructions(f.getBody(), true);
            while (ii.hasNext()) { Instruction ins = ii.next(); String s = ins.toString(); if (s.contains("0x170]")) out.println("   " + f.getName() + " " + ins.getAddress() + "  " + s); }
        }
        // also global search on refs to all functions passed as pointer near 0x170... skip. Check refs to the candidate callback functions: StrategyClient vft[5] FUN_00774fc0 and others in 0x0077xxxx
        for (long fa : new long[]{0x00774fc0L, 0x00774f30L, 0x00775100L, 0x00779a70L}) {
            Function f = getFunctionAt(toAddr(fa)); if (f == null) continue; StringBuilder sb = new StringBuilder();
            ReferenceIterator ri = rm.getReferencesTo(f.getEntryPoint()); while (ri.hasNext()) { Reference r = ri.next(); Function c = getFunctionContaining(r.getFromAddress()); sb.append(r.getFromAddress() + "(" + r.getReferenceType() + (c == null ? "" : " in " + c.getName()) + ") "); }
            out.println("   refs to " + f.getName() + ": " + sb);
        }

        out.println("\n==== 3. CALLEE STRINGS ====");
        for (long fa : CALLEE_STRS) {
            Function f = getFunctionAt(toAddr(fa)); if (f == null) continue; out.println("DRIVER " + f.getName() + "@" + f.getEntryPoint() + " sz=" + f.getBody().getNumAddresses() + "  strs: " + strsOf(f));
            for (Function c : f.getCalledFunctions(monitor)) { if (c.getBody().getNumAddresses() < 40) continue; String s = strsOf(c); StringBuilder cs = new StringBuilder(); int n = 0; for (Function cc : c.getCallingFunctions(monitor)) if (n++ < 4) cs.append(cc.getName() + " ");
                out.println("   -> " + c.getName() + "@" + c.getEntryPoint() + " sz=" + c.getBody().getNumAddresses() + " callers(" + n + "):[" + cs.toString().trim() + "]  strs: " + s); }
        }
        for (long d : DECOMP) { Function f = getFunctionAt(toAddr(d)); if (f != null) toDump.add(f); }
        out.println("\n==== 4. DUMPED ====");
        for (Function f : toDump) { out.println(f.getName() + " @ " + f.getEntryPoint() + " sz=" + f.getBody().getNumAddresses()); dumpFunc(f); }
        out.close(); decomp.dispose(); println("done");
    }
}