backup

Genius3/raw-feature-extractor/convert_pkl_to_json.py

@@ -3,58 +3,90 @@ import pickle as pk
 import re
 import json
 import os
+from tqdm import tqdm
 
 
-def convert():
+def convert(start, end):
-    # for workflow in range(0, 20):
+    for workflow in range(start, end):
-        workflow = 0
+        # workflow = 0
         cfg_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_cfg".format(workflow)
         output_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_json".format(workflow)
-        dot_path = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_dot".format(workflow)
+        dot_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_dot".format(workflow)
+        log_path = "D:\\hkn\\infected\\datasets\\logging\\convert_log{}.log".format(workflow)
+        process_log_path = "D:\\hkn\\infected\\datasets\\logging\\convert_process_log{}.log".format(workflow)
 
-        for cfg in os.listdir(cfg_dir):
+        with open(log_path, 'a+') as log, open(process_log_path, 'a+') as process_log:
-            name = cfg[:-4]  # 纯文件名，不带后缀
+            logged = log.readline()
-            cfg_file = open(os.path.join(cfg_dir, name + '.ida'), 'r')
+            if logged == '':
-            data = pk.load(cfg_file)
+                log_index = 0
-            cfg_file.close()
+            else:
+                log_index = int(logged)
 
-            # 打开dot文件获取fcg
+            for index, cfg in enumerate(tqdm(os.listdir(cfg_dir))):
-            raw_function_edges = []
+                if index < log_index:
-            with open(os.path.join(dot_path, name + '.dot'), 'r') as dot:
+                    continue
-                for line in dot:
-                    if '->' in line:
-                        raw_function_edges.append(re.findall(r'\b\d+\b', line))
 
-            # 为当前pe文件创建json对象
+                name = cfg[:-4]  # 纯文件名，不带后缀
-            json_obj = {
+                cfg_file = open(os.path.join(cfg_dir, name + '.ida'), 'r')
-                'hash': data.binary_name[11:],
+                try:
-                'function_number': data.raw_graph_list.__len__(),
+                    data = pk.load(cfg_file)
-                'function_edges': [[d[0] for d in raw_function_edges], [d[1] for d in raw_function_edges]],
+                except EOFError:
-                'acfg_list': [],
+                    process_log.write("index {}, {} process failed. EOFError occurred.\n".format(index, cfg))
-                'function_names': []
+                    continue
-            }
+                except ValueError:
-            # 读取pkl文件，一个acfg由一个函数分解而来
+                    process_log.write("index {}, {} process failed. ValueError occurred.\n".format(index, cfg))
-            for acfg in data.raw_graph_list:
+                    continue
-                # 这里2是因为Genius框架提取特征时将后代数量放在2
-                offspring = [d.get('v')[2] for d in acfg.g.node.values()]
-                # 将后代数量的特征放入bb_features中
-                for i, f in enumerate(acfg.bb_features):
-                    f.append(offspring[i])
 
-                acfg_item = {
+                cfg_file.close()
-                    'block_number': acfg.g.__len__(),
-                    'block_edges': [[d[0] for d in acfg.g.edges], [d[1] for d in acfg.g.edges]],
-                    'block_features': acfg.bb_features
-                }
 
-                json_obj['acfg_list'].append(acfg_item)
+                dot_file_path = os.path.join(dot_dir, name + '.dot')
-                json_obj['function_names'].append(acfg.funcname)
+                if not os.path.exists(dot_file_path):
+                    process_log.write("index {}, {} process failed. dot file not exists.\n".format(index, cfg))
+                else:
+                    # 打开dot文件获取fcg
+                    raw_function_edges = []
+                    with open(dot_file_path, 'r') as dot:
+                        for line in dot:
+                            if '->' in line:
+                                raw_function_edges.append(re.findall(r'\b\d+\b', line))
 
-            # 将结果写入json本地文件
+                    # 为当前pe文件创建json对象
-            result = json.dumps(json_obj)
+                    json_obj = {
-            with open(os.path.join(output_dir, name + '.jsonl'), 'w') as out:
+                        'hash': data.binary_name[11:],
-                out.write(result)
+                        'function_number': data.raw_graph_list.__len__(),
+                        'function_edges': [[d[0] for d in raw_function_edges], [d[1] for d in raw_function_edges]],
+                        'acfg_list': [],
+                        'function_names': []
+                    }
+                    # 读取pkl文件，一个acfg由一个函数分解而来
+                    for acfg in data.raw_graph_list:
+                        # 这里2是因为Genius框架提取特征时将后代数量放在2
+                        offspring = [d.get('v')[2] for d in acfg.g.node.values()]
+                        # 将后代数量的特征放入bb_features中
+                        for i, offs in enumerate(offspring):
+                            acfg.bb_features[i].append(offs)
+
+                        acfg_item = {
+                            'block_number': acfg.g.__len__(),
+                            'block_edges': [[d[0] for d in acfg.g.edges], [d[1] for d in acfg.g.edges]],
+                            'block_features': acfg.bb_features
+                        }
+
+                        json_obj['acfg_list'].append(acfg_item)
+                        json_obj['function_names'].append(acfg.funcname)
+
+                    # 将结果写入json本地文件
+                    result = json.dumps(json_obj, ensure_ascii=False)
+
+                    with open(os.path.join(output_dir, name + '.jsonl'), 'w') as out:
+                        out.write(result)
+
+                    log.truncate(0)
+                    log.seek(0)
+                    log.write(str(index))
+                    log.flush()
+                    process_log.write("index {}, {} process done.\n".format(index, cfg))
 
 
 if __name__ == '__main__':
-    convert()
+    convert(20, 35)

Genius3/raw-feature-extractor/ida_batch.py

@@ -6,6 +6,8 @@ from tqdm import tqdm
 import time
 
 # 单个pe文件处理超时/s
+# 多次处理，一批数据中只有少量文件会超时
+# 所有数据处理完成后可以对这些数据再进行一次更长超时时间的处理，若仍然超时则放弃
 TIMEOUT = 60
 
 
@@ -14,7 +16,7 @@ def call_preprocess(cmd_line):
 
 
 def batch_mode():
-    for workflow in range(0, 1):
+    for workflow in range(1, 20):
         # workflow = 0
         pe_dir = 'D:\\hkn\\infected\\datasets\\virusshare_infected{}'.format(workflow)
         # for test
@@ -36,7 +38,8 @@ def batch_mode():
                 # for test
                 # cmd_line = r'idaq64 -c -A -S"D:\hkn\project_folder\Gencoding3\Genius3\raw-feature-extractor\preprocessing_ida.py {}" -oF:\iout {}'.format(
                 #     workflow, os.path.join(pe_dir, pe))
-                cmd_line = r'idaq64 -c -A -S"D:\hkn\project_folder\Gencoding3\Genius3\raw-feature-extractor\preprocessing_ida.py {}" -oF:\iout {}'.format(workflow, os.path.join(pe_dir, pe))
+                cmd_line = r'idaq64 -c -A -S"D:\hkn\project_folder\Gencoding3\Genius3\raw-feature-extractor\preprocessing_ida.py {}" -oF:\iout {}'.format(
+                    workflow, os.path.join(pe_dir, pe))
 
                 p = multiprocessing.Process(target=call_preprocess, args=[cmd_line])
                 p.start()
@@ -51,7 +54,8 @@ def batch_mode():
 
                 if flag_kill:
                     subprocess.call('taskkill /im idaq64.exe /f')
-                    process_log.write("index {}, {} in workflow {} stuck, process terminated.\n".format(index, pe, workflow))
+                    process_log.write(
+                        "index {}, {} in workflow {} stuck, process terminated.\n".format(index, pe, workflow))
                 else:
                     # 正常运行结束
                     log.truncate(0)
@@ -59,6 +63,14 @@ def batch_mode():
                     log.write(str(index))
                     log.flush()
                     process_log.write("index {}, {} process done.\n".format(index, pe))
+    # 一次workflow结束后将所有副产物删除
+    delete_output()
+
+
+def delete_output():
+    out_dir = 'F:\\iout'
+    os.rmdir(out_dir)
+    os.mkdir(out_dir)
 
 
 # 注意：该py文件必须放在IDA的根目录下，且必须使用cmd命令执行，否则无法链接到python库

Genius3/raw-feature-extractor/read_idaFILE.py

@@ -15,7 +15,7 @@ def print_obj(obj):
 # sub_10F20 308  反编译代码有字符串，但是这个特征提取里没有字符串 constant，可能是间接引用的，不识别。看了下所有函数的特征，几乎都没有字符串常量，可能都是写在别的地方然后引用的。
 # sub_166C4 393
 if __name__ == '__main__':
-    testpath = "D:\\hkn\\infected\\datasets\\virusshare_infected0_cfg\\VirusShare_cd53c6637ca75ac5fc1cbe6d2ced41a1.ida"
+    testpath = "D:\\hkn\\infected\\datasets\\virusshare_infected11_cfg\\VirusShare_5c088a2a6e0391b7c6ab22e4648eab3a.ida"
     fr = open(testpath, 'r')
     data = pickle.load(fr) #一个二进制文件的acfgs
     fr.close()

Genius3/raw-feature-extractor/test.py

@@ -63,8 +63,22 @@ def clock():
         subprocess.call('taskkill /im idaq64.exe /f')
 
 
+def delete_error():
+    for workflow in range(0, 35):
+        convert_log_path = "D:\\hkn\\infected\\datasets\\logging\\convert_process_log{}.log".format(workflow)
+        json_dir = "D:\\hkn\\infected\\datasets\\virusshare_infected{}_json".format(workflow)
+
+        with open(convert_log_path, 'r') as log:
+            for line in log:
+                if 'Error occurred' in line:
+                    name = line[line.find(',') + 2: line.find('.')] + '.jsonl'
+                    # print(os.path.join(json_dir, name))
+                    if os.path.exists(os.path.join(json_dir, name)):
+                        os.remove(os.path.join(json_dir, name))
+
+
 if __name__ == '__main__':
     # gen_dir()
     # change_max_item_lines()
-    subprocess.call('taskkill /im idaq64.exe /f')
+    # subprocess.call('taskkill /im idaq64.exe /f')
-
+    delete_error()