// Recover MSVC 6 RTTI in sacred.exe: for every class with a TypeDescriptor
// (".?AV<Name>@@"), find its CompleteObjectLocators and vtables, create a
// function for every vtable entry and name default-named ones
// "<Class>::vfNN" (NN = byte offset in the vtable). Labels the vtables
// "<Class>::vftable" (plus "_at_<offset>" for secondary base subobjects).
//
// Ghidra's own RecoverClassesFromRTTIScript finds nothing in this binary.
// Run once, headless, without -readOnly (see tools/README.md).
// @category sacred_re

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import ghidra.app.cmd.disassemble.DisassembleCommand;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.mem.Memory;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.program.model.symbol.Namespace;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.model.symbol.SymbolTable;

public class ApplyRtti extends GhidraScript {

    private Memory memory;

    @Override
    protected void run() throws Exception {
        memory = currentProgram.getMemory();
        MemoryBlock text = memory.getBlock(".text");
        MemoryBlock rdata = memory.getBlock(".rdata");
        SymbolTable symbols = currentProgram.getSymbolTable();

        // 1. TypeDescriptors: the mangled name sits 8 bytes in. Search all blocks.
        Map<Long, String> typeDescriptors = new HashMap<>();
        byte[] prefix = ".?AV".getBytes("ISO-8859-1");
        for (MemoryBlock block : memory.getBlocks()) {
            if (!block.isInitialized()) {
                continue;
            }
            Address cursor = block.getStart();
            while (cursor != null && cursor.compareTo(block.getEnd()) < 0) {
                Address hit = memory.findBytes(cursor, block.getEnd(), prefix, null, true, monitor);
                if (hit == null) {
                    break;
                }
                String mangled = readCString(hit, 512);
                if (mangled.endsWith("@@")) {
                    typeDescriptors.put(hit.getOffset() - 8, mangled.substring(4, mangled.length() - 2));
                }
                cursor = hit.add(1);
            }
        }
        println("type descriptors: " + typeDescriptors.size());

        // 2. CompleteObjectLocators in .rdata: {0, offset, cdOffset, pTD, pCHD}.
        Map<Long, long[]> locators = new HashMap<>(); // col -> {td, subobjectOffset}
        long rdataStart = rdata.getStart().getOffset();
        long rdataEnd = rdata.getEnd().getOffset();
        for (long a = rdataStart; a + 20 <= rdataEnd; a += 4) {
            if (readU32(a) != 0) {
                continue;
            }
            long td = readU32(a + 12);
            if (typeDescriptors.containsKey(td)) {
                locators.put(a, new long[] {td, readU32(a + 4)});
            }
        }
        println("object locators: " + locators.size());

        // 3. vtables: a dword in .rdata pointing at a COL; entries follow it.
        int vtables = 0;
        int created = 0;
        int named = 0;
        for (long a = rdataStart; a + 8 <= rdataEnd; a += 4) {
            long[] locator = locators.get(readU32(a));
            if (locator == null) {
                continue;
            }
            String className = typeDescriptors.get(locator[0]);
            long vtable = a + 4;
            Namespace namespace = symbols.getOrCreateNameSpace(currentProgram.getGlobalNamespace(),
                    className, SourceType.ANALYSIS);
            String suffix = locator[1] == 0 ? "" : String.format("_at_%x", locator[1]);
            symbols.createLabel(toAddr(vtable), "vftable" + suffix, namespace, SourceType.ANALYSIS);
            vtables++;

            List<Long> entries = new ArrayList<>();
            for (long e = vtable; e + 4 <= rdataEnd; e += 4) {
                long target = readU32(e);
                if (!text.contains(toAddr(target))) {
                    break;
                }
                if (e != vtable && locators.containsKey(target)) {
                    break;
                }
                entries.add(target);
            }
            for (int i = 0; i < entries.size(); i++) {
                Address entry = toAddr(entries.get(i));
                Function function = getFunctionAt(entry);
                if (function == null) {
                    new DisassembleCommand(entry, null, true).applyTo(currentProgram, monitor);
                    function = createFunction(entry, null);
                    if (function == null) {
                        continue;
                    }
                    created++;
                }
                if (function.getSymbol().getSource() == SourceType.DEFAULT) {
                    function.setName(String.format("vf%02x%s", i * 4, suffix), SourceType.ANALYSIS);
                    function.setParentNamespace(namespace);
                    named++;
                }
            }
        }
        println("vtables: " + vtables + ", functions created: " + created + ", named: " + named);
    }

    private long readU32(long address) throws Exception {
        return memory.getInt(toAddr(address)) & 0xffffffffL;
    }

    private String readCString(Address address, int max) throws Exception {
        StringBuilder text = new StringBuilder();
        for (int i = 0; i < max; i++) {
            byte b = memory.getByte(address.add(i));
            if (b == 0) {
                break;
            }
            text.append((char) (b & 0xff));
        }
        return text.toString();
    }
}