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