加入多线程与日志记录

Genius3/raw-feature-extractor/log_utils.py

@@ -0,0 +1,26 @@
+import logging
+import os
+
+
+def setup_logger(name, log_file, level=logging.INFO):
+    """Function setup as many loggers as you want"""
+
+    formatter = logging.Formatter('%(asctime)s %(levelname)s %(message)s')
+
+    handler = logging.FileHandler(log_file)
+    handler.setFormatter(formatter)
+
+    # Also add a stream handler for console output
+    stream_handler = logging.StreamHandler()
+    stream_handler.setFormatter(formatter)
+
+    logger = logging.getLogger(name)
+    logger.setLevel(level)
+    logger.addHandler(handler)
+    logger.addHandler(stream_handler)
+
+    # If the file already exists, clear its contents to start fresh
+    if os.path.exists(log_file):
+        open(log_file, 'w').close()
+
+    return logger

Genius3/raw-feature-extractor/preprocessing_ida.py

@@ -31,13 +31,13 @@ def preprocess():
     idaapi.autoWait()
 
     # 生成pe文件的cfg列表
-    cfgs = get_func_cfgs_c(FirstSeg())
+    # cfgs = get_func_cfgs_c(FirstSeg())
     # 将cfg保存为.ida
-    pickle.dump(cfgs, open(cfg_path, 'w'))
+    # pickle.dump(cfgs, open(cfg_path, 'w'))
 
     # 生成pe文件的fcg，保存为.dot文件
     # idc.GenCallGdl(gdl_path, 'Call Gdl', idc.CHART_GEN_GDL) 这个生成gdl文件，网上几乎找不到gdl这个格式
-    idc.GenCallGdl(gdl_path, 'Call Gdl', idaapi.CHART_GEN_DOT)
+    # idc.GenCallGdl(gdl_path, 'Call Gdl', idaapi.CHART_GEN_DOT)
 
     # 生成.asm文件
     idc.GenerateFile(idc.OFILE_ASM, asm_path, 0, idc.BADADDR, 0)

Genius3/raw-feature-extractor/thread.py

@@ -0,0 +1,86 @@
+# coding=utf-8
+import os
+import subprocess
+import threading
+import time
+
+
+# 设置最大并发线程数
+max_threads = 20
+
+# 创建一个锁和条件变量用于线程同步
+thread_lock = threading.Lock()
+active_threads = 0
+threads_completed = 0
+condition = threading.Condition(thread_lock)
+timer_event = threading.Event()
+
+def execute_command(cmd):
+    """
+    在子线程中执行给定的命令。
+    """
+    global active_threads, threads_completed
+
+    process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+    stdout, stderr = process.communicate()
+
+    with condition:
+        # print("Command: %s" % cmd)
+        if stdout:
+            print("stdout: %s" % stdout.decode('gbk'))
+        if stderr:
+            print("stderr: %s" % stderr.decode('gbk'))
+
+        # 当前线程完成任务后释放一个线程位置
+        active_threads -= 1
+        threads_completed += 1
+        condition.notify_all()
+        print threads_completed
+    return (stdout, stderr)  # 返回结果（可选）
+
+def timer_thread():
+    print 'start timer thread'
+    while True:
+        with condition:
+            # 每隔1秒检查一次并输出已完成命令的数量
+            while not timer_event.is_set() and threads_completed < len(commands):
+                timer_event.wait(1)
+                print "done file: %d" % threads_completed
+
+        # 如果所有命令都已完成，则停止计时器线程
+        if threads_completed == len(commands):
+            timer_event.set()
+            break
+
+
+if __name__ == '__main__':
+    # timer = threading.Thread(target=timer_thread)
+    # timer.start()
+    # timer_event.clear()
+
+    # 样本文件夹
+    sample_dir = "D:/bishe/dataset/sample_20230130_458"
+    # 创建并启动线程
+    commands = []
+
+    for file in os.listdir(sample_dir):
+        com = r'D:\IDA_Pro_v6.8\idaq64.exe -c -A -S"D:\bishe\Gencoding_KE\Genius3\raw-feature-extractor\preprocessing_ida.py 0" -oD:\bishe\dataset\out '
+        commands.append(com+"D:/bishe/dataset/sample_20230130_458/"+file)
+    threads = []
+    for cmd in commands:
+        while active_threads >= max_threads:
+            with condition:
+                # 等待有线程完成任务
+                condition.wait()
+
+        thread = threading.Thread(target=execute_command, args=(cmd,))
+        thread.start()
+        active_threads += 1
+        threads.append(thread)
+
+
+    # 等待所有线程完成
+    for thread in threads:
+        thread.join()
+
+    print("所有命令已执行完毕.")