From 4b58fae27ea8b3ac275067bdda161ec057d198d4 Mon Sep 17 00:00:00 2001
From: Sean Morgan <sean@protectai.com>
Date: Tue, 28 May 2024 12:02:16 -0700
Subject: [PATCH] Add runpy to unsafe operators

---
 src/picklescan/scanner.py |  1 +
 tests/test_scanner.py     | 23 ++++++++++++++++++++---
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/src/picklescan/scanner.py b/src/picklescan/scanner.py
index 0539736..145bd49 100644
--- a/src/picklescan/scanner.py
+++ b/src/picklescan/scanner.py
@@ -114,6 +114,7 @@ def __str__(self) -> str:
     "socket": "*",
     "subprocess": "*",
     "sys": "*",
+    "runpy": "*",  # Includes runpy._run_code
     "operator": "attrgetter",  # Ex of code execution: operator.attrgetter("system")(__import__("os"))("echo pwned")
     "pickle": "*",
     "_pickle": "*",
diff --git a/tests/test_scanner.py b/tests/test_scanner.py
index 19847d4..c34948e 100644
--- a/tests/test_scanner.py
+++ b/tests/test_scanner.py
@@ -3,6 +3,7 @@
 import importlib
 import io
 import os
+import runpy
 import pickle
 import pytest
 import requests
@@ -83,6 +84,11 @@ def __reduce__(self):
         return pickle.loads, (b"I12345\n.",)  # Loads the value 12345
 
 
+class Malicious14:
+    def __reduce__(self):
+        return runpy._run_code, ("print('456')",)
+
+
 class HTTPResponse:
     def __init__(self, status, data=None):
         self.status = status
@@ -336,6 +342,9 @@ def initialize_pickle_files():
     initialize_pickle_file(
         f"{_root_path}/data/malicious13b.pkl", Malicious13(), 4
     )  # pickle module serialized as _pickle
+    initialize_pickle_file(
+        f"{_root_path}/data/malicious14.pkl", Malicious14(), 4
+    )  # runpy
 
     initialize_zip_file(
         f"{_root_path}/data/malicious1.zip",
@@ -552,6 +561,13 @@ def test_scan_file_path():
         scan_file_path(f"{_root_path}/data/bad_pytorch.pt"), bad_pytorch
     )
 
+    malicious14 = ScanResult(
+        [Global("runpy", "_run_code", SafetyLevel.Dangerous)], 1, 1, 1
+    )
+    compare_scan_results(
+        scan_file_path(f"{_root_path}/data/malicious14.pkl"), malicious14
+    )
+
 
 def test_scan_directory_path():
     sr = ScanResult(
@@ -578,6 +594,7 @@ def test_scan_directory_path():
             Global("requests.api", "get", SafetyLevel.Dangerous),
             Global("builtins", "eval", SafetyLevel.Dangerous),
             Global("builtins", "eval", SafetyLevel.Dangerous),
+            Global("runpy", "_run_code", SafetyLevel.Dangerous),
             Global("socket", "create_connection", SafetyLevel.Dangerous),
             Global("collections", "OrderedDict", SafetyLevel.Innocuous),
             Global("torch._utils", "_rebuild_tensor_v2", SafetyLevel.Innocuous),
@@ -594,9 +611,9 @@ def test_scan_directory_path():
             Global("_pickle", "loads", SafetyLevel.Dangerous),
             Global("_codecs", "encode", SafetyLevel.Suspicious),
         ],
-        scanned_files=27,
-        issues_count=25,
-        infected_files=22,
+        scanned_files=28,
+        issues_count=26,
+        infected_files=23,
         scan_err=True,
     )
     compare_scan_results(scan_directory_path(f"{_root_path}/data/"), sr)