Fix handling of MEMOIZE and add CTF test

src/picklescan/scanner.py

@@ -189,10 +189,11 @@ def _list_globals(data: IO[bytes], multiple_pickles=True) -> Set[Tuple[str, str]
             op_name = op[0].name
             op_value = op[1]
 
-            if op_name in ["MEMOIZE", "PUT", "BINPUT", "LONG_BINPUT"] and n > 0:
+            if op_name == "MEMOIZE" and n > 0:
+                memo[len(memo)] = ops[n - 1][1]
+            elif op_name in ["PUT", "BINPUT", "LONG_BINPUT"] and n > 0:
                 memo[op_value] = ops[n - 1][1]
-
-            if op_name in ("GLOBAL", "INST"):
+            elif op_name in ("GLOBAL", "INST"):
                 globals.add(tuple(op_value.split(" ", 1)))
             elif op_name == "STACK_GLOBAL":
                 values = []

tests/data/malicious11.pkl

@@ -0,0 +1,12 @@
+Vos
+p2
+0Vsystem
+p3
+0Vtorch
+p0
+0VLongStorage
+p1
+0g2
+g3
+�(Vcat flag.txt
+tR.

tests/test_scanner.py

@@ -208,6 +208,36 @@ def initialize_pickle_files():
         b"(S'raise RuntimeError(\"Injection running\")'\ni__builtin__\nexec\n.",
     )
 
+    # Malicious Pickle from Capture-the-Flag challenge 'Misc/Safe Pickle' at https://imaginaryctf.org/Challenges
+    # GitHub Issue: https://github.com/mmaitre314/picklescan/issues/22
+    initialize_data_file(
+        f"{_root_path}/data/malicious11.pkl",
+        b"".join(
+            [
+                pickle.UNICODE + b"os\n",
+                pickle.PUT + b"2\n",
+                pickle.POP,
+                pickle.UNICODE + b"system\n",
+                pickle.PUT + b"3\n",
+                pickle.POP,
+                pickle.UNICODE + b"torch\n",
+                pickle.PUT + b"0\n",
+                pickle.POP,
+                pickle.UNICODE + b"LongStorage\n",
+                pickle.PUT + b"1\n",
+                pickle.POP,
+                pickle.GET + b"2\n",
+                pickle.GET + b"3\n",
+                pickle.STACK_GLOBAL,
+                pickle.MARK,
+                pickle.UNICODE + b"cat flag.txt\n",
+                pickle.TUPLE,
+                pickle.REDUCE,
+                pickle.STOP,
+            ]
+        ),
+    )
+
     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)
@@ -469,10 +499,11 @@ def test_scan_directory_path():
             Global("torch._utils", "_rebuild_tensor", SafetyLevel.Suspicious),
             Global("torch", "_utils", SafetyLevel.Suspicious),
             Global("__builtin__", "exec", SafetyLevel.Dangerous),
+            Global("os", "system", SafetyLevel.Dangerous),
         ],
-        22,
-        20,
-        17,
+        23,
+        21,
+        18,
     )
     compare_scan_results(scan_directory_path(f"{_root_path}/data/"), sr)