#!/usr/bin/env python3
"""PyGhidra driver: map AIS sections, analyze, decompile checksum targets.

Runs with the project venv (has pyghidra): ./venv/bin/python /tmp/opencode/ghidra-lab/pyais.py
Env: GHIDRA_INSTALL_DIR + JAVA_HOME must be set (see launch command).
Writes /tmp/opencode/ghidra-lab/out/{dis,dec}_<ADDR>.txt
Stdlib + pyghidra + jpype only.
"""
import os
import struct
import sys

UNPACK = "/tmp/ais-unpack"
OUT = "/tmp/opencode/ghidra-lab/out"
SECT_IMAGE = "/tmp/ais-unpack/ais_sect1_addrC7074800.bin"
PROJ_LOC = "/tmp/opencode/ghidra-lab/pyproj"
PROJ_NAME = "LofiPy"
LANG = "C6000:LE:32:default"
ENTRY = "C70A28A0"
TARGETS = ["C70A0A60", "C708E4E4", "C7086400", "C70A28A0", "C707841C"]


def parse_ais(path):
    d = open(path, "rb").read()
    assert d[:4] == b"TIPA", "bad AIS magic"
    segs, i, n = [], 4, len(d)
    entry = None
    while i + 12 <= n:
        op = d[i:i + 4]
        if op == bytes([0x01, 0x59, 0x53, 0x58]):
            addr, sz = struct.unpack("<II", d[i + 4:i + 12])
            segs.append((addr, d[i + 12:i + 12 + sz]))
            i += 12 + sz
        elif op == bytes([0x06, 0x59, 0x53, 0x58]):
            entry = struct.unpack("<I", d[i + 4:i + 8])[0]
            break
        else:
            i += 4
    return segs, entry


def main():
    os.makedirs(OUT, exist_ok=True)
    import jpype
    import pyghidra

    segs, entry = parse_ais("/var/home/dejvino/Downloads/Lofi-12XT/ic300.bin")
    print("segs=%d entry=%08X" % (len(segs), entry), flush=True)

    pyghidra.start()
    with pyghidra.open_program(
        SECT_IMAGE,
        project_location=PROJ_LOC,
        project_name=PROJ_NAME,
        analyze=False,
        language=LANG,
    ) as flat:
        program = flat.getCurrentProgram()
        print("program=" + program.getName()
              + " lang=" + program.getLanguageID().toString(), flush=True)
        JByte = jpype.JArray(jpype.JByte)
        try:
            from java.io import ByteArrayInputStream
        except ImportError:
            import jpype.imports  # noqa
            from java.io import ByteArrayInputStream
        from ghidra.util.task import TaskMonitor
        from ghidra.program.model.symbol import SourceType

        # 1. map blocks at DDR addresses
        with pyghidra.transaction(program, "map AIS"):
            mem = program.getMemory()
            for idx, (addr_int, blob) in enumerate(segs):
                addr = flat.toAddr("0x%08X" % addr_int)
                if mem.getBlock(addr) is not None:
                    print("ais%d already mapped" % idx, flush=True)
                    continue
                stream = ByteArrayInputStream(JByte(bytes(blob)))
                blk = mem.createInitializedBlock(
                    "ais%d" % idx, addr, stream, len(blob),
                    TaskMonitor.DUMMY, False)
                blk.setRead(True)
                blk.setWrite(False)
                blk.setExecute(True)
                print("mapped ais%d %08X len=%d"
                      % (idx, addr_int, len(blob)), flush=True)
        # 2. entry + analyze
        with pyghidra.transaction(program, "entry"):
            eaddr = flat.toAddr("0x" + ENTRY)
            try:
                program.getSymbolTable().addExternalEntryPoint(eaddr)
            except Exception as e:
                print("entry point: " + str(e), flush=True)
            flat.disassemble(eaddr)
            if flat.getFunctionAt(eaddr) is None:
                flat.createFunction(eaddr, "ais_entry")
        print("analyzing...", flush=True)
        flat.analyzeAll(program)
        print("analysis done", flush=True)
        # 3. decompile + disassembly dump per target
        from ghidra.app.decompiler.flatapi import FlatDecompilerAPI

        dapi = FlatDecompilerAPI(flat)
        try:
            for t in TARGETS:
                addr = flat.toAddr("0x" + t)
                try:
                    fn = flat.getFunctionAt(addr)
                    if fn is None:
                        flat.disassemble(addr)
                        try:
                            fn = flat.createFunction(addr, "sub_" + t)
                        except Exception as e:
                            print(t + " createFunction: " + str(e),
                                  flush=True)
                    # disassembly window
                    try:
                        start = flat.toAddr("0x%08X"
                                            % (int(t, 16) - 64))
                        it = program.getListing().getInstructions(start,
                                                                  True)
                        lines, n = [], 0
                        while it.hasNext() and n < 150:
                            ins = it.next()
                            lines.append("%08X %s %s" % (
                                ins.getAddress().getOffset(),
                                ins.getMnemonicString(), ins.toString()))
                            n += 1
                        open(os.path.join(OUT, "dis_" + t + ".txt"),
                             "w").write("\n".join(lines) + "\n")
                    except Exception as e:
                        print(t + " disasm: " + str(e), flush=True)
                    # decompile
                    if fn is not None:
                        try:
                            c = dapi.decompile(fn)
                            if not c:
                                c = "DECOMPILE_NULL"
                            open(os.path.join(OUT, "dec_" + t + ".txt"),
                                 "w").write(c if c else "DECOMPILE_NULL")
                            print(t + " decompiled %d chars"
                                  % (len(c) if c else 0), flush=True)
                        except Exception as e:
                            print(t + " decompile: " + str(e), flush=True)
                            open(os.path.join(OUT, "dec_" + t + ".txt"),
                                 "w").write("DECOMPILE_ERROR: " + str(e))
                except Exception as e:
                    print(t + " FAILED: " + str(e), flush=True)
        finally:
            dapi.dispose()
    print("ALL_DONE -> " + OUT, flush=True)


if __name__ == "__main__":
    sys.exit(main())