# Lofi12XT_Map.py — Ghidra-side script (Jython2 + Ghidrathon compatible).
# Run headless as -postScript. Two modes:
#   Mode A (map):  Lofi12XT_Map.py <unpack_dir>
#     Creates one memory block per sect*.bin at its DDR address (from headers.json),
#     marks mmtd entry point, disassembles + makes a function there.
#   Mode B (vtables): Lofi12XT_Map.py vtables <vtables.csv>
#     Labels each vtable run, converts entries to pointers, bookmarks them.
# No f-strings (Jython 2.7 safe). No third-party deps.
import json
import os

from ghidra.program.model.symbol import SourceType
from ghidra.program.model.data import PointerDataType
from java.io import File


def log(msg):
    print("[Lofi12XT] " + msg)


def map_sects(unpack_dir):
    info = json.load(open(os.path.join(unpack_dir, "headers.json")))
    mem = currentProgram.getMemory()
    for s in info["sects"]:
        name = "sect%d" % s["index"]
        addr = toAddr(s["addr_hex"])
        fn = os.path.join(unpack_dir,
                          "sect%d_addr%s.bin" % (s["index"], s["addr_hex"]))
        size = s["len"]
        if mem.getBlock(addr) is not None:
            log(name + " already mapped at " + s["addr_hex"] + ", skip")
            continue
        mem.createInitializedBlock(name, addr, File(fn), size,
                                   "from lofi_unpack", "", False)
        blk = mem.getBlock(addr)
        is_code = (size > 1000000)  # only the big sect is code
        blk.setRead(True)
        blk.setWrite(not is_code)
        blk.setExecute(is_code)
        log("mapped %s %s len=%d exec=%s" % (name, s["addr_hex"], size, is_code))

    entry = toAddr(info["mmtd"]["entry_hex"])
    currentProgram.getSymbolTable().addExternalEntryPoint(entry)
    createLabel(entry, "fw_entry", True, SourceType.IMPORTED)
    disassemble(entry)
    createFunction(entry, "fw_entry")
    log("entry " + info["mmtd"]["entry_hex"] + " marked + function created")


def map_vtables(csv_path):
    st = currentProgram.getSymbolTable()
    bm = currentProgram.getBookmarkManager()
    count = 0
    with open(csv_path) as f:
        header = f.readline()
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split(",", 3)
            load = toAddr(parts[1])
            entries = parts[3].split(";")
            createLabel(load, "vtable_%s" % parts[1], True,
                        SourceType.ANALYSIS)
            for k, e in enumerate(entries):
                ea = load.add(k * 4)
                try:
                    createData(ea, PointerDataType.dataType)
                except Exception:
                    try:
                        createDword(ea)
                    except Exception:
                        pass  # labels/bookmarks below still land
                try:
                    createLabel(toAddr(e), "vfunc_%s_%d" % (parts[1], k),
                                False, SourceType.ANALYSIS)
                except Exception:
                    pass
            bm.setBookmark(load, "Note", "vtable",
                           "vtable %s n=%d" % (parts[1], len(entries)))
            count += 1
    log("labeled %d vtables from %s" % (count, csv_path))


args = getScriptArgs()
if len(args) == 1:
    map_sects(args[0])
elif len(args) == 2 and args[0] == "vtables":
    map_vtables(args[1])
else:
    log("usage: Lofi12XT_Map.py <unpack_dir>  OR  Lofi12XT_Map.py vtables <csv>")