From 40001cd1caa9e041b1bce1b80f3707056cd8be52 Mon Sep 17 00:00:00 2001
From: Matthieu Maitre <mmaitre314@users.noreply.github.com>
Date: Sat, 23 Sep 2023 11:57:50 -0700
Subject: [PATCH] Handle INST opcode (#21)

* Handling INST opcode
---
 setup.cfg                  |  2 +-
 src/picklescan/scanner.py  |  2 +-
 tests/data/malicious10.pkl |  4 ++++
 tests/test_scanner.py      | 24 +++++++++++++++++++++---
 4 files changed, 27 insertions(+), 5 deletions(-)
 create mode 100644 tests/data/malicious10.pkl

diff --git a/setup.cfg b/setup.cfg
index e1a2bdd..a331304 100644
--- 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 = mmaitre314@users.noreply.github.com
 description = Security scanner detecting Python Pickle files performing suspicious actions
diff --git a/src/picklescan/scanner.py b/src/picklescan/scanner.py
index 8412093..7363d09 100644
--- 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
new file mode 100644
index 0000000..fa94fbe
--- /dev/null
+++ b/tests/data/malicious10.pkl
@@ -0,0 +1,4 @@
+(S'raise RuntimeError("Injection running")'
+i__builtin__
+exec
+.
\ No newline at end of file
diff --git a/tests/test_scanner.py b/tests/test_scanner.py
index 843f749..41aa253 100644
--- a/tests/test_scanner.py
+++ b/tests/test_scanner.py
@@ -198,6 +198,16 @@ 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 +423,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 +468,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)