#!/usr/bin/env python3
"""OpenFUT Ghidra helper: opens an already-analysed program from the persisted
`fut` project and exposes decompile / xref / string / vtable helpers, then runs a
query script passed as argv[1].

Run with the restored toolchain:

    GHIDRA_INSTALL_DIR=/home/alex/ghidra/ghidra_11.1.2_PUBLIC \
      /home/alex/re-venv/bin/python tools/re/ghidra_env.py <query.py>

Target program defaults to CardsDLL (the FUT UI, where the kit-selector filter
lives). Override for powdll (the EASFC/POW layer):

    GHIDRA_PROG=powdll.dll ... ghidra_env.py <query.py>
"""
import os, sys

os.environ.setdefault("GHIDRA_INSTALL_DIR", "/home/alex/ghidra/ghidra_11.1.2_PUBLIC")
# Ghidra 11.1.2 does not bundle the in-tree PyGhidra module that the pip
# `pyghidra` 2.x/3.x require, so use the standalone `pyhidra` package (same API).
try:
    import pyhidra as _pg
except ImportError:
    import pyghidra as _pg
_pg.start(verbose=False)

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

PROJ_DIR = os.environ.get("GHIDRA_PROJ_DIR", "/home/alex/ghidra_projects")
PROJ = os.environ.get("GHIDRA_PROJ", "fut")
PROG = os.environ.get("GHIDRA_PROG", "cardsdll.dll")

# Open the ALREADY-ANALYSED program straight from the persisted project.
# pyhidra.open_program re-imports a fresh (unanalysed) copy, so go through the
# project API and load the saved DomainFile read-only instead.
from ghidra.base.project import GhidraProject               # noqa: E402
_project = GhidraProject.openProject(PROJ_DIR, PROJ, True)
prog = _project.openProgram("/", PROG, True)  # (folder, name, readOnly)
flat = None
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 = []
    for r in refs.getReferencesTo(addr(a)):
        fr = r.getFromAddress()
        f = fm.getFunctionContaining(fr)
        out.append((int(fr.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):
    buf = _JBYTE(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."""
    if isinstance(pattern, str):
        pattern = pattern.encode()
    hits = []
    for b in mem.getBlocks():
        if b.getName() not in blocks:
            continue
        start = b.getStart()
        size = int(b.getSize())
        data = read_bytes(int(start.getOffset()), size)
        i = data.find(pattern)
        while i != -1:
            hits.append(int(start.getOffset()) + i)
            i = data.find(pattern, i + 1)
    return hits


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


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:
        with open(sys.argv[1]) as fh:
            code = fh.read()
        exec(compile(code, sys.argv[1], "exec"), globals())
    os._exit(0)