From 27d33d2d129dbe4808bd2b8e9371384cb9ebebfa Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Mon, 25 Sep 2023 18:28:40 -0700 Subject: [PATCH 1/5] Bitwarden secret test --- .github/workflows/test.yml | 6 + tests/test_modelscan.py | 1456 ++++++++++++++++++------------------ tests/test_secrets | 4 + 3 files changed, 738 insertions(+), 728 deletions(-) create mode 100644 tests/test_secrets diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index c8fc5356..2407aaa6 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -34,6 +34,12 @@ jobs: if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' run: | make install-test + - name: Get Secrets + uses: bitwarden/sm-action@v1 + with: + access_token: ${{ secrets.BW_ACCESS_TOKEN }} + secrets: | + 6b0baeba-4bd1-4c7d-b0c4-b0850005549d > BW_SECRET_1 - name: Run Tests run: | make test diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index d6395226..ad2226b8 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -1,737 +1,737 @@ -import aiohttp -import http.client -import importlib -import io -import numpy as np -import os -from pathlib import Path -import pickle -import pytest -import requests # type: ignore[import] -import socket -import subprocess -import sys -import tensorflow as tf -from tensorflow import keras -from typing import Any, List, Set -from test_utils import generate_dill_unsafe_file -import zipfile +# import aiohttp +# import http.client +# import importlib +# import io +# import numpy as np +# import os +# from pathlib import Path +# import pickle +# import pytest +# import requests # type: ignore[import] +# import socket +# import subprocess +# import sys +# import tensorflow as tf +# from tensorflow import keras +# from typing import Any, List, Set +# from test_utils import generate_dill_unsafe_file +# import zipfile -from modelscan.modelscan import Modelscan -from modelscan.cli import cli -from modelscan.issues import ( - Issue, - IssueCode, - IssueSeverity, - OperatorIssueDetails, -) -from modelscan.tools.picklescanner import ( - scan_pickle_bytes, - scan_numpy, -) +# from modelscan.modelscan import Modelscan +# from modelscan.cli import cli +# from modelscan.issues import ( +# Issue, +# IssueCode, +# IssueSeverity, +# OperatorIssueDetails, +# ) +# from modelscan.tools.picklescanner import ( +# scan_pickle_bytes, +# scan_numpy, +# ) -class Malicious1: - def __reduce__(self) -> Any: - return eval, ("print('456')",) +# class Malicious1: +# def __reduce__(self) -> Any: +# return eval, ("print('456')",) -class Malicious2: - def __reduce__(self) -> Any: - return os.system, ("ls -la",) - - -class Malicious3: - def __reduce__(self) -> Any: - return http.client.HTTPSConnection, ("github.com",) - - -malicious3_pickle_bytes = pickle.dumps( - Malicious3(), protocol=0 -) # Malicious3 needs to be pickled before HTTPSConnection is mocked below - - -class Malicious4: - def __reduce__(self) -> Any: - return requests.get, ("https://github.com",) - - -class Malicious5: - def __reduce__(self) -> Any: - return aiohttp.ClientSession, tuple() - - -class Malicious6: - def __reduce__(self) -> Any: - return socket.create_connection, (("github.com", 80),) - - -class Malicious7: - def __reduce__(self) -> Any: - return subprocess.run, (["ls", "-l"],) - - -class Malicious8: - def __reduce__(self) -> Any: - return sys.exit, (0,) - - -def initialize_pickle_file(path: str, obj: Any, version: int) -> None: - if not os.path.exists(path): - with open(path, "wb") as file: - pickle.dump(obj, file, protocol=version) - - -def initialize_data_file(path: str, data: Any) -> None: - if not os.path.exists(path): - with open(path, "wb") as file: - file.write(data) - - -def initialize_zip_file(path: str, file_name: str, data: Any) -> None: - if not os.path.exists(path): - with zipfile.ZipFile(path, "w") as zip: - zip.writestr(file_name, data) - - -def initialize_numpy_file(path: str) -> None: - import numpy as np - - # create numpy object array - with open(path, "wb") as f: - data = [(1, 2), (3, 4)] - x = np.empty((2, 2), dtype=object) - x[:] = data - np.save(f, x) - - -@pytest.fixture(scope="session") -def zip_file_path(tmp_path_factory: Any) -> Any: - tmp = tmp_path_factory.mktemp("zip") - initialize_zip_file( - f"{tmp}/test.zip", - "data.pkl", - pickle.dumps(Malicious1(), protocol=4), - ) - return tmp +# class Malicious2: +# def __reduce__(self) -> Any: +# return os.system, ("ls -la",) + + +# class Malicious3: +# def __reduce__(self) -> Any: +# return http.client.HTTPSConnection, ("github.com",) + + +# malicious3_pickle_bytes = pickle.dumps( +# Malicious3(), protocol=0 +# ) # Malicious3 needs to be pickled before HTTPSConnection is mocked below + + +# class Malicious4: +# def __reduce__(self) -> Any: +# return requests.get, ("https://github.com",) + + +# class Malicious5: +# def __reduce__(self) -> Any: +# return aiohttp.ClientSession, tuple() + + +# class Malicious6: +# def __reduce__(self) -> Any: +# return socket.create_connection, (("github.com", 80),) + + +# class Malicious7: +# def __reduce__(self) -> Any: +# return subprocess.run, (["ls", "-l"],) + + +# class Malicious8: +# def __reduce__(self) -> Any: +# return sys.exit, (0,) + + +# def initialize_pickle_file(path: str, obj: Any, version: int) -> None: +# if not os.path.exists(path): +# with open(path, "wb") as file: +# pickle.dump(obj, file, protocol=version) + + +# def initialize_data_file(path: str, data: Any) -> None: +# if not os.path.exists(path): +# with open(path, "wb") as file: +# file.write(data) + + +# def initialize_zip_file(path: str, file_name: str, data: Any) -> None: +# if not os.path.exists(path): +# with zipfile.ZipFile(path, "w") as zip: +# zip.writestr(file_name, data) + + +# def initialize_numpy_file(path: str) -> None: +# import numpy as np + +# # create numpy object array +# with open(path, "wb") as f: +# data = [(1, 2), (3, 4)] +# x = np.empty((2, 2), dtype=object) +# x[:] = data +# np.save(f, x) + + +# @pytest.fixture(scope="session") +# def zip_file_path(tmp_path_factory: Any) -> Any: +# tmp = tmp_path_factory.mktemp("zip") +# initialize_zip_file( +# f"{tmp}/test.zip", +# "data.pkl", +# pickle.dumps(Malicious1(), protocol=4), +# ) +# return tmp -@pytest.fixture(scope="session") -def pickle_file_path(tmp_path_factory: Any) -> Any: - tmp = tmp_path_factory.mktemp("test_files") - os.makedirs(f"{tmp}/data", exist_ok=True) +# @pytest.fixture(scope="session") +# def pickle_file_path(tmp_path_factory: Any) -> Any: +# tmp = tmp_path_factory.mktemp("test_files") +# os.makedirs(f"{tmp}/data", exist_ok=True) - # Test with Pickle versions 0, 3, and 4: - # - Pickle versions 0, 1, 2 have built-in functions under '__builtin__' while versions 3 and 4 have them under 'builtins' - # - Pickle versions 0, 1, 2, 3 use 'GLOBAL' opcode while 4 uses 'STACK_GLOBAL' opcode - for version in (0, 3, 4): - initialize_pickle_file( - f"{tmp}/data/benign0_v{version}.pkl", ["a", "b", "c"], version - ) - initialize_pickle_file( - f"{tmp}/data/malicious1_v{version}.pkl", Malicious1(), version - ) - initialize_pickle_file( - f"{tmp}/data/malicious2_v{version}.pkl", Malicious2(), version - ) +# # Test with Pickle versions 0, 3, and 4: +# # - Pickle versions 0, 1, 2 have built-in functions under '__builtin__' while versions 3 and 4 have them under 'builtins' +# # - Pickle versions 0, 1, 2, 3 use 'GLOBAL' opcode while 4 uses 'STACK_GLOBAL' opcode +# for version in (0, 3, 4): +# initialize_pickle_file( +# f"{tmp}/data/benign0_v{version}.pkl", ["a", "b", "c"], version +# ) +# initialize_pickle_file( +# f"{tmp}/data/malicious1_v{version}.pkl", Malicious1(), version +# ) +# initialize_pickle_file( +# f"{tmp}/data/malicious2_v{version}.pkl", Malicious2(), version +# ) - # Malicious Pickle from https://sensepost.com/cms/resources/conferences/2011/sour_pickles/BH_US_11_Slaviero_Sour_Pickles.pdf - initialize_data_file( - f"{tmp}/data/malicious0.pkl", - b'c__builtin__\nglobals\n(tRp100\n0c__builtin__\ncompile\n(S\'fl=open("/etc/passwd");picklesmashed=fl.read();' - + b"'\nS''\nS'exec'\ntRp101\n0c__builtin__\neval\n(g101\ng100\ntRp102\n0c__builtin__\ngetattr\n(c__builtin__\n" - + b"dict\nS'get'\ntRp103\n0c__builtin__\napply\n(g103\n(g100\nS'picklesmashed'\nltRp104\n0g104\n.", - ) - - initialize_data_file(f"{tmp}/data/malicious3.pkl", malicious3_pickle_bytes) - initialize_pickle_file(f"{tmp}/data/malicious4.pickle", Malicious4(), 4) - initialize_pickle_file(f"{tmp}/data/malicious5.pickle", Malicious5(), 4) - initialize_data_file( - f"{tmp}/data/malicious6.pkl", - pickle.dumps(["a", "b", "c"]) + pickle.dumps(Malicious4()), - ) - 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_zip_file( - f"{tmp}/data/malicious1.zip", - "data.pkl", - pickle.dumps(Malicious1(), protocol=4), - ) - - initialize_numpy_file(f"{tmp}/data/object_array.npy") - - return tmp - - -@pytest.fixture(scope="session") -def keras_file_path(tmp_path_factory: Any) -> Any: - # Create a simple model. - inputs = keras.Input(shape=(32,)) - outputs = keras.layers.Dense(1)(inputs) - keras_model = keras.Model(inputs, outputs) - keras_model.compile(optimizer="adam", loss="mean_squared_error") - - # Train the model. - test_input = np.random.random((128, 32)) - test_target = np.random.random((128, 1)) - keras_model.fit(test_input, test_target) - - tmp = tmp_path_factory.mktemp("keras") - with open(f"{tmp}/safe", "wb") as fo: - pickle.dump(keras_model, fo) - keras_model.save(f"{tmp}/safe.h5") - - # Inject code with the command - command = "exec" - malicious_code = 'print("Malicious code!")' - - generate_dill_unsafe_file(keras_model, command, malicious_code, f"{tmp}/unsafe") - attack = ( - lambda x: exec( # type: ignore[func-returns-value] - """import http.client -import json -import os -conn = http.client.HTTPSConnection("protectai.com")""" - ) - or x - ) - input_to_new_layer = keras.layers.Dense(1)(keras_model.layers[-1].output) - new_layer = keras.layers.Lambda(attack)(input_to_new_layer) - - malicious_model = tf.keras.Model(inputs=keras_model.inputs, outputs=[new_layer]) - malicious_model.compile(optimizer="adam", loss="mean_squared_error") - - malicious_model.save(f"{tmp}/unsafe.h5") - - return tmp - - -def compare_results(resultList: List[Issue], expectedSet: Set[Issue]) -> None: - for result in resultList: - assert result in expectedSet - - -def test_scan_pickle_bytes() -> None: - expected = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails("builtins", "eval", "file.pkl"), - ) - ] - assert ( - scan_pickle_bytes(io.BytesIO(pickle.dumps(Malicious1())), "file.pkl")[0] - == expected - ) - - -def test_scan_zip(zip_file_path: Any) -> None: - expected = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{zip_file_path}/test.zip:data.pkl" - ), - ) - ] - - ms = Modelscan() - ms._scan_zip(f"{zip_file_path}/test.zip") - assert ms.issues.all_issues == expected - - -def test_scan_numpy(pickle_file_path: Any) -> None: - expected = { - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "numpy.core.multiarray", "_reconstruct", "object_array.npy" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails("numpy", "ndarray", "object_array.npy"), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails("numpy", "dtype", "object_array.npy"), - ), - } - with open(f"{pickle_file_path}/data/object_array.npy", "rb") as f: - compare_results( - scan_numpy(io.BytesIO(f.read()), "object_array.npy")[0], expected - ) - - -def test_scan_file_path(pickle_file_path: Any) -> None: - benign = Modelscan() - benign.scan_path(Path(f"{pickle_file_path}/data/benign0_v3.pkl")) - assert benign.issues.all_issues == [] - - malicious0 = Modelscan() - expected_malicious0 = { - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "__builtin__", "dict", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "apply", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "eval", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "compile", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "__builtin__", "globals", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "getattr", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - } - malicious0.scan_path(Path(f"{pickle_file_path}/data/malicious0.pkl")) - compare_results(malicious0.issues.all_issues, expected_malicious0) - - expected_malicious1_v0 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "eval", f"{pickle_file_path}/data/malicious1_v0.pkl" - ), - ) - ] - expected_malicious1_v3 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{pickle_file_path}/data/malicious1_v3.pkl" - ), - ) - ] - expected_malicious1_v4 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{pickle_file_path}/data/malicious1_v4.pkl" - ), - ) - ] - expected_malicious1 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{pickle_file_path}/data/malicious1.zip:data.pkl" - ), - ) - ] - malicious1_v0 = Modelscan() - malicious1_v3 = Modelscan() - malicious1_v4 = Modelscan() - malicious1 = Modelscan() - malicious1_v0.scan_path(Path(f"{pickle_file_path}/data/malicious1_v0.pkl")) - malicious1_v3.scan_path(Path(f"{pickle_file_path}/data/malicious1_v3.pkl")) - malicious1_v4.scan_path(Path(f"{pickle_file_path}/data/malicious1_v4.pkl")) - malicious1.scan_path(Path(f"{pickle_file_path}/data/malicious1.zip")) - assert malicious1_v0.issues.all_issues == expected_malicious1_v0 - assert malicious1_v3.issues.all_issues == expected_malicious1_v3 - assert malicious1_v4.issues.all_issues == expected_malicious1_v4 - assert malicious1.issues.all_issues == expected_malicious1 - - expected_malicious2_v0 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "posix", "system", f"{pickle_file_path}/data/malicious2_v0.pkl" - ), - ) - ] - expected_malicious2_v3 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "posix", "system", f"{pickle_file_path}/data/malicious2_v3.pkl" - ), - ) - ] - expected_malicious2_v4 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "posix", "system", f"{pickle_file_path}/data/malicious2_v4.pkl" - ), - ) - ] - malicious2_v0 = Modelscan() - malicious2_v3 = Modelscan() - malicious2_v4 = Modelscan() - malicious2_v0.scan_path(Path(f"{pickle_file_path}/data/malicious2_v0.pkl")) - malicious2_v3.scan_path(Path(f"{pickle_file_path}/data/malicious2_v3.pkl")) - malicious2_v4.scan_path(Path(f"{pickle_file_path}/data/malicious2_v4.pkl")) - assert malicious2_v0.issues.all_issues == expected_malicious2_v0 - assert malicious2_v3.issues.all_issues == expected_malicious2_v3 - assert malicious2_v4.issues.all_issues == expected_malicious2_v4 - - expected_malicious3 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "httplib", - "HTTPSConnection", - Path(f"{pickle_file_path}/data/malicious3.pkl"), - ), - ) - ] - malicious3 = Modelscan() - malicious3.scan_path(Path(f"{pickle_file_path}/data/malicious3.pkl")) - assert malicious3.issues.all_issues == expected_malicious3 - - expected_malicious4 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "requests.api", "get", f"{pickle_file_path}/data/malicious4.pickle" - ), - ) - ] - malicious4 = Modelscan() - malicious4.scan_path(Path(f"{pickle_file_path}/data/malicious4.pickle")) - assert malicious4.issues.all_issues == expected_malicious4 - - expected_malicious5 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "aiohttp.client", - "ClientSession", - f"{pickle_file_path}/data/malicious5.pickle", - ), - ) - ] - malicious5 = Modelscan() - malicious5.scan_path(Path(f"{pickle_file_path}/data/malicious5.pickle")) - assert malicious5.issues.all_issues == expected_malicious5 - - expected_malicious6 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "requests.api", "get", f"{pickle_file_path}/data/malicious6.pkl" - ), - ) - ] - malicious6 = Modelscan() - malicious6.scan_path(Path(f"{pickle_file_path}/data/malicious6.pkl")) - assert malicious6.issues.all_issues == expected_malicious6 - - expected_malicious7 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "socket", "create_connection", f"{pickle_file_path}/data/malicious7.pkl" - ), - ) - ] - malicious7 = Modelscan() - malicious7.scan_path(Path(f"{pickle_file_path}/data/malicious7.pkl")) - assert malicious7.issues.all_issues == expected_malicious7 - - expected_malicious8 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "subprocess", "run", f"{pickle_file_path}/data/malicious8.pkl" - ), - ) - ] - malicious8 = Modelscan() - malicious8.scan_path(Path(f"{pickle_file_path}/data/malicious8.pkl")) - assert malicious8.issues.all_issues == expected_malicious8 - - expected_malicious9 = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "sys", "exit", f"{pickle_file_path}/data/malicious9.pkl" - ), - ) - ] - malicious9 = Modelscan() - malicious9.scan_path(Path(f"{pickle_file_path}/data/malicious9.pkl")) - assert malicious9.issues.all_issues == expected_malicious9 - - -def test_scan_directory_path(pickle_file_path: str) -> None: - expected = { - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{pickle_file_path}/data/malicious1.zip:data.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "subprocess", "run", f"{pickle_file_path}/data/malicious8.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "sys", "exit", f"{pickle_file_path}/data/malicious9.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "requests.api", "get", f"{pickle_file_path}/data/malicious4.pickle" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{pickle_file_path}/data/malicious1_v3.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "eval", f"{pickle_file_path}/data/malicious1_v0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "builtins", "eval", f"{pickle_file_path}/data/malicious1_v4.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "numpy", "ndarray", f"{pickle_file_path}/data/object_array.npy" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "numpy", "dtype", f"{pickle_file_path}/data/object_array.npy" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "numpy", "dtype", f"{pickle_file_path}/data/object_array.npy" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "numpy.core.multiarray", - "_reconstruct", - f"{pickle_file_path}/data/object_array.npy", - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "aiohttp.client", - "ClientSession", - f"{pickle_file_path}/data/malicious5.pickle", - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "posix", "system", f"{pickle_file_path}/data/malicious2_v4.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "socket", "create_connection", f"{pickle_file_path}/data/malicious7.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "requests.api", "get", f"{pickle_file_path}/data/malicious6.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "compile", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "eval", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "__builtin__", "globals", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "apply", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", "getattr", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "__builtin__", "dict", f"{pickle_file_path}/data/malicious0.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "posix", "system", f"{pickle_file_path}/data/malicious2_v3.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.HIGH, - OperatorIssueDetails( - "httplib", "HTTPSConnection", f"{pickle_file_path}/data/malicious3.pkl" - ), - ), - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "posix", "system", f"{pickle_file_path}/data/malicious2_v0.pkl" - ), - ), - } - ms = Modelscan() - p = Path(f"{pickle_file_path}/data/") - ms.scan_path(p) - compare_results(ms.issues.all_issues, expected) - - -def test_scan_huggingface_model() -> None: - expected = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.CRITICAL, - OperatorIssueDetails( - "__builtin__", - "eval", - "https://huggingface.co/ykilcher/totally-harmless-model/resolve/main/pytorch_model.bin:archive/data.pkl", - ), - ) - ] - ms = Modelscan() - ms.scan_huggingface_model("ykilcher/totally-harmless-model") - assert ms.issues.all_issues == expected - - -# def test_scan_tf() -> None: - - -def test_scan_keras(keras_file_path: Any) -> None: - ms = Modelscan() - ms.scan_path(Path(f"{keras_file_path}/safe.h5")) - assert ms.issues.all_issues == [] - - expected = [ - Issue( - IssueCode.UNSAFE_OPERATOR, - IssueSeverity.MEDIUM, - OperatorIssueDetails( - "Keras", - "Lambda", - f"{keras_file_path}/unsafe.h5", - ), - ) - ] - ms.scan_path(Path(f"{keras_file_path}/unsafe.h5")) - assert ms.issues.all_issues == expected - - -def test_main(pickle_file_path: Any) -> None: - argv = sys.argv - try: - sys.argv = ["modelscan", "-p", f"{pickle_file_path}/data/benign0_v3.pkl"] - assert cli() == 0 - importlib.import_module("modelscan.scanner") - except SystemExit: - pass - finally: - sys.argv = argv +# # Malicious Pickle from https://sensepost.com/cms/resources/conferences/2011/sour_pickles/BH_US_11_Slaviero_Sour_Pickles.pdf +# initialize_data_file( +# f"{tmp}/data/malicious0.pkl", +# b'c__builtin__\nglobals\n(tRp100\n0c__builtin__\ncompile\n(S\'fl=open("/etc/passwd");picklesmashed=fl.read();' +# + b"'\nS''\nS'exec'\ntRp101\n0c__builtin__\neval\n(g101\ng100\ntRp102\n0c__builtin__\ngetattr\n(c__builtin__\n" +# + b"dict\nS'get'\ntRp103\n0c__builtin__\napply\n(g103\n(g100\nS'picklesmashed'\nltRp104\n0g104\n.", +# ) + +# initialize_data_file(f"{tmp}/data/malicious3.pkl", malicious3_pickle_bytes) +# initialize_pickle_file(f"{tmp}/data/malicious4.pickle", Malicious4(), 4) +# initialize_pickle_file(f"{tmp}/data/malicious5.pickle", Malicious5(), 4) +# initialize_data_file( +# f"{tmp}/data/malicious6.pkl", +# pickle.dumps(["a", "b", "c"]) + pickle.dumps(Malicious4()), +# ) +# 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_zip_file( +# f"{tmp}/data/malicious1.zip", +# "data.pkl", +# pickle.dumps(Malicious1(), protocol=4), +# ) + +# initialize_numpy_file(f"{tmp}/data/object_array.npy") + +# return tmp + + +# @pytest.fixture(scope="session") +# def keras_file_path(tmp_path_factory: Any) -> Any: +# # Create a simple model. +# inputs = keras.Input(shape=(32,)) +# outputs = keras.layers.Dense(1)(inputs) +# keras_model = keras.Model(inputs, outputs) +# keras_model.compile(optimizer="adam", loss="mean_squared_error") + +# # Train the model. +# test_input = np.random.random((128, 32)) +# test_target = np.random.random((128, 1)) +# keras_model.fit(test_input, test_target) + +# tmp = tmp_path_factory.mktemp("keras") +# with open(f"{tmp}/safe", "wb") as fo: +# pickle.dump(keras_model, fo) +# keras_model.save(f"{tmp}/safe.h5") + +# # Inject code with the command +# command = "exec" +# malicious_code = 'print("Malicious code!")' + +# generate_dill_unsafe_file(keras_model, command, malicious_code, f"{tmp}/unsafe") +# attack = ( +# lambda x: exec( # type: ignore[func-returns-value] +# """import http.client +# import json +# import os +# conn = http.client.HTTPSConnection("protectai.com")""" +# ) +# or x +# ) +# input_to_new_layer = keras.layers.Dense(1)(keras_model.layers[-1].output) +# new_layer = keras.layers.Lambda(attack)(input_to_new_layer) + +# malicious_model = tf.keras.Model(inputs=keras_model.inputs, outputs=[new_layer]) +# malicious_model.compile(optimizer="adam", loss="mean_squared_error") + +# malicious_model.save(f"{tmp}/unsafe.h5") + +# return tmp + + +# def compare_results(resultList: List[Issue], expectedSet: Set[Issue]) -> None: +# for result in resultList: +# assert result in expectedSet + + +# def test_scan_pickle_bytes() -> None: +# expected = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails("builtins", "eval", "file.pkl"), +# ) +# ] +# assert ( +# scan_pickle_bytes(io.BytesIO(pickle.dumps(Malicious1())), "file.pkl")[0] +# == expected +# ) + + +# def test_scan_zip(zip_file_path: Any) -> None: +# expected = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{zip_file_path}/test.zip:data.pkl" +# ), +# ) +# ] + +# ms = Modelscan() +# ms._scan_zip(f"{zip_file_path}/test.zip") +# assert ms.issues.all_issues == expected + + +# def test_scan_numpy(pickle_file_path: Any) -> None: +# expected = { +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "numpy.core.multiarray", "_reconstruct", "object_array.npy" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails("numpy", "ndarray", "object_array.npy"), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails("numpy", "dtype", "object_array.npy"), +# ), +# } +# with open(f"{pickle_file_path}/data/object_array.npy", "rb") as f: +# compare_results( +# scan_numpy(io.BytesIO(f.read()), "object_array.npy")[0], expected +# ) + + +# def test_scan_file_path(pickle_file_path: Any) -> None: +# benign = Modelscan() +# benign.scan_path(Path(f"{pickle_file_path}/data/benign0_v3.pkl")) +# assert benign.issues.all_issues == [] + +# malicious0 = Modelscan() +# expected_malicious0 = { +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "__builtin__", "dict", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "apply", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "eval", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "compile", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "__builtin__", "globals", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "getattr", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# } +# malicious0.scan_path(Path(f"{pickle_file_path}/data/malicious0.pkl")) +# compare_results(malicious0.issues.all_issues, expected_malicious0) + +# expected_malicious1_v0 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "eval", f"{pickle_file_path}/data/malicious1_v0.pkl" +# ), +# ) +# ] +# expected_malicious1_v3 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{pickle_file_path}/data/malicious1_v3.pkl" +# ), +# ) +# ] +# expected_malicious1_v4 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{pickle_file_path}/data/malicious1_v4.pkl" +# ), +# ) +# ] +# expected_malicious1 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{pickle_file_path}/data/malicious1.zip:data.pkl" +# ), +# ) +# ] +# malicious1_v0 = Modelscan() +# malicious1_v3 = Modelscan() +# malicious1_v4 = Modelscan() +# malicious1 = Modelscan() +# malicious1_v0.scan_path(Path(f"{pickle_file_path}/data/malicious1_v0.pkl")) +# malicious1_v3.scan_path(Path(f"{pickle_file_path}/data/malicious1_v3.pkl")) +# malicious1_v4.scan_path(Path(f"{pickle_file_path}/data/malicious1_v4.pkl")) +# malicious1.scan_path(Path(f"{pickle_file_path}/data/malicious1.zip")) +# assert malicious1_v0.issues.all_issues == expected_malicious1_v0 +# assert malicious1_v3.issues.all_issues == expected_malicious1_v3 +# assert malicious1_v4.issues.all_issues == expected_malicious1_v4 +# assert malicious1.issues.all_issues == expected_malicious1 + +# expected_malicious2_v0 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "posix", "system", f"{pickle_file_path}/data/malicious2_v0.pkl" +# ), +# ) +# ] +# expected_malicious2_v3 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "posix", "system", f"{pickle_file_path}/data/malicious2_v3.pkl" +# ), +# ) +# ] +# expected_malicious2_v4 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "posix", "system", f"{pickle_file_path}/data/malicious2_v4.pkl" +# ), +# ) +# ] +# malicious2_v0 = Modelscan() +# malicious2_v3 = Modelscan() +# malicious2_v4 = Modelscan() +# malicious2_v0.scan_path(Path(f"{pickle_file_path}/data/malicious2_v0.pkl")) +# malicious2_v3.scan_path(Path(f"{pickle_file_path}/data/malicious2_v3.pkl")) +# malicious2_v4.scan_path(Path(f"{pickle_file_path}/data/malicious2_v4.pkl")) +# assert malicious2_v0.issues.all_issues == expected_malicious2_v0 +# assert malicious2_v3.issues.all_issues == expected_malicious2_v3 +# assert malicious2_v4.issues.all_issues == expected_malicious2_v4 + +# expected_malicious3 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "httplib", +# "HTTPSConnection", +# Path(f"{pickle_file_path}/data/malicious3.pkl"), +# ), +# ) +# ] +# malicious3 = Modelscan() +# malicious3.scan_path(Path(f"{pickle_file_path}/data/malicious3.pkl")) +# assert malicious3.issues.all_issues == expected_malicious3 + +# expected_malicious4 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "requests.api", "get", f"{pickle_file_path}/data/malicious4.pickle" +# ), +# ) +# ] +# malicious4 = Modelscan() +# malicious4.scan_path(Path(f"{pickle_file_path}/data/malicious4.pickle")) +# assert malicious4.issues.all_issues == expected_malicious4 + +# expected_malicious5 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "aiohttp.client", +# "ClientSession", +# f"{pickle_file_path}/data/malicious5.pickle", +# ), +# ) +# ] +# malicious5 = Modelscan() +# malicious5.scan_path(Path(f"{pickle_file_path}/data/malicious5.pickle")) +# assert malicious5.issues.all_issues == expected_malicious5 + +# expected_malicious6 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "requests.api", "get", f"{pickle_file_path}/data/malicious6.pkl" +# ), +# ) +# ] +# malicious6 = Modelscan() +# malicious6.scan_path(Path(f"{pickle_file_path}/data/malicious6.pkl")) +# assert malicious6.issues.all_issues == expected_malicious6 + +# expected_malicious7 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "socket", "create_connection", f"{pickle_file_path}/data/malicious7.pkl" +# ), +# ) +# ] +# malicious7 = Modelscan() +# malicious7.scan_path(Path(f"{pickle_file_path}/data/malicious7.pkl")) +# assert malicious7.issues.all_issues == expected_malicious7 + +# expected_malicious8 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "subprocess", "run", f"{pickle_file_path}/data/malicious8.pkl" +# ), +# ) +# ] +# malicious8 = Modelscan() +# malicious8.scan_path(Path(f"{pickle_file_path}/data/malicious8.pkl")) +# assert malicious8.issues.all_issues == expected_malicious8 + +# expected_malicious9 = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "sys", "exit", f"{pickle_file_path}/data/malicious9.pkl" +# ), +# ) +# ] +# malicious9 = Modelscan() +# malicious9.scan_path(Path(f"{pickle_file_path}/data/malicious9.pkl")) +# assert malicious9.issues.all_issues == expected_malicious9 + + +# def test_scan_directory_path(pickle_file_path: str) -> None: +# expected = { +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{pickle_file_path}/data/malicious1.zip:data.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "subprocess", "run", f"{pickle_file_path}/data/malicious8.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "sys", "exit", f"{pickle_file_path}/data/malicious9.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "requests.api", "get", f"{pickle_file_path}/data/malicious4.pickle" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{pickle_file_path}/data/malicious1_v3.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "eval", f"{pickle_file_path}/data/malicious1_v0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "builtins", "eval", f"{pickle_file_path}/data/malicious1_v4.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "numpy", "ndarray", f"{pickle_file_path}/data/object_array.npy" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "numpy", "dtype", f"{pickle_file_path}/data/object_array.npy" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "numpy", "dtype", f"{pickle_file_path}/data/object_array.npy" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "numpy.core.multiarray", +# "_reconstruct", +# f"{pickle_file_path}/data/object_array.npy", +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "aiohttp.client", +# "ClientSession", +# f"{pickle_file_path}/data/malicious5.pickle", +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "posix", "system", f"{pickle_file_path}/data/malicious2_v4.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "socket", "create_connection", f"{pickle_file_path}/data/malicious7.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "requests.api", "get", f"{pickle_file_path}/data/malicious6.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "compile", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "eval", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "__builtin__", "globals", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "apply", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", "getattr", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "__builtin__", "dict", f"{pickle_file_path}/data/malicious0.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "posix", "system", f"{pickle_file_path}/data/malicious2_v3.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.HIGH, +# OperatorIssueDetails( +# "httplib", "HTTPSConnection", f"{pickle_file_path}/data/malicious3.pkl" +# ), +# ), +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "posix", "system", f"{pickle_file_path}/data/malicious2_v0.pkl" +# ), +# ), +# } +# ms = Modelscan() +# p = Path(f"{pickle_file_path}/data/") +# ms.scan_path(p) +# compare_results(ms.issues.all_issues, expected) + + +# def test_scan_huggingface_model() -> None: +# expected = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.CRITICAL, +# OperatorIssueDetails( +# "__builtin__", +# "eval", +# "https://huggingface.co/ykilcher/totally-harmless-model/resolve/main/pytorch_model.bin:archive/data.pkl", +# ), +# ) +# ] +# ms = Modelscan() +# ms.scan_huggingface_model("ykilcher/totally-harmless-model") +# assert ms.issues.all_issues == expected + + +# # def test_scan_tf() -> None: + + +# def test_scan_keras(keras_file_path: Any) -> None: +# ms = Modelscan() +# ms.scan_path(Path(f"{keras_file_path}/safe.h5")) +# assert ms.issues.all_issues == [] + +# expected = [ +# Issue( +# IssueCode.UNSAFE_OPERATOR, +# IssueSeverity.MEDIUM, +# OperatorIssueDetails( +# "Keras", +# "Lambda", +# f"{keras_file_path}/unsafe.h5", +# ), +# ) +# ] +# ms.scan_path(Path(f"{keras_file_path}/unsafe.h5")) +# assert ms.issues.all_issues == expected + + +# def test_main(pickle_file_path: Any) -> None: +# argv = sys.argv +# try: +# sys.argv = ["modelscan", "-p", f"{pickle_file_path}/data/benign0_v3.pkl"] +# assert cli() == 0 +# importlib.import_module("modelscan.scanner") +# except SystemExit: +# pass +# finally: +# sys.argv = argv diff --git a/tests/test_secrets b/tests/test_secrets new file mode 100644 index 00000000..8a6bee09 --- /dev/null +++ b/tests/test_secrets @@ -0,0 +1,4 @@ +import os + +def test_scan_pickle_bytes() -> None: + assert os.environ.get("BW_SECRET_1") == "Secret-Value-123" \ No newline at end of file From a89ddcd752ea5f82819169d927db9e75eac03240 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Mon, 25 Sep 2023 18:36:52 -0700 Subject: [PATCH 2/5] test fix --- tests/test_secrets.py | 4 + tests/test_utils.py | 414 +++++++++++++++++++++--------------------- 2 files changed, 211 insertions(+), 207 deletions(-) create mode 100644 tests/test_secrets.py diff --git a/tests/test_secrets.py b/tests/test_secrets.py new file mode 100644 index 00000000..8a6bee09 --- /dev/null +++ b/tests/test_secrets.py @@ -0,0 +1,4 @@ +import os + +def test_scan_pickle_bytes() -> None: + assert os.environ.get("BW_SECRET_1") == "Secret-Value-123" \ No newline at end of file diff --git a/tests/test_utils.py b/tests/test_utils.py index 34b3eeb5..691520a5 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,236 +1,236 @@ -import dill -import os -import pickle -import struct -from typing import Any, Tuple -import os +# import dill +# import os +# import pickle +# import struct +# from typing import Any, Tuple +# import os - -class PickleInject: - """Pickle injection""" - - def __init__(self, inj_objs: Any, first: bool = True): - self.__name__ = "pickle_inject" - self.inj_objs = inj_objs - self.first = first - - class _Pickler(pickle._Pickler): - """Reimplementation of Pickler with support for injection""" - - def __init__( - self, file: Any, protocol: Any, inj_objs: Any, first: bool = True - ) -> None: - """ - file: File object with write attribute - protocol: Pickle protocol - Currently the default protocol is 4: https://docs.python.org/3/library/pickle.html - inj_objs: _joblibInject object that has both the command, and the code to be injected - first: Boolean object to determine if inj_objs should be serialized before the safe file or after the safe file. - """ - super().__init__(file, protocol) - self.inj_objs = inj_objs - self.first = first - - def dump(self, obj: Any) -> None: - """Pickle data, inject object before or after""" - if self.proto >= 2: # type: ignore[attr-defined] - self.write(pickle.PROTO + struct.pack("= 4: # type: ignore[attr-defined] - self.framer.start_framing() # type: ignore[attr-defined] - - # Inject the object(s) before the user-supplied data? - if self.first: - # Pickle injected objects - for inj_obj in self.inj_objs: - self.save(inj_obj) # type: ignore[attr-defined] - - # Pickle user-supplied data - self.save(obj) # type: ignore[attr-defined] - - # Inject the object(s) after the user-supplied data? - if not self.first: - # Pickle injected objects - for inj_obj in self.inj_objs: - self.save(inj_obj) # type: ignore[attr-defined] - - self.write(pickle.STOP) # type: ignore[attr-defined] - self.framer.end_framing() # type: ignore[attr-defined] - - def Pickler(self, file: Any, protocol: Any) -> _Pickler: - # Initialise the pickler interface with the injected object - return self._Pickler(file, protocol, self.inj_objs) - - class _PickleInject: - """Base class for pickling injected commands""" - - def __init__(self, args: Any, command: Any = None) -> None: - self.command = command - self.args = args - - def __reduce__(self) -> Tuple[Any, Any]: - """ - In general, the __reduce__ function is used by pickle to serialize objects. - If defined for an object, pickle would override its default __reduce__ function and serialize the object as outlined by the custom specified __reduce__ function, - The object returned by __reduce__ here is a callable: (self.command), and the tuple: with first element (self.args) is the code to be executed by self.command. - """ - return self.command, (self.args,) - - class System(_PickleInject): - """Create os.system command""" - - def __init__(self, args: Any) -> None: - super().__init__(args, command=os.system) - - class Exec(_PickleInject): - """Create exec command""" - - def __init__(self, args: Any) -> None: - super().__init__(args, command=exec) - - class Eval(_PickleInject): - """Create eval command""" - - def __init__(self, args: Any) -> None: - super().__init__(args, command=eval) - - class RunPy(_PickleInject): - """Create runpy command""" - - def __init__(self, args: Any) -> None: - import runpy - - super().__init__(args, command=runpy._run_code) # type: ignore[attr-defined] - - def __reduce__(self) -> Tuple[Any, Any]: - return self.command, (self.args, {}) + +# class PickleInject: +# """Pickle injection""" + +# def __init__(self, inj_objs: Any, first: bool = True): +# self.__name__ = "pickle_inject" +# self.inj_objs = inj_objs +# self.first = first + +# class _Pickler(pickle._Pickler): +# """Reimplementation of Pickler with support for injection""" + +# def __init__( +# self, file: Any, protocol: Any, inj_objs: Any, first: bool = True +# ) -> None: +# """ +# file: File object with write attribute +# protocol: Pickle protocol - Currently the default protocol is 4: https://docs.python.org/3/library/pickle.html +# inj_objs: _joblibInject object that has both the command, and the code to be injected +# first: Boolean object to determine if inj_objs should be serialized before the safe file or after the safe file. +# """ +# super().__init__(file, protocol) +# self.inj_objs = inj_objs +# self.first = first + +# def dump(self, obj: Any) -> None: +# """Pickle data, inject object before or after""" +# if self.proto >= 2: # type: ignore[attr-defined] +# self.write(pickle.PROTO + struct.pack("= 4: # type: ignore[attr-defined] +# self.framer.start_framing() # type: ignore[attr-defined] + +# # Inject the object(s) before the user-supplied data? +# if self.first: +# # Pickle injected objects +# for inj_obj in self.inj_objs: +# self.save(inj_obj) # type: ignore[attr-defined] + +# # Pickle user-supplied data +# self.save(obj) # type: ignore[attr-defined] + +# # Inject the object(s) after the user-supplied data? +# if not self.first: +# # Pickle injected objects +# for inj_obj in self.inj_objs: +# self.save(inj_obj) # type: ignore[attr-defined] + +# self.write(pickle.STOP) # type: ignore[attr-defined] +# self.framer.end_framing() # type: ignore[attr-defined] + +# def Pickler(self, file: Any, protocol: Any) -> _Pickler: +# # Initialise the pickler interface with the injected object +# return self._Pickler(file, protocol, self.inj_objs) + +# class _PickleInject: +# """Base class for pickling injected commands""" + +# def __init__(self, args: Any, command: Any = None) -> None: +# self.command = command +# self.args = args + +# def __reduce__(self) -> Tuple[Any, Any]: +# """ +# In general, the __reduce__ function is used by pickle to serialize objects. +# If defined for an object, pickle would override its default __reduce__ function and serialize the object as outlined by the custom specified __reduce__ function, +# The object returned by __reduce__ here is a callable: (self.command), and the tuple: with first element (self.args) is the code to be executed by self.command. +# """ +# return self.command, (self.args,) + +# class System(_PickleInject): +# """Create os.system command""" + +# def __init__(self, args: Any) -> None: +# super().__init__(args, command=os.system) + +# class Exec(_PickleInject): +# """Create exec command""" + +# def __init__(self, args: Any) -> None: +# super().__init__(args, command=exec) + +# class Eval(_PickleInject): +# """Create eval command""" + +# def __init__(self, args: Any) -> None: +# super().__init__(args, command=eval) + +# class RunPy(_PickleInject): +# """Create runpy command""" + +# def __init__(self, args: Any) -> None: +# import runpy + +# super().__init__(args, command=runpy._run_code) # type: ignore[attr-defined] + +# def __reduce__(self) -> Tuple[Any, Any]: +# return self.command, (self.args, {}) -def get_pickle_payload(command: str, malicious_code: str) -> Any: - if command == "system": - payload: Any = PickleInject.System(malicious_code) - elif command == "exec": - payload = PickleInject.Exec(malicious_code) - elif command == "eval": - payload = PickleInject.Eval(malicious_code) - elif command == "runpy": - payload = PickleInject.RunPy(malicious_code) - return payload +# def get_pickle_payload(command: str, malicious_code: str) -> Any: +# if command == "system": +# payload: Any = PickleInject.System(malicious_code) +# elif command == "exec": +# payload = PickleInject.Exec(malicious_code) +# elif command == "eval": +# payload = PickleInject.Eval(malicious_code) +# elif command == "runpy": +# payload = PickleInject.RunPy(malicious_code) +# return payload -def generate_unsafe_pickle_file( - safe_model: Any, command: str, malicious_code: str, unsafe_model_path: str -) -> None: - payload = get_pickle_payload(command, malicious_code) - pickle_protocol = 4 - file_for_unsafe_model = open(unsafe_model_path, "wb") - mypickler = PickleInject._Pickler(file_for_unsafe_model, pickle_protocol, [payload]) - mypickler.dump(safe_model) - file_for_unsafe_model.close() +# def generate_unsafe_pickle_file( +# safe_model: Any, command: str, malicious_code: str, unsafe_model_path: str +# ) -> None: +# payload = get_pickle_payload(command, malicious_code) +# pickle_protocol = 4 +# file_for_unsafe_model = open(unsafe_model_path, "wb") +# mypickler = PickleInject._Pickler(file_for_unsafe_model, pickle_protocol, [payload]) +# mypickler.dump(safe_model) +# file_for_unsafe_model.close() -class DillInject: - """Code injection using Dill Pickler""" +# class DillInject: +# """Code injection using Dill Pickler""" - def __init__(self, inj_objs: Any, first: bool = True): - self.__name__ = "dill_inject" - self.inj_objs = inj_objs - self.first = first +# def __init__(self, inj_objs: Any, first: bool = True): +# self.__name__ = "dill_inject" +# self.inj_objs = inj_objs +# self.first = first - class _Pickler(dill._dill.Pickler): # type: ignore[misc] - """Reimplementation of Pickler with support for injection""" +# class _Pickler(dill._dill.Pickler): # type: ignore[misc] +# """Reimplementation of Pickler with support for injection""" - def __init__(self, file: Any, protocol: Any, inj_objs: Any, first: bool = True): - super().__init__(file, protocol) - self.inj_objs = inj_objs - self.first = first +# def __init__(self, file: Any, protocol: Any, inj_objs: Any, first: bool = True): +# super().__init__(file, protocol) +# self.inj_objs = inj_objs +# self.first = first - def dump(self, obj: Any) -> None: - """Pickle data, inject object before or after""" - if self.proto >= 2: - self.write(pickle.PROTO + struct.pack("= 4: - self.framer.start_framing() +# def dump(self, obj: Any) -> None: +# """Pickle data, inject object before or after""" +# if self.proto >= 2: +# self.write(pickle.PROTO + struct.pack("= 4: +# self.framer.start_framing() - # Inject the object(s) before the user-supplied data? - if self.first: - # Pickle injected objects - for inj_obj in self.inj_objs: - self.save(inj_obj) +# # Inject the object(s) before the user-supplied data? +# if self.first: +# # Pickle injected objects +# for inj_obj in self.inj_objs: +# self.save(inj_obj) - # Pickle user-supplied data - self.save(obj) +# # Pickle user-supplied data +# self.save(obj) - # Inject the object(s) after the user-supplied data? - if not self.first: - # Pickle injected objects - for inj_obj in self.inj_objs: - self.save(inj_obj) +# # Inject the object(s) after the user-supplied data? +# if not self.first: +# # Pickle injected objects +# for inj_obj in self.inj_objs: +# self.save(inj_obj) - self.write(pickle.STOP) - self.framer.end_framing() +# self.write(pickle.STOP) +# self.framer.end_framing() - def DillPickler(self, file: Any, protocol: Any) -> _Pickler: - # Initialise the pickler interface with the injected object - return self._Pickler(file, protocol, self.inj_objs) +# def DillPickler(self, file: Any, protocol: Any) -> _Pickler: +# # Initialise the pickler interface with the injected object +# return self._Pickler(file, protocol, self.inj_objs) - class _DillInject: - """Base class for pickling injected commands""" +# class _DillInject: +# """Base class for pickling injected commands""" - def __init__(self, args: Any, command: Any = None): - self.command = command - self.args = args +# def __init__(self, args: Any, command: Any = None): +# self.command = command +# self.args = args - def __reduce__(self) -> Tuple[Any, Any]: - return self.command, (self.args,) +# def __reduce__(self) -> Tuple[Any, Any]: +# return self.command, (self.args,) - class System(_DillInject): - """Create os.system command""" +# class System(_DillInject): +# """Create os.system command""" - def __init__(self, args: Any): - super().__init__(args, command=os.system) +# def __init__(self, args: Any): +# super().__init__(args, command=os.system) - class Exec(_DillInject): - """Create exec command""" +# class Exec(_DillInject): +# """Create exec command""" - def __init__(self, args: Any): - super().__init__(args, command=exec) +# def __init__(self, args: Any): +# super().__init__(args, command=exec) - class Eval(_DillInject): - """Create eval command""" +# class Eval(_DillInject): +# """Create eval command""" - def __init__(self, args: Any): - super().__init__(args, command=eval) - - class RunPy(_DillInject): - """Create runpy command""" - - def __init__(self, args: Any): - import runpy - - super().__init__(args, command=runpy._run_code) # type: ignore[attr-defined] - - def __reduce__(self) -> Any: - return self.command, (self.args, {}) +# def __init__(self, args: Any): +# super().__init__(args, command=eval) + +# class RunPy(_DillInject): +# """Create runpy command""" + +# def __init__(self, args: Any): +# import runpy + +# super().__init__(args, command=runpy._run_code) # type: ignore[attr-defined] + +# def __reduce__(self) -> Any: +# return self.command, (self.args, {}) -def get_dill_payload(command: str, malicious_code: str) -> Any: - payload: Any - if command == "system": - payload = DillInject.System(malicious_code) - elif command == "exec": - payload = DillInject.Exec(malicious_code) - elif command == "eval": - payload = DillInject.Eval(malicious_code) - elif command == "runpy": - payload = DillInject.RunPy(malicious_code) - return payload +# def get_dill_payload(command: str, malicious_code: str) -> Any: +# payload: Any +# if command == "system": +# payload = DillInject.System(malicious_code) +# elif command == "exec": +# payload = DillInject.Exec(malicious_code) +# elif command == "eval": +# payload = DillInject.Eval(malicious_code) +# elif command == "runpy": +# payload = DillInject.RunPy(malicious_code) +# return payload -def generate_dill_unsafe_file( - safe_model: Any, command: str, malicious_code: str, unsafe_model_path: str -) -> None: - payload = get_dill_payload(command, malicious_code) - pickle_protocol = 4 - file_for_unsafe_model = open(unsafe_model_path, "wb") - mypickler = DillInject._Pickler(file_for_unsafe_model, pickle_protocol, [payload]) - mypickler.dump(safe_model) - file_for_unsafe_model.close() +# def generate_dill_unsafe_file( +# safe_model: Any, command: str, malicious_code: str, unsafe_model_path: str +# ) -> None: +# payload = get_dill_payload(command, malicious_code) +# pickle_protocol = 4 +# file_for_unsafe_model = open(unsafe_model_path, "wb") +# mypickler = DillInject._Pickler(file_for_unsafe_model, pickle_protocol, [payload]) +# mypickler.dump(safe_model) +# file_for_unsafe_model.close() From 5f69103c392846785d3f37323c7a6bfcec03b443 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Tue, 26 Sep 2023 16:19:08 -0700 Subject: [PATCH 3/5] Add OP --- .github/workflows/test.yml | 10 +++++++++- tests/test_secrets.py | 7 +++++-- 2 files changed, 14 insertions(+), 3 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2407aaa6..7b7c08e4 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -34,12 +34,20 @@ jobs: if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' run: | make install-test - - name: Get Secrets + - name: Get Bitwarden Secrets uses: bitwarden/sm-action@v1 with: access_token: ${{ secrets.BW_ACCESS_TOKEN }} secrets: | 6b0baeba-4bd1-4c7d-b0c4-b0850005549d > BW_SECRET_1 + - name: Configure 1Password Service Account + uses: 1password/load-secrets-action/configure@v1 + with: + service-account-token: ${{ secrets.OP_SERVICE_ACCOUNT_TOKEN }} + - name: Load secret + uses: 1password/load-secrets-action@v1 + env: + OP_SECRET: op://app-cicd/Better-Than-Bitwarden?/password - name: Run Tests run: | make test diff --git a/tests/test_secrets.py b/tests/test_secrets.py index 8a6bee09..fd5e2151 100644 --- a/tests/test_secrets.py +++ b/tests/test_secrets.py @@ -1,4 +1,7 @@ import os -def test_scan_pickle_bytes() -> None: - assert os.environ.get("BW_SECRET_1") == "Secret-Value-123" \ No newline at end of file +def test_bw_secret() -> None: + assert os.environ.get("BW_SECRET_1") == "Secret-Value-123" + +def test_op_secret() -> None: + assert os.environ.get("OP_SECRET_1") == "IQXya3HSTY9Suwtsp18p" \ No newline at end of file From af6a4587b474a2a586cee769769f5f229d1766e1 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Tue, 26 Sep 2023 16:21:51 -0700 Subject: [PATCH 4/5] format --- .github/workflows/test.yml | 2 +- tests/test_secrets.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7b7c08e4..02146a75 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -47,7 +47,7 @@ jobs: - name: Load secret uses: 1password/load-secrets-action@v1 env: - OP_SECRET: op://app-cicd/Better-Than-Bitwarden?/password + OP_SECRET: op://app-cicd/Better-Than-Bitwarden?/Password - name: Run Tests run: | make test diff --git a/tests/test_secrets.py b/tests/test_secrets.py index fd5e2151..21a6cd00 100644 --- a/tests/test_secrets.py +++ b/tests/test_secrets.py @@ -1,7 +1,9 @@ import os + def test_bw_secret() -> None: assert os.environ.get("BW_SECRET_1") == "Secret-Value-123" + def test_op_secret() -> None: - assert os.environ.get("OP_SECRET_1") == "IQXya3HSTY9Suwtsp18p" \ No newline at end of file + assert os.environ.get("OP_SECRET_1") == "IQXya3HSTY9Suwtsp18p" From 6cdc76c73f4afa124333c2ae8547195353771285 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Tue, 26 Sep 2023 16:33:57 -0700 Subject: [PATCH 5/5] quotes --- .github/workflows/test.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 02146a75..200c1d44 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -47,7 +47,7 @@ jobs: - name: Load secret uses: 1password/load-secrets-action@v1 env: - OP_SECRET: op://app-cicd/Better-Than-Bitwarden?/Password + OP_SECRET: op://app-cicd/"Better-Than-Bitwarden?"/Password - name: Run Tests run: | make test