#!/usr/bin/env python3
"""Dump disassembly + decompilation for functions of an imported program.

Plain PyGhidra API (venv with pyghidra). Usage:
  GHIDRA_INSTALL_DIR=... venv/bin/python tools/ghidra_scripts/dump_decomp.py \
      <elf> <project_dir> <project_name> [entry_only|all|0xADDR ...]
"""
import sys

import pyghidra

pyghidra.start()

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

elf, proj_dir, proj_name = sys.argv[1], sys.argv[2], sys.argv[3]
mode = sys.argv[4] if len(sys.argv) > 4 else 'entry_only'

with pyghidra.program_context(pyghidra.open_project(
        __import__('os').path.abspath(proj_dir), proj_name),
        '/' + __import__('os').path.basename(elf)) as prog:
    fm = prog.getFunctionManager()

    def addr(a):
        return prog.getAddressFactory().getAddress(hex(a) if isinstance(a, int) else a)

    dec = DecompInterface()
    dec.openProgram(prog)

    all_funcs = list(fm.getFunctions(True))
    by_ep = {int('%x' % f.getEntryPoint().getOffset(), 16): f
             for f in all_funcs}
    targets = []
    if mode == 'entry_only':
        targets = all_funcs[:1]
    elif mode == 'all':
        targets = all_funcs
    else:
        for s in sys.argv[4:]:
            targets.append(by_ep.get(int(s, 16)))
        targets = [t for t in targets if t is not None]

    # entry listing: first 40 instructions
    ent = fm.getFunctions(True).next()
    listing = prog.getListing()
    it = listing.getInstructions(ent.getEntryPoint(), True)
    print('=== ENTRY LISTING %s ===' % ent.getEntryPoint())
    for i in range(40):
        x = it.next()
        if x is None:
            break
        print('%s  %s' % (x.getAddress(), x))

    for f in targets:
        if f is None:
            continue
        r = dec.decompileFunction(f, 60, ConsoleTaskMonitor())
        print('\n=== %s %s ===' % (f.getEntryPoint(), f.getName()))
        if r.decompileCompleted():
            print(r.getDecompiledFunction().getC())
        else:
            print('// decompile failed:', r.getErrorMessage())