From 0ae787002c27aae195c5b310c0aca4a91ba8f209 Mon Sep 17 00:00:00 2001
From: tritium0041 <haichuanyu243@gmail.com>
Date: Tue, 14 May 2024 21:31:31 +0800
Subject: [PATCH] =?UTF-8?q?update:=E5=AE=8C=E5=96=84=E8=B0=83=E7=94=A8?=
 =?UTF-8?q?=E6=96=B9=E5=BC=8F=EF=BC=8C=E5=88=A0=E9=99=A4=E5=A4=9A=E4=BD=99?=
 =?UTF-8?q?=E4=BB=A3=E7=A0=81?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 detection/backdoor_detection.py |  7 +++++
 detection/pickle_detection.py   | 48 +++++++++++++++------------------
 2 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/detection/backdoor_detection.py b/detection/backdoor_detection.py
index ef3e32b..e50963c 100644
--- a/detection/backdoor_detection.py
+++ b/detection/backdoor_detection.py
@@ -3,6 +3,8 @@ from typing import Dict, List, Tuple
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet
 from reportlab.platypus import Paragraph, Spacer, SimpleDocTemplate
+
+from detection.pickle_detection import pickleDataDetection
 from .Regexdetection import find_dangerous_functions
 from .GPTdetection import detectGPT
 from .utils import *
@@ -176,6 +178,9 @@ def main():
     parser.add_argument(
         "-m", "--mode", help="Mode of operation:[regex,llm]", default="regex"
     )
+    parser.add_argument(
+        "-p","--pickle",help="analyze the pickle file",default=None
+    )
     args = parser.parse_args()
     output_format = "txt"  # Default output format
     output_file = None
@@ -191,6 +196,8 @@ def main():
             )
             output_file = args.output.rsplit(".", 1)[0] + ".txt"
     # 如果未指定输出文件，则输出到 stdout；否则写入文件
+    if args.pickle:
+        pickleDataDetection(args.pickle, output_file)
     process_path(args.path, output_format, args.mode, output_file)
 
 
diff --git a/detection/pickle_detection.py b/detection/pickle_detection.py
index b833c1c..4730f79 100644
--- a/detection/pickle_detection.py
+++ b/detection/pickle_detection.py
@@ -1,4 +1,5 @@
 import io
+import json
 import os
 import pickletools
 import pickle
@@ -87,10 +88,9 @@ class pickleScanner():
         self.file = file
 
     def find_class(self, module, name):
-        print(module, name)
         if module.decode() in dangerous_modules or name.decode() in dangerous_names:
             # self.maliciousCount += 1
-            self.maliciousModule.append((module, name))
+            self.maliciousModule.append((module.decode(), name.decode()))
 
     def load(self):
         self._unframer = _Unframer(self._file_read, self._file_readline)
@@ -128,30 +128,24 @@ class pickleScanner():
 
 
     def output(self):
-        if self.ReduceCount > 0 or len(self.maliciousModule) > 0:
-            print("The pickle file maybe contains malicious code")
-            print(f"The number of REDUCE opcodes is {self.ReduceCount}")
-            print("The malicious options are: ", self.maliciousModule)
-        else:
-            print("The pickle file is safe")
+        return {
+        "ReduceCount": self.ReduceCount,
+        "maliciousModule": self.maliciousModule
+    }
 
+def pickleDataDetection(file,output_file=None):
+    '''
+    :param file: pickle file path
+    '''
+    with open(file, "rb") as file:
+        pickscan = pickleScanner(file)
+        pickscan.load()
+    res = pickscan.output()
+    if output_file:
+        with open(output_file, "w") as file:
+            file.writelines(json.dumps(res))
+    else:
+        print(json.dumps(res))
 
-
-
-class test:
-    a = 1
-    b = 2
-    def __reduce__(self):
-        return (__import__("os").system,('calc',))
-
-
-data = pickle.dumps(test(),protocol=2)
-print(data)
-print(pickletools.dis(data))
-with open("test.pkl", "wb") as file:
-    file.write(data)
-
-with open("test.pkl", "rb") as file:
-    pickscan = pickleScanner(file)
-    pickscan.load()
-pickscan.output()
\ No newline at end of file
+if __name__ == '__main__':
+    pickleDataDetection("test.pkl")
\ No newline at end of file