From 98aa1f890fcece12abe8b687b29a82b794141bbf Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 10:46:54 -0700 Subject: [PATCH 1/8] Port picklescan PR #29 --- modelscan/settings.py | 3 +++ tests/test_modelscan.py | 27 +++++++++++++++++++++++++++ 2 files changed, 30 insertions(+) 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..db3dca9 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): + 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,22 @@ 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", + IssueSeverity.CRITICAL, + f"{file_path}/data/malicious15.pkl", + ), + ) + ] + malicious15 = ModelScan() + malicious15.scan(Path(f"{file_path}/data/malicious15.pkl")) + assert malicious15.issues.all_issues == expected_malicious15 + def test_scan_directory_path(file_path: str) -> None: expected = { From 8ec5642cca942f27eb85e85eaa0c588332e50b2b Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 10:49:46 -0700 Subject: [PATCH 2/8] update test --- tests/test_modelscan.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index db3dca9..4572ea1 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1022,6 +1022,16 @@ def test_scan_pickle_operators(file_path: Any) -> None: 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", + ), ) ] malicious15 = ModelScan() From 390ac6be11dca8d270385eec78928572a08203e6 Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 10:50:55 -0700 Subject: [PATCH 3/8] Fix lint --- tests/test_modelscan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index 4572ea1..8d69fcc 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1032,7 +1032,7 @@ def test_scan_pickle_operators(file_path: Any) -> None: IssueSeverity.CRITICAL, f"{file_path}/data/malicious15.pkl", ), - ) + ), ] malicious15 = ModelScan() malicious15.scan(Path(f"{file_path}/data/malicious15.pkl")) From f3b3eeb2ead250d38a62c2f89d3c4b1ed34c5e1d Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 10:54:56 -0700 Subject: [PATCH 4/8] Fix test_scan_directory_path --- tests/test_modelscan.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index 8d69fcc..fb102f1 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1302,6 +1302,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/") @@ -1320,6 +1340,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", From 6766e2aa4cc522a9c2815e49a459d1e218d41f32 Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 10:58:37 -0700 Subject: [PATCH 5/8] Fix test_scan_pickle_operators --- tests/test_modelscan.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index fb102f1..711e92b 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1018,7 +1018,7 @@ def test_scan_pickle_operators(file_path: Any) -> None: IssueSeverity.CRITICAL, OperatorIssueDetails( "bdb", - "Bdb", + "Bdb.run", IssueSeverity.CRITICAL, f"{file_path}/data/malicious15.pkl", ), @@ -1028,7 +1028,7 @@ def test_scan_pickle_operators(file_path: Any) -> None: IssueSeverity.CRITICAL, OperatorIssueDetails( "bdb", - "Bdb.run", + "Bdb", IssueSeverity.CRITICAL, f"{file_path}/data/malicious15.pkl", ), From 2f9487b42741010c78ab034869c0adb374807553 Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 11:30:10 -0700 Subject: [PATCH 6/8] Try stabilize tests --- tests/test_modelscan.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index 711e92b..d32b1d1 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -90,7 +90,7 @@ def __reduce__(self) -> Any: class Malicious15: - def __reduce__(self): + def __reduce__(self) -> Any: bd = bdb.Bdb() return bdb.Bdb.run, ( bd, @@ -1036,7 +1036,7 @@ def test_scan_pickle_operators(file_path: Any) -> None: ] malicious15 = ModelScan() malicious15.scan(Path(f"{file_path}/data/malicious15.pkl")) - assert malicious15.issues.all_issues == expected_malicious15 + assert sorted(malicious15.issues.all_issues) == sorted(expected_malicious15) def test_scan_directory_path(file_path: str) -> None: From 06b563df816a7588a76202f06e86a01f1a304283 Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 11:36:15 -0700 Subject: [PATCH 7/8] update key --- tests/test_modelscan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index d32b1d1..dd06b0f 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1036,7 +1036,7 @@ def test_scan_pickle_operators(file_path: Any) -> None: ] malicious15 = ModelScan() malicious15.scan(Path(f"{file_path}/data/malicious15.pkl")) - assert sorted(malicious15.issues.all_issues) == sorted(expected_malicious15) + assert sorted(malicious15.issues.all_issues, key=str) == sorted(expected_malicious15, key=str) def test_scan_directory_path(file_path: str) -> None: From 70da7a7b240b272afd4c871fb1ba181273c12d46 Mon Sep 17 00:00:00 2001 From: Matthieu Maitre Date: Sat, 14 Sep 2024 11:38:31 -0700 Subject: [PATCH 8/8] black --- tests/test_modelscan.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index dd06b0f..464d26c 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1036,7 +1036,9 @@ def test_scan_pickle_operators(file_path: Any) -> None: ] malicious15 = ModelScan() malicious15.scan(Path(f"{file_path}/data/malicious15.pkl")) - assert sorted(malicious15.issues.all_issues, key=str) == sorted(expected_malicious15, key=str) + assert sorted(malicious15.issues.all_issues, key=str) == sorted( + expected_malicious15, key=str + ) def test_scan_directory_path(file_path: str) -> None: