#!/usr/bin/env python3
"""OpenFUT Ghidra helper: opens the analysed cardsdll.dll program once and exposes
decompile / xref / vtable helpers, so each RE question is a small python file
instead of a JVM restart + OSGi compile.

Usage:  ghidra_env.py <query.py>     -- runs <query.py> with the helpers in scope
        (see tools/ghidra_queries/ for worked examples)

WHY PYGHIDRA: this box's Ghidra 12.1.2 cannot compile .java scripts at all --
analyzeHeadless -postScript Foo.java dies with "Failed to get OSGi bundle
containing script" for EVERY script, including ones that ran before (it is the
in-process OSGi/javac path that is broken, not the scripts). PyGhidra bypasses it.
Setup once:
    python3 -m venv gvenv
    gvenv/bin/pip install --no-index \
        --find-links /opt/ghidra/Ghidra/Features/PyGhidra/pypkg/dist pyghidra
    gvenv/bin/python ghidra_env.py <query.py>

Two traps this file already works around:
  * open_program(..., nested_project_location=False) -- otherwise pyghidra creates
    a NEW empty project at <loc>/<name>/ and re-imports (losing the analysis).
  * os._exit(0) at the end -- JVM teardown under jpype deadlocks forever.
  * read_bytes() uses a Java byte[]; passing a Python bytearray to Memory.getBytes
    silently reads NOTHING and every scan comes back with 0 hits.
"""
import os, sys

os.environ.setdefault("GHIDRA_INSTALL_DIR", "/opt/ghidra")
import pyghidra

pyghidra.start(verbose=False)

from ghidra.app.decompiler import DecompInterface           # noqa: E402
from ghidra.util.task import ConsoleTaskMonitor             # noqa: E402

DLL = "/tmp/fut/cardsdll.dll"
PROJ_DIR, PROJ = "/tmp/ghidra_fut", "cardsdll"

# nested_project_location=False -> use /tmp/ghidra_fut/cardsdll.gpr itself (the
# already-analysed project) instead of creating /tmp/ghidra_fut/cardsdll/.
_ctx = pyghidra.open_program(DLL, project_location=PROJ_DIR, project_name=PROJ,
                             analyze=False, program_name="cardsdll.dll",
                             nested_project_location=False)
flat = _ctx.__enter__()
prog = flat.getCurrentProgram()
mon = ConsoleTaskMonitor()
fm = prog.getFunctionManager()
listing = prog.getListing()
mem = prog.getMemory()
refs = prog.getReferenceManager()

_dec = DecompInterface()
_dec.openProgram(prog)


def addr(a):
    return prog.getAddressFactory().getDefaultAddressSpace().getAddress(int(a))


def func(a):
    return fm.getFunctionContaining(addr(a)) if not hasattr(a, "getEntryPoint") else a


def dec(a, timeout=180):
    """Decompiled C for the function containing address a."""
    f = func(a)
    if f is None:
        return "// no function at %#x" % int(a)
    r = _dec.decompileFunction(f, timeout, mon)
    if r is None or not r.decompileCompleted():
        return "// decompile failed for %s" % f.getName()
    return str(r.getDecompiledFunction().getC())


def xrefs_to(a):
    """[(from_addr, reftype, containing_function_name, entry)] for refs to a."""
    out = []
    it = refs.getReferencesTo(addr(a))
    while it.hasNext():
        r = it.next()
        f = fm.getFunctionContaining(r.getFromAddress())
        out.append((int(r.getFromAddress().getOffset()), str(r.getReferenceType()),
                    f.getName() if f else "?",
                    int(f.getEntryPoint().getOffset()) if f else 0))
    return out


def qword(a):
    return mem.getLong(addr(a)) & 0xFFFFFFFFFFFFFFFF


def dword(a):
    return mem.getInt(addr(a)) & 0xFFFFFFFF


import jpype                                                    # noqa: E402
_JBYTE = jpype.JArray(jpype.JByte)


def read_bytes(a, n):
    """Bulk read n bytes at a. MUST use a Java byte[] -- passing a Python
    bytearray to Memory.getBytes silently reads nothing (this bug quietly
    zeroed several earlier scans)."""
    buf = _JBYTE(int(n))
    got = mem.getBytes(addr(a), buf)
    return bytes((int(x) & 0xFF) for x in buf[:got])


def find_all(pattern, blocks=(".text", ".rdata", ".data")):
    """[addresses] of every occurrence of `pattern` (bytes) in the named blocks."""
    hits = []
    for b in mem.getBlocks():
        if b.getName() not in blocks or not b.isInitialized():
            continue
        s = int(b.getStart().getOffset())
        size = int(b.getEnd().getOffset()) - s + 1
        off = 0
        chunk = 1 << 20
        while off < size:
            ln = min(chunk, size - off)
            try:
                data = read_bytes(s + off, ln)
            except Exception:
                off += ln
                continue
            i = data.find(pattern)
            while i != -1:
                hits.append(s + off + i)
                i = data.find(pattern, i + 1)
            off += ln - (len(pattern) - 1) if ln == chunk else ln
    return hits


def rd_str(a, maxlen=200):
    b = bytearray()
    p = int(a)
    for _ in range(maxlen):
        c = mem.getByte(addr(p)) & 0xFF
        if c == 0:
            break
        b.append(c)
        p += 1
    return b.decode("utf-8", "replace")


def vtable(a, n=64):
    """[(slot_offset, target_addr, function_name)] reading n qwords at a."""
    out = []
    for i in range(n):
        try:
            t = qword(int(a) + i * 8)
        except Exception:
            break
        f = fm.getFunctionAt(addr(t)) if 0x180000000 <= t < 0x181000000 else None
        out.append((i * 8, t, f.getName() if f else ""))
    return out


def fname(a):
    f = func(a)
    return f.getName() if f else "?"


def callees(a):
    f = func(a)
    return sorted({(int(c.getEntryPoint().getOffset()), c.getName())
                   for c in f.getCalledFunctions(mon)}) if f else []


def callers(a):
    f = func(a)
    return sorted({(int(c.getEntryPoint().getOffset()), c.getName())
                   for c in f.getCallingFunctions(mon)}) if f else []


if __name__ == "__main__":
    if len(sys.argv) > 1:
        g = dict(globals())
        g["__name__"] = "__main__"
        exec(open(sys.argv[1]).read(), g)
    else:
        print("loaded:", prog.getName(), fm.getFunctionCount(), "functions")
    sys.stdout.flush()
    # JVM teardown deadlocks under jpype here -- skip it, all output is flushed.
    os._exit(0)