Handling INST opcode

mmaitre314 · Sep 23, 2023 · bfa1fa0 · bfa1fa0
1 parent efcebea
commit bfa1fa0
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 5 deletions.
diff --git a/setup.cfg b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = picklescan
-version = 0.0.11
+version = 0.0.12
 author = Matthieu Maitre
 author_email = [email protected]
 description = Security scanner detecting Python Pickle files performing suspicious actions

diff --git a/src/picklescan/scanner.py b/src/picklescan/scanner.py
@@ -192,7 +192,7 @@ def _list_globals(data: IO[bytes], multiple_pickles=True) -> Set[Tuple[str, str]
             if op_name in ["MEMOIZE", "PUT", "BINPUT", "LONG_BINPUT"] and n > 0:
                 memo[len(memo)] = ops[n - 1][1]
 
-            if op_name == "GLOBAL":
+            if op_name in ("GLOBAL", "INST"):
                 globals.add(tuple(op_value.split(" ", 1)))
             elif op_name == "STACK_GLOBAL":
                 values = []

diff --git a/tests/data/malicious10.pkl b/tests/data/malicious10.pkl
@@ -0,0 +1,4 @@
+(S'raise RuntimeError("Injection running")'
+i__builtin__
+exec
+.
diff --git a/tests/test_scanner.py b/tests/test_scanner.py
@@ -198,6 +198,13 @@ def initialize_pickle_files():
         + b"dict\nS'get'\ntRp103\n0c__builtin__\napply\n(g103\n(g100\nS'picklesmashed'\nltRp104\n0g104\n.",
     )
 
+    # Malicious Pickle with an INST opcode
+    #     0: (    MARK
+    #     1: S        STRING     'raise RuntimeError("Injection running")'
+    #    44: i        INST       '__builtin__ exec' (MARK at 0)
+    #    62: .    STOP
+    initialize_data_file(f"{_root_path}/data/malicious10.pkl", b'(S\'raise RuntimeError("Injection running")\'\ni__builtin__\nexec\n.')
+
     initialize_data_file(f"{_root_path}/data/malicious3.pkl", malicious3_pickle_bytes)
     initialize_pickle_file(f"{_root_path}/data/malicious4.pickle", Malicious4(), 4)
     initialize_pickle_file(f"{_root_path}/data/malicious5.pickle", Malicious5(), 4)
@@ -413,6 +420,13 @@ def test_scan_file_path():
         scan_file_path(f"{_root_path}/data/malicious9.pkl"), malicious9
     )
 
+    malicious10 = ScanResult(
+        [Global("__builtin__", "exec", SafetyLevel.Dangerous)], 1, 1, 1
+    )
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious10.pkl"), malicious10
+    )
+
     bad_pytorch = ScanResult([], 0, 0, 0, True)
     compare_scan_results(
         scan_file_path(f"{_root_path}/data/bad_pytorch.pt"), bad_pytorch
@@ -451,10 +465,11 @@ def test_scan_directory_path():
             Global("_rebuild_tensor", "unknown", SafetyLevel.Dangerous),
             Global("torch._utils", "_rebuild_tensor", SafetyLevel.Suspicious),
             Global("torch", "_utils", SafetyLevel.Suspicious),
+            Global("__builtin__", "exec", SafetyLevel.Dangerous),
         ],
-        21,
-        19,
-        16,
+        22,
+        20,
+        17,
     )
     compare_scan_results(scan_directory_path(f"{_root_path}/data/"), sr)