From b2b6a2b7a6ba3be897c0256ffead07e607d9593b Mon Sep 17 00:00:00 2001
From: huihun <781165206@qq.com>
Date: Mon, 22 Apr 2024 13:14:01 +0800
Subject: [PATCH] =?UTF-8?q?=E5=88=A4=E6=96=AD=E7=A9=BA=E6=8C=87=E4=BB=A4?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 bert/obtain_inst_vec.py | 7 +++++--
 exe2json.py             | 1 -
 my_utils.py             | 2 ++
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/bert/obtain_inst_vec.py b/bert/obtain_inst_vec.py
index 40d7bf8..85606b2 100644
--- a/bert/obtain_inst_vec.py
+++ b/bert/obtain_inst_vec.py
@@ -105,8 +105,11 @@ def generate_inst_vec(inst, method="mean"):
 
 
 def bb2vec(item):
-    tmp = generate_inst_vec(item['opcode'], method="mean")
-    return item['addr'], list(np.mean(tmp.detach().numpy(), axis=0).astype(float))
+    try:
+        tmp = generate_inst_vec(item['opcode'], method="mean")
+        return True, item['addr'], list(np.mean(tmp.detach().numpy(), axis=0).astype(float))
+    except Exception as e:
+        return False, item['addr'], item['opcode'], e
 
 
 if __name__ == "__main__":
diff --git a/exe2json.py b/exe2json.py
index 3445d73..2a1a54f 100644
--- a/exe2json.py
+++ b/exe2json.py
@@ -247,7 +247,6 @@ def exe_to_json(file_path):
             logger.error(f"cfg错误：{cfg_Operation_message}")
         return
     # json写入
-    os.makedirs(output_path, exist_ok=True)
     save_json(os.path.join(output_path, file_fingerprint + '.jsonl'), json_obj)
     logger.info(f"解析完成，文件名{file_path}")
     return
diff --git a/my_utils.py b/my_utils.py
index b090d19..8194626 100644
--- a/my_utils.py
+++ b/my_utils.py
@@ -82,6 +82,8 @@ def multi_thread_disorder(func, thread_num=THREAD_FULL, **args):
 
 def save_json(filename, data):
     data = json.dumps(data)
+    # 确保路径存在
+    os.makedirs(os.path.dirname(filename), exist_ok=True)
     file = open(filename, 'w')
     file.write(data)
     file.close()