diff --git a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py index edb2b8b..9aea15d 100644 --- a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py +++ b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py @@ -37,7 +37,6 @@ import idaapi import idc import json -import ntpath import os import time @@ -205,7 +204,7 @@ def get_basic_blocks(fva): def get_bb_disasm(bb, md, prefix): """Return the (nomalized) disassembly for a BasicBlock.""" - b64_bytes = base64.b64encode(idc.get_bytes(bb.va, bb.size)) + b64_bytes = base64.b64encode(idc.get_bytes(bb.va, bb.size)).decode() bb_heads, bb_mnems, bb_disasm, bb_norm = \ capstone_disassembly(md, bb.va, bb.size, prefix) return b64_bytes, bb_heads, bb_mnems, bb_disasm, bb_norm @@ -222,7 +221,7 @@ def run_acfg_disasm(idb_path, fva_list, output_dir): output_dict = dict() output_dict[idb_path] = dict() - procname = idaapi.get_inf_structure().procName.lower() + procname = idaapi.get_inf_structure().procname.lower() bitness = get_bitness() output_dict[idb_path]['arch'] = convert_procname_to_str(procname, bitness) md, prefix = initialize_capstone(procname, bitness) @@ -272,7 +271,7 @@ def run_acfg_disasm(idb_path, fva_list, output_dir): print("[!] Exception: skipping function fva: %d" % fva) print(e) - out_name = ntpath.basename(idb_path.replace(".i64", "_acfg_disasm.json")) + out_name = os.path.basename(idb_path.replace(".i64", "_acfg_disasm.json")) with open(os.path.join(output_dir, out_name), "w") as f_out: json.dump(output_dict, f_out) @@ -280,12 +279,12 @@ def run_acfg_disasm(idb_path, fva_list, output_dir): if __name__ == '__main__': if not idaapi.get_plugin_options("acfg_disasm"): print("[!] -Oacfg_disasm option is missing") - idc.Exit(1) + idaapi.qexit(1) plugin_options = idaapi.get_plugin_options("acfg_disasm").split(":") if len(plugin_options) != 3: print("[!] -Oacfg_disasm:INPUT_JSON:IDB_PATH:OUTPUT_DIR is required") - idc.Exit(1) + idaapi.qexit(1) input_json = plugin_options[0] idb_path = plugin_options[1] @@ -296,10 +295,10 @@ def run_acfg_disasm(idb_path, fva_list, output_dir): if idb_path not in selected_functions: print("[!] Error! IDB path (%s) not in %s" % (idb_path, input_json)) - idc.Exit(1) + idaapi.qexit(1) fva_list = selected_functions[idb_path] print("[D] Found %d addresses" % len(fva_list)) run_acfg_disasm(idb_path, fva_list, output_dir) - idc.Exit(0) + idaapi.qexit(0) diff --git a/IDA_scripts/IDA_acfg_disasm/test_acfg_disasm.py b/IDA_scripts/IDA_acfg_disasm/test_acfg_disasm.py index 192ec24..5407b0c 100644 --- a/IDA_scripts/IDA_acfg_disasm/test_acfg_disasm.py +++ b/IDA_scripts/IDA_acfg_disasm/test_acfg_disasm.py @@ -41,13 +41,20 @@ class TestIDAAcfgDisasm(unittest.TestCase): + maxDiff = None + def remove_elaspesed_time(self, jd): """Elapsed time will be different in any run.""" for idb, addr in jd.items(): for va in addr: - if isinstance(jd[idb][va], dict) \ - and 'elapsed_time' in jd[idb][va]: + if not isinstance(jd[idb][va], dict): + continue + if 'elapsed_time' in jd[idb][va]: jd[idb][va]['elapsed_time'] = -1 + if 'nodes' in jd[idb][va]: + jd[idb][va]['nodes'] = sorted(jd[idb][va]['nodes']) + if 'edges' in jd[idb][va]: + jd[idb][va]['edges'] = sorted(jd[idb][va]['edges']) return jd def test_acfg_disasm(self): diff --git a/IDA_scripts/IDA_acfg_disasm/test_large_acfg_disasm.py b/IDA_scripts/IDA_acfg_disasm/test_large_acfg_disasm.py index b9ce15b..942ff93 100644 --- a/IDA_scripts/IDA_acfg_disasm/test_large_acfg_disasm.py +++ b/IDA_scripts/IDA_acfg_disasm/test_large_acfg_disasm.py @@ -41,13 +41,20 @@ class TestIDAAcfgDisasm(unittest.TestCase): + maxDiff = None + def remove_elaspesed_time(self, jd): """Elapsed time will be different in any run.""" for idb, addr in jd.items(): for va in addr: - if isinstance(jd[idb][va], dict) \ - and 'elapsed_time' in jd[idb][va]: + if not isinstance(jd[idb][va], dict): + continue + if 'elapsed_time' in jd[idb][va]: jd[idb][va]['elapsed_time'] = -1 + if 'nodes' in jd[idb][va]: + jd[idb][va]['nodes'] = sorted(jd[idb][va]['nodes']) + if 'edges' in jd[idb][va]: + jd[idb][va]['edges'] = sorted(jd[idb][va]['edges']) return jd def test_acfg_disasm_large(self): diff --git a/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py b/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py index d2f5c4b..ed20b1e 100644 --- a/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py +++ b/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py @@ -36,7 +36,6 @@ import idautils import idc import json -import ntpath import os import time @@ -224,7 +223,7 @@ def run_acfg_features(idb_path, fva_list, output_dir): output_dict = dict() output_dict[idb_path] = dict() - procname = idaapi.get_inf_structure().procName.lower() + procname = idaapi.get_inf_structure().procname.lower() bitness = get_bitness() md, arch = initialize_capstone(procname, bitness) @@ -265,7 +264,7 @@ def run_acfg_features(idb_path, fva_list, output_dir): print("[!] Exception: skipping function fva: %d" % fva) print(e) - out_name = ntpath.basename(idb_path.replace(".i64", "_acfg_features.json")) + out_name = os.path.basename(idb_path.replace(".i64", "_acfg_features.json")) with open(os.path.join(output_dir, out_name), "w") as f_out: json.dump(output_dict, f_out) @@ -273,12 +272,12 @@ def run_acfg_features(idb_path, fva_list, output_dir): if __name__ == '__main__': if not idaapi.get_plugin_options("acfg_features"): print("[!] -Oacfg_features option is missing") - idc.Exit(1) + idaapi.qexit(1) plugin_options = idaapi.get_plugin_options("acfg_features").split(":") if len(plugin_options) != 3: print("[!] -Oacfg_features:INPUT_JSON:IDB_PATH:OUTPUT_DIR is required") - idc.Exit(1) + idaapi.qexit(1) input_json = plugin_options[0] idb_path = plugin_options[1] @@ -289,10 +288,10 @@ def run_acfg_features(idb_path, fva_list, output_dir): if idb_path not in selected_functions: print("[!] Error! IDB path (%s) not in %s" % (idb_path, input_json)) - idc.Exit(1) + idaapi.qexit(1) fva_list = selected_functions[idb_path] print("[D] Found %d addresses" % len(fva_list)) run_acfg_features(idb_path, fva_list, output_dir) - idc.Exit(0) + idaapi.qexit(0) diff --git a/IDA_scripts/IDA_acfg_features/core/bb_features.py b/IDA_scripts/IDA_acfg_features/core/bb_features.py index fe05060..53963e2 100644 --- a/IDA_scripts/IDA_acfg_features/core/bb_features.py +++ b/IDA_scripts/IDA_acfg_features/core/bb_features.py @@ -30,7 +30,7 @@ import idautils -from architecture import ARCH_MNEM +from .architecture import ARCH_MNEM def get_bb_strings(bb, string_list): diff --git a/IDA_scripts/IDA_acfg_features/core/ff_features.py b/IDA_scripts/IDA_acfg_features/core/ff_features.py index aa8a30f..c6a4778 100644 --- a/IDA_scripts/IDA_acfg_features/core/ff_features.py +++ b/IDA_scripts/IDA_acfg_features/core/ff_features.py @@ -56,6 +56,8 @@ def get_size_local_vars(fva): Return: the size of local variables """ + if hasattr(idc, 'get_func_attr'): + return idc.get_func_attr(fva, idc.FUNCATTR_FRSIZE) return idc.GetFrameLvarSize(fva) diff --git a/IDA_scripts/IDA_acfg_features/test_acfg_features.py b/IDA_scripts/IDA_acfg_features/test_acfg_features.py index 26e3c1b..9276494 100644 --- a/IDA_scripts/IDA_acfg_features/test_acfg_features.py +++ b/IDA_scripts/IDA_acfg_features/test_acfg_features.py @@ -41,13 +41,18 @@ class TestIDAAcfgFeatures(unittest.TestCase): + maxDiff = None + def remove_elaspesed_time(self, jd): """Elapsed time will be different in any run.""" for idb, addr in jd.items(): for va in addr: - if isinstance(jd[idb][va], dict) \ - and 'elapsed_time' in jd[idb][va]: + if 'elapsed_time' in jd[idb][va]: jd[idb][va]['elapsed_time'] = -1 + if 'nodes' in jd[idb][va]: + jd[idb][va]['nodes'] = sorted(jd[idb][va]['nodes']) + if 'edges' in jd[idb][va]: + jd[idb][va]['edges'] = sorted(jd[idb][va]['edges']) return jd def test_acfg_features(self): diff --git a/IDA_scripts/IDA_acfg_features/test_large_acfg_features.py b/IDA_scripts/IDA_acfg_features/test_large_acfg_features.py index fed00f5..d532bc8 100644 --- a/IDA_scripts/IDA_acfg_features/test_large_acfg_features.py +++ b/IDA_scripts/IDA_acfg_features/test_large_acfg_features.py @@ -41,13 +41,18 @@ class TestIDAAcfgFeatures(unittest.TestCase): + maxDiff = None + def remove_elaspesed_time(self, jd): """Elapsed time will be different in any run.""" for idb, addr in jd.items(): for va in addr: - if isinstance(jd[idb][va], dict) \ - and 'elapsed_time' in jd[idb][va]: + if 'elapsed_time' in jd[idb][va]: jd[idb][va]['elapsed_time'] = -1 + if 'nodes' in jd[idb][va]: + jd[idb][va]['nodes'] = sorted(jd[idb][va]['nodes']) + if 'edges' in jd[idb][va]: + jd[idb][va]['edges'] = sorted(jd[idb][va]['edges']) return jd def test_acfg_features_large(self): diff --git a/IDA_scripts/IDA_flowchart/IDA_flowchart.py b/IDA_scripts/IDA_flowchart/IDA_flowchart.py index 635a6cf..41ca9c7 100644 --- a/IDA_scripts/IDA_flowchart/IDA_flowchart.py +++ b/IDA_scripts/IDA_flowchart/IDA_flowchart.py @@ -117,7 +117,7 @@ def get_function_hashopcodes(fva): # Create a string with the opcodes opc_string = ''.join(opc_list) - opc_string = opc_string.upper() + opc_string = opc_string.encode('utf-8').upper() # Get the sha256 hash hashopcodes = hashlib.sha256(opc_string).hexdigest() @@ -183,15 +183,15 @@ def analyze_functions(idb_path, output_csv): if __name__ == "__main__": if not idaapi.get_plugin_options("flowchart"): print("[!] -Oflowchart option is missing") - idc.Exit(1) + idaapi.qexit(1) plugin_options = idaapi.get_plugin_options("flowchart").split(':') if len(plugin_options) != 2: print("[!] -Oflowchart:IDB_PATH:OUTPUT_CSV is required") - idc.Exit(1) + idaapi.qexit(1) idb_path = plugin_options[0] output_csv = plugin_options[1] analyze_functions(idb_path, output_csv) - idc.Exit(0) + idaapi.qexit(0) diff --git a/Models/Catalog1/IDA_catalog1.py b/Models/Catalog1/IDA_catalog1.py index d72068d..8f14d20 100644 --- a/Models/Catalog1/IDA_catalog1.py +++ b/Models/Catalog1/IDA_catalog1.py @@ -111,12 +111,12 @@ def run_catalog1(idb_path, fva_list, sig_size, output_csv): if __name__ == '__main__': if not idaapi.get_plugin_options("catalog1"): print("[!] -Ocatalog1 option is missing") - idc.Exit(1) + idaapi.qexit(1) plugin_options = idaapi.get_plugin_options("catalog1").split(':') if len(plugin_options) != 4: print("[!] -Ocatalog1:INPUT_JSON:IDB_PATH:SIG_SIZE:OUTPUT_CSV is required") - idc.Exit(1) + idaapi.qexit(1) input_json = plugin_options[0] idb_path = plugin_options[1] @@ -128,10 +128,10 @@ def run_catalog1(idb_path, fva_list, sig_size, output_csv): if idb_path not in selected_functions: print("[!] Error! IDB path (%s) not in %s" % (idb_path, input_json)) - idc.Exit(1) + idaapi.qexit(1) fva_list = selected_functions[idb_path] print("[D] Found %d addresses" % len(fva_list)) run_catalog1(idb_path, fva_list, sig_size, output_csv) - idc.Exit(0) + idaapi.qexit(0) diff --git a/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py b/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py index daf1ce4..1f9b875 100644 --- a/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py +++ b/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py @@ -32,7 +32,6 @@ import idc import json import networkx as nx -import ntpath import os import pickle import time @@ -196,7 +195,7 @@ def run_codeCMR(idb_path, fva_list, output_dir): if not os.path.isdir(output_dir): os.mkdir(output_dir) - output_name = ntpath.basename( + output_name = os.path.basename( idb_path.replace(".i64", "").replace(".idb", "")) output_name += "_codeCMR.pkl" output_path = os.path.join(output_dir, output_name) @@ -233,12 +232,12 @@ def run_codeCMR(idb_path, fva_list, output_dir): if __name__ == '__main__': if not idaapi.get_plugin_options("codeCMR"): print("[!] -OcodeCMR option is missing") - idc.Exit(1) + idaapi.qexit(1) plugin_options = idaapi.get_plugin_options("codeCMR").split(":") if len(plugin_options) != 3: print("[!] -Ofss:INPUT_JSON:IDB_PATH:OUTPUT_DIR") - idc.Exit(1) + idaapi.qexit(1) input_json = plugin_options[0] idb_path = plugin_options[1] @@ -249,7 +248,7 @@ def run_codeCMR(idb_path, fva_list, output_dir): if idb_path not in selected_functions: print("[!] Error! IDB path (%s) not in %s" % (idb_path, input_json)) - idc.Exit(1) + idaapi.qexit(1) fva_list = selected_functions[idb_path] print("[D] Found %d addresses" % len(fva_list)) diff --git a/Models/functionsimsearch/IDA_fss/IDA_fss.py b/Models/functionsimsearch/IDA_fss/IDA_fss.py index eee8330..9367822 100644 --- a/Models/functionsimsearch/IDA_fss/IDA_fss.py +++ b/Models/functionsimsearch/IDA_fss/IDA_fss.py @@ -38,7 +38,6 @@ import idautils import idc import json -import ntpath import os import traceback @@ -62,7 +61,7 @@ def initialize_capstone(): https://github.com/williballenthin/python-idb/blob/ 2de7df8356ee2d2a96a795343e59848c1b4cb45b/idb/idapython.py#L874 """ - procname = idaapi.get_inf_structure().procName.lower() + procname = idaapi.get_inf_structure().procname.lower() bitness = get_bitness() md = None prefix = "UNK_" @@ -102,7 +101,7 @@ def initialize_capstone(): def get_call_mnemonics(): """Return different call instructions based on the arch.""" - procname = idaapi.get_inf_structure().procName.lower() + procname = idaapi.get_inf_structure().procname.lower() print('[D] procName = {}'.format(procname)) # Default choice @@ -268,7 +267,7 @@ def get_flowgraph_from(address, use_capstone): for ii in idautils.Heads(block.start_ea, block.end_ea): instructions.append(( ii, - idc.GetMnem(ii), + idaapi.ua_mnem(ii), # FIXME: only two operands? # It's ok for x86/64 but it will not work on other archs. (idc.print_operand(ii, 0).replace("+var_", "-0x"), @@ -308,7 +307,7 @@ def run_fss(idb_path, fva_list, output_dir, use_capstone): if not os.path.isdir(output_dir): os.mkdir(output_dir) - out_json_name = ntpath.basename(idb_path.replace(".i64", "")) + out_json_name = os.path.basename(idb_path.replace(".i64", "")) out_json_name += "_Capstone_{}_fss.json".format(use_capstone) out_json_path = os.path.join(output_dir, out_json_name) @@ -335,12 +334,12 @@ def run_fss(idb_path, fva_list, output_dir, use_capstone): if __name__ == '__main__': if not idaapi.get_plugin_options("fss"): print("[!] -Ofss option is missing") - idc.Exit(1) + idaapi.qexit(1) plugin_options = idaapi.get_plugin_options("fss").split(":") if len(plugin_options) != 4: print("[!] -Ofss:INPUT_JSON:IDB_PATH:OUTPUT_DIR:USE_CAPSTONE") - idc.Exit(1) + idaapi.qexit(1) input_json = plugin_options[0] idb_path = plugin_options[1] @@ -356,10 +355,10 @@ def run_fss(idb_path, fva_list, output_dir, use_capstone): if idb_path not in selected_functions: print("[!] Error! IDB path (%s) not in %s" % (idb_path, input_json)) - idc.Exit(1) + idaapi.qexit(1) fva_list = selected_functions[idb_path] print("[D] Found %d addresses" % len(fva_list)) run_fss(idb_path, fva_list, output_dir, use_capstone) - idc.Exit(0) + idaapi.qexit(0)