diff --git a/modelscan/settings.py b/modelscan/settings.py index 395dfbe..f202a8f 100644 --- 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 index 0953ad3..464d26c 100644 --- 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",