diff --git a/src/picklescan/scanner.py b/src/picklescan/scanner.py index 62e9210..e27b58f 100644 --- a/src/picklescan/scanner.py +++ b/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 = [] diff --git a/tests/data/malicious11.pkl b/tests/data/malicious11.pkl new file mode 100644 index 0000000..f8645b4 --- /dev/null +++ b/tests/data/malicious11.pkl @@ -0,0 +1,12 @@ +Vos +p2 +0Vsystem +p3 +0Vtorch +p0 +0VLongStorage +p1 +0g2 +g3 +“(Vcat flag.txt +tR. \ No newline at end of file diff --git a/tests/test_scanner.py b/tests/test_scanner.py index 41aa253..04f3457 100644 --- a/tests/test_scanner.py +++ b/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)