from idautils import *
from idaapi import *
from idc import *
from idc_bc695 import *
import ida_frame
import ida_funcs
import ida_typeinf
import sys
import idaapi
import ida_nalt
import idautils
import json

filename = GetInputFile()
decom_buf = list()

idaapi.auto_wait()

#idaapi.load_plugin('hexrays')
#idaapi.load_plugin('hexx64')
idaapi.load_plugin('hexarm')
if not idaapi.init_hexrays_plugin():
    print('Unable to load Hex-rays')

kb = {}

#callees = dict()
for function_ea in idautils.Functions():
    f_name = idc.get_func_name(function_ea)

    func_type = idc.get_type(function_ea)

    tif = ida_typeinf.tinfo_t()
    ida_nalt.get_tinfo(tif, function_ea)
    ftd = ida_typeinf.func_type_data_t()
    tif.get_func_details(ftd)

    kb[function_ea] = {}
    kb[function_ea]['name'] = f_name
    kb[function_ea]['type'] = ftd.rettype.dstr()
    kb[function_ea]['caller'] = {}
    for ref_ea in idautils.CodeRefsTo(function_ea, 0):
        caller_name = idc.get_func_name(ref_ea)
        kb[function_ea]['caller'][ref_ea] = caller_name

for function_ea in idautils.Functions():
    orig_f_name = idc.get_func_name(function_ea)
    f = get_func(function_ea)
    try:
        cfunc = decompile(f)
        vuu = open_pseudocode(function_ea, 0)
        if cfunc is None:
            continue

        kb[function_ea]['var'] = {}
        for v in cfunc.arguments:
            loc = v.location

            reg = print_vdloc(loc, 1)
            kb[function_ea]['var'][reg] = [v.name, str(v.type())]

        for v in cfunc.lvars:
            if v in cfunc.arguments:
                continue
            loc = v.location
            width = v.width
            if 'r' in print_vdloc(loc, 1):
                continue
            loc = int(print_vdloc(loc, 1).replace('^',''), 16)
            for i in range(width):
                kb[function_ea]['var'][loc+i] = [v.name, str(v.type())]

    except Exception as e:
        continue

with open(filename + '_kb.json', 'w') as outfile:
    json.dump(kb, outfile, indent=2)

print ('Done')
ida_pro.qexit(0)
