Merge branch 'main' into dependabot/pip/rich-13.8.0

protectai · Sep 16, 2024 · f2a3699 · f2a3699
2 parents fe3e9ac + 2d63621
commit f2a3699
Show file tree

Hide file tree

Showing 2 changed files with 63 additions and 0 deletions.
diff --git a/modelscan/settings.py b/modelscan/settings.py
@@ -125,6 +125,9 @@ class SupportedModelFormats:
             ],
             "pty": "*",
             "pickle": "*",
+            "bdb": "*",
+            "pdb": "*",
+            "shutil": "*",
         },
         "HIGH": {
             "webbrowser": "*",  # Includes webbrowser.open()

diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py
@@ -1,4 +1,5 @@
 import aiohttp
+import bdb
 import http.client
 import importlib
 import io
@@ -88,6 +89,15 @@ def __reduce__(self) -> Any:
         return sys.exit, (0,)
 
 
+class Malicious15:
+    def __reduce__(self) -> Any:
+        bd = bdb.Bdb()
+        return bdb.Bdb.run, (
+            bd,
+            'import os\nos.system("whoami")',
+        )
+
+
 def malicious12_gen() -> bytes:
     p = pickle.PROTO + b"\x05"
 
@@ -272,6 +282,7 @@ def file_path(tmp_path_factory: Any) -> Any:
     initialize_pickle_file(f"{tmp}/data/malicious7.pkl", Malicious6(), 4)
     initialize_pickle_file(f"{tmp}/data/malicious8.pkl", Malicious7(), 4)
     initialize_pickle_file(f"{tmp}/data/malicious9.pkl", Malicious8(), 4)
+    initialize_pickle_file(f"{tmp}/data/malicious15.pkl", Malicious15(), 4)
 
     # Malicious Pickle from Capture-the-Flag challenge 'Misc/Safe Pickle' at https://imaginaryctf.org/Challenges
     # GitHub Issue: https://github.com/mmaitre314/picklescan/issues/22
@@ -1001,6 +1012,34 @@ def test_scan_pickle_operators(file_path: Any) -> None:
     malicious14.scan(Path(f"{file_path}/data/malicious14.pkl"))
     assert malicious14.issues.all_issues == expected_malicious14
 
+    expected_malicious15 = [
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "bdb",
+                "Bdb.run",
+                IssueSeverity.CRITICAL,
+                f"{file_path}/data/malicious15.pkl",
+            ),
+        ),
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "bdb",
+                "Bdb",
+                IssueSeverity.CRITICAL,
+                f"{file_path}/data/malicious15.pkl",
+            ),
+        ),
+    ]
+    malicious15 = ModelScan()
+    malicious15.scan(Path(f"{file_path}/data/malicious15.pkl"))
+    assert sorted(malicious15.issues.all_issues, key=str) == sorted(
+        expected_malicious15, key=str
+    )
+
 
 def test_scan_directory_path(file_path: str) -> None:
     expected = {
@@ -1265,6 +1304,26 @@ def test_scan_directory_path(file_path: str) -> None:
                 f"{file_path}/data/malicious14.pkl",
             ),
         ),
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "bdb",
+                "Bdb",
+                IssueSeverity.CRITICAL,
+                f"{file_path}/data/malicious15.pkl",
+            ),
+        ),
+        Issue(
+            IssueCode.UNSAFE_OPERATOR,
+            IssueSeverity.CRITICAL,
+            OperatorIssueDetails(
+                "bdb",
+                "Bdb.run",
+                IssueSeverity.CRITICAL,
+                f"{file_path}/data/malicious15.pkl",
+            ),
+        ),
     }
     ms = ModelScan()
     p = Path(f"{file_path}/data/")
@@ -1283,6 +1342,7 @@ def test_scan_directory_path(file_path: str) -> None:
         "malicious12.pkl",
         "malicious13.pkl",
         "malicious14.pkl",
+        "malicious15.pkl",
         "malicious1_v0.dill",
         "malicious1_v3.dill",
         "malicious1_v4.dill",