IDA pro compatibility tweaks

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,20 +271,20 @@ 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)
 
 
 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)

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):

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):

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,20 +264,20 @@ 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)
 
 
 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)

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):

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)
 
 

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):

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):

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)

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)

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))

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)