From fa02a1863d43bd2b7db50dac8c1e3a6f97bf68f8 Mon Sep 17 00:00:00 2001 From: Josh Bundt Date: Sun, 30 Oct 2022 20:43:18 +0000 Subject: [PATCH 1/4] IDA pro compatibility tweaks - use universal `idaapi.qexit` - ensure byte/string compability for Python 2/3 - sort nodes/edges to help tests pass With these changes, the unit tests for the IDA scripts pass for IDA Pro versions 7.2 - 8.1. This is using Python2 for v7.2 and Python3 for v7.4+, but they should be interchangable. --- IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py | 10 +++++----- IDA_scripts/IDA_acfg_disasm/test_acfg_disasm.py | 11 +++++++++-- IDA_scripts/IDA_acfg_disasm/test_large_acfg_disasm.py | 11 +++++++++-- IDA_scripts/IDA_acfg_features/IDA_acfg_features.py | 8 ++++---- IDA_scripts/IDA_acfg_features/core/bb_features.py | 2 +- IDA_scripts/IDA_acfg_features/core/ff_features.py | 2 ++ IDA_scripts/IDA_acfg_features/test_acfg_features.py | 9 +++++++-- .../IDA_acfg_features/test_large_acfg_features.py | 9 +++++++-- IDA_scripts/IDA_flowchart/IDA_flowchart.py | 8 ++++---- 9 files changed, 48 insertions(+), 22 deletions(-) diff --git a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py index edb2b8b..c76f3ec 100644 --- a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py +++ b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py @@ -205,7 +205,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 @@ -280,12 +280,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 +296,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..c3a68b3 100644 --- a/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py +++ b/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py @@ -273,12 +273,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 +289,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) From a048d926e8d9ba2b0bf3524a68018b9e4ef91275 Mon Sep 17 00:00:00 2001 From: Josh Bundt Date: Wed, 16 Nov 2022 14:13:47 +0000 Subject: [PATCH 2/4] remove Windows specific `ntpath` usage --- IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py | 3 +-- IDA_scripts/IDA_acfg_features/IDA_acfg_features.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py index c76f3ec..ec6bf8c 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 @@ -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) diff --git a/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py b/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py index c3a68b3..e26a18b 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 @@ -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) From 025d74feaaf3cf92d81870b0319638fd5a6ae716 Mon Sep 17 00:00:00 2001 From: Josh Bundt Date: Thu, 17 Nov 2022 03:51:15 +0000 Subject: [PATCH 3/4] remove two more Windows ntpath uses --- Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py | 3 +-- Models/functionsimsearch/IDA_fss/IDA_fss.py | 3 +-- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py b/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py index daf1ce4..c878ecc 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) diff --git a/Models/functionsimsearch/IDA_fss/IDA_fss.py b/Models/functionsimsearch/IDA_fss/IDA_fss.py index eee8330..84a5e08 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 @@ -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) From d9aa9f6ab527f8f4b150d1ffcf035508f6a02b86 Mon Sep 17 00:00:00 2001 From: Josh Bundt Date: Thu, 17 Nov 2022 20:51:17 +0000 Subject: [PATCH 4/4] replace more deprecated and old IDA python code - `idaapi.ua_mnem` is available since v6.1 - `idaapi.qexit` is available since v6.1 - `procName` is not deprecated and issues a warning now. `info.procname` is available in versions 7+ --- IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py | 2 +- IDA_scripts/IDA_acfg_features/IDA_acfg_features.py | 2 +- Models/Catalog1/IDA_catalog1.py | 8 ++++---- Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py | 6 +++--- Models/functionsimsearch/IDA_fss/IDA_fss.py | 14 +++++++------- 5 files changed, 16 insertions(+), 16 deletions(-) diff --git a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py index ec6bf8c..9aea15d 100644 --- a/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py +++ b/IDA_scripts/IDA_acfg_disasm/IDA_acfg_disasm.py @@ -221,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) diff --git a/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py b/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py index e26a18b..ed20b1e 100644 --- a/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py +++ b/IDA_scripts/IDA_acfg_features/IDA_acfg_features.py @@ -223,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) 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 c878ecc..1f9b875 100644 --- a/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py +++ b/Models/CodeCMR/IDA_CodeCMR/IDA_codeCMR.py @@ -232,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] @@ -248,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 84a5e08..9367822 100644 --- a/Models/functionsimsearch/IDA_fss/IDA_fss.py +++ b/Models/functionsimsearch/IDA_fss/IDA_fss.py @@ -61,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_" @@ -101,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 @@ -267,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"), @@ -334,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] @@ -355,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)