From 11d6252ea79cf9d6912d94165afe095c1a75e9f8 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 13:19:50 -0700 Subject: [PATCH 01/21] codecov report and yml --- .github/workflows/coverage.yml | 2 +- codecov.yml | 0 2 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 codecov.yml diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 4dea8c4d..65d4dd79 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -50,7 +50,7 @@ jobs: OP_SECRET: op://app-cicd/"Better-Than-Bitwarden?"/Password - name: Run Coverage run: | - make test + make coverage - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: diff --git a/codecov.yml b/codecov.yml new file mode 100644 index 00000000..e69de29b From cd5a27248f2c3e9e54e320dc3e3bb82fe857c389 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 13:20:42 -0700 Subject: [PATCH 02/21] comment out hf test --- tests/test_modelscan.py | 30 +++++++++++++++--------------- 1 file changed, 15 insertions(+), 15 deletions(-) diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index d6395226..6a37f5d2 100644 --- a/tests/test_modelscan.py +++ b/tests/test_modelscan.py @@ -685,21 +685,21 @@ def test_scan_directory_path(pickle_file_path: str) -> None: 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_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: From eba56f878c3ea1c89d7da59f4e6f90cbb0ba3231 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 13:25:52 -0700 Subject: [PATCH 03/21] blog format --- .github/workflows/coverage.yml | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 65d4dd79..f57c7326 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -11,7 +11,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: ["3.8", "3.9", "3.10"] + python-version: ["3.8"] steps: - uses: actions/checkout@v2 @@ -50,7 +50,8 @@ jobs: OP_SECRET: op://app-cicd/"Better-Than-Bitwarden?"/Password - name: Run Coverage run: | - make coverage + pip install coverage + coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: From 8ae689aaafe239f92adbef1ff3e7a8b8dd934a47 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 13:31:08 -0700 Subject: [PATCH 04/21] make install-test --- .github/workflows/coverage.yml | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index f57c7326..f7ba9307 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -48,6 +48,10 @@ jobs: uses: 1password/load-secrets-action@v1 env: OP_SECRET: op://app-cicd/"Better-Than-Bitwarden?"/Password + - name: Install Dependencies + if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' + run: | + make install-test - name: Run Coverage run: | pip install coverage From 74d46946895d0168cf991f951252e6b360eeb358 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 13:34:52 -0700 Subject: [PATCH 05/21] poetry run --- .github/workflows/coverage.yml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index f7ba9307..69d2eb52 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -54,8 +54,8 @@ jobs: make install-test - name: Run Coverage run: | - pip install coverage - coverage run -m pytest + pip install coverage + poetry run coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: From b6d25c8ce4b0b6344952d1183e99d0f54e195143 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 13:43:26 -0700 Subject: [PATCH 06/21] make --- .github/workflows/coverage.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 69d2eb52..dfa1993f 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -54,8 +54,7 @@ jobs: make install-test - name: Run Coverage run: | - pip install coverage - poetry run coverage run -m pytest + make coverage - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: From f24f0bda469b62151cdb27556ebbfe64145388d5 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:17:37 -0700 Subject: [PATCH 07/21] overhaul --- .github/workflows/coverage.yml | 43 +- tests/test_cov.py | 11 + tests/test_modelscan.py | 1432 ++++++++++++++++---------------- 3 files changed, 733 insertions(+), 753 deletions(-) create mode 100644 tests/test_cov.py diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index dfa1993f..1ac0c1ce 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -9,52 +9,21 @@ on: jobs: test: runs-on: ubuntu-latest - strategy: - matrix: - python-version: ["3.8"] + env: + OS: ubuntu-latest + PYTHON: '3.9' steps: - uses: actions/checkout@v2 - - name: Set up Python ${{ matrix.python-version }} + - name: Set up Python uses: actions/setup-python@v2 - with: - python-version: ${{ matrix.python-version }} - - uses: snok/install-poetry@v1 - with: - virtualenvs-create: true - virtualenvs-in-project: true - installer-parallel: true - - name: Load cached venv - id: cached-poetry-dependencies - uses: actions/cache@v3 - with: - path: .venv - key: venv-test-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }} - name: Install Dependencies - if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' run: | make install-test - - 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: Install Dependencies - if: steps.cached-poetry-dependencies.outputs.cache-hit != 'true' - run: | - make install-test - name: Run Coverage run: | - make coverage + pip install coverage + poetry run coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: diff --git a/tests/test_cov.py b/tests/test_cov.py new file mode 100644 index 00000000..620ba278 --- /dev/null +++ b/tests/test_cov.py @@ -0,0 +1,11 @@ +import os +from pathlib import Path +from modelscan.modelscan import Modelscan + + +def coverage() -> None: + try: + ms = Modelscan() + ms.scan_path(Path("")) + except Exception: + pass diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index 6a37f5d2..905a80ca 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",) -@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) +# class Malicious3: +# def __reduce__(self) -> Any: +# return http.client.HTTPSConnection, ("github.com",) - # 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: +# 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) + +# # 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( -# "__builtin__", -# "eval", -# "https://huggingface.co/ykilcher/totally-harmless-model/resolve/main/pytorch_model.bin:archive/data.pkl", +# "builtins", "eval", f"{zip_file_path}/test.zip:data.pkl" # ), # ) # ] + # ms = Modelscan() -# ms.scan_huggingface_model("ykilcher/totally-harmless-model") +# 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_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 +# 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 From 112585971eecc1b065dc1546ee35eb7fd1ea960d Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:19:26 -0700 Subject: [PATCH 08/21] no poetry --- .github/workflows/coverage.yml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 1ac0c1ce..098c176d 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -17,13 +17,10 @@ jobs: - uses: actions/checkout@v2 - name: Set up Python uses: actions/setup-python@v2 - - name: Install Dependencies - run: | - make install-test - name: Run Coverage run: | pip install coverage - poetry run coverage run -m pytest + coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: From 2662ca1ec3da752d8a24503981458c4c3d28f627 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:21:18 -0700 Subject: [PATCH 09/21] put back poetry to make --- .github/workflows/coverage.yml | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 098c176d..a7a06a02 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -17,6 +17,14 @@ jobs: - uses: actions/checkout@v2 - name: Set up Python uses: actions/setup-python@v2 + - uses: snok/install-poetry@v1 + with: + virtualenvs-create: true + virtualenvs-in-project: true + installer-parallel: true + - name: Install Dependencies + run: | + make install-test - name: Run Coverage run: | pip install coverage From 57a0e6f134b283d8283ed723c6fb9badeb616f03 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:22:59 -0700 Subject: [PATCH 10/21] Update coverage.yml --- .github/workflows/coverage.yml | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index a7a06a02..284da75c 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -22,6 +22,12 @@ jobs: virtualenvs-create: true virtualenvs-in-project: true installer-parallel: true + - name: Load cached venv + id: cached-poetry-dependencies + uses: actions/cache@v3 + with: + path: .venv + key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }} - name: Install Dependencies run: | make install-test From ee95c8e329e34b5390347453cbda56450ab00dff Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:23:44 -0700 Subject: [PATCH 11/21] pytest --- .github/workflows/coverage.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 284da75c..6c1945ab 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -34,6 +34,7 @@ jobs: - name: Run Coverage run: | pip install coverage + pip install pytest coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 From 727f1a492b68877f2844c242a93db7475b85a994 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:25:19 -0700 Subject: [PATCH 12/21] python name --- .github/workflows/coverage.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 6c1945ab..fd188bea 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -17,6 +17,9 @@ jobs: - uses: actions/checkout@v2 - name: Set up Python uses: actions/setup-python@v2 + id: setup-python + with: + python-version: "3.8" - uses: snok/install-poetry@v1 with: virtualenvs-create: true @@ -30,7 +33,7 @@ jobs: key: venv-${{ runner.os }}-${{ steps.setup-python.outputs.python-version }}-${{ hashFiles('**/poetry.lock') }} - name: Install Dependencies run: | - make install-test + make install-test - name: Run Coverage run: | pip install coverage From d16fc61b908d2068b78ac13452cd23d80f18ba87 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:27:55 -0700 Subject: [PATCH 13/21] test --- tests/test_cov.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_cov.py b/tests/test_cov.py index 620ba278..869b25a7 100644 --- a/tests/test_cov.py +++ b/tests/test_cov.py @@ -3,7 +3,7 @@ from modelscan.modelscan import Modelscan -def coverage() -> None: +def test_coverage() -> None: try: ms = Modelscan() ms.scan_path(Path("")) From ca092c34c6d46cb986299fa26f001f76931d7a56 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:30:03 -0700 Subject: [PATCH 14/21] Update test_utils.py --- tests/test_utils.py | 414 ++++++++++++++++++++++---------------------- 1 file changed, 207 insertions(+), 207 deletions(-) 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 189c50253925337279baa1a332d1517b39d74ff3 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:33:05 -0700 Subject: [PATCH 15/21] empty test --- tests/test_cov.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/tests/test_cov.py b/tests/test_cov.py index 869b25a7..59f31677 100644 --- a/tests/test_cov.py +++ b/tests/test_cov.py @@ -1,11 +1,12 @@ -import os -from pathlib import Path -from modelscan.modelscan import Modelscan +# import os +# from pathlib import Path +# from modelscan.modelscan import Modelscan def test_coverage() -> None: - try: - ms = Modelscan() - ms.scan_path(Path("")) - except Exception: - pass + # try: + # ms = Modelscan() + # ms.scan_path(Path("")) + # except Exception: + # pass + pass From 9495acdb437e8e631d69fbcc172aeffd5c4a0d1e Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:47:09 -0700 Subject: [PATCH 16/21] poetry --- .github/workflows/coverage.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index fd188bea..d8024ea0 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -38,7 +38,7 @@ jobs: run: | pip install coverage pip install pytest - coverage run -m pytest + poetry run coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: From 19ade95728dd95b2460b82e47ac0ffd5a6f7a77e Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 14:55:36 -0700 Subject: [PATCH 17/21] poetry add --- .github/workflows/coverage.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index d8024ea0..9ed7e675 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -38,8 +38,10 @@ jobs: run: | pip install coverage pip install pytest + poetry add coverage codecov poetry run coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: token: ${{ secrets.CODECOV_TOKEN }} + verbose: true From 69df0e87aa96b1ced1c3e1f0498400c8908c37a3 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 15:00:06 -0700 Subject: [PATCH 18/21] make coverage --- .github/workflows/coverage.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 9ed7e675..4a4fe2fd 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -36,10 +36,8 @@ jobs: make install-test - name: Run Coverage run: | - pip install coverage - pip install pytest poetry add coverage codecov - poetry run coverage run -m pytest + make coverage - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: From 5445b7733fdbd81fbd93cf36dccfc11c568a6338 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 15:06:03 -0700 Subject: [PATCH 19/21] poetry add coverage and codecov to test group --- .github/workflows/coverage.yml | 3 +- poetry.lock | 120 ++++++++++++++++++++------------- pyproject.toml | 2 + 3 files changed, 75 insertions(+), 50 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 4a4fe2fd..15d2fbcd 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -36,8 +36,7 @@ jobs: make install-test - name: Run Coverage run: | - poetry add coverage codecov - make coverage + poetry run coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 with: diff --git a/poetry.lock b/poetry.lock index 2d571c61..9f7d3433 100644 --- a/poetry.lock +++ b/poetry.lock @@ -378,6 +378,21 @@ files = [ [package.dependencies] colorama = {version = "*", markers = "platform_system == \"Windows\""} +[[package]] +name = "codecov" +version = "2.1.13" +description = "Hosted coverage reports for GitHub, Bitbucket and Gitlab" +optional = false +python-versions = ">=2.7, !=3.0.*, !=3.1.*, !=3.2.*, !=3.3.*" +files = [ + {file = "codecov-2.1.13-py2.py3-none-any.whl", hash = "sha256:c2ca5e51bba9ebb43644c43d0690148a55086f7f5e6fd36170858fa4206744d5"}, + {file = "codecov-2.1.13.tar.gz", hash = "sha256:2362b685633caeaf45b9951a9b76ce359cd3581dd515b430c6c3f5dfb4d92a8c"}, +] + +[package.dependencies] +coverage = "*" +requests = ">=2.7.9" + [[package]] name = "colorama" version = "0.4.6" @@ -706,60 +721,69 @@ six = "*" [[package]] name = "grpcio" -version = "1.58.0" +version = "1.59.0" description = "HTTP/2-based RPC framework" optional = true python-versions = ">=3.7" files = [ - {file = "grpcio-1.58.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:3e6bebf1dfdbeb22afd95650e4f019219fef3ab86d3fca8ebade52e4bc39389a"}, - {file = "grpcio-1.58.0-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:cde11577d5b6fd73a00e6bfa3cf5f428f3f33c2d2878982369b5372bbc4acc60"}, - {file = "grpcio-1.58.0-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:a2d67ff99e70e86b2be46c1017ae40b4840d09467d5455b2708de6d4c127e143"}, - {file = "grpcio-1.58.0-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1ed979b273a81de36fc9c6716d9fb09dd3443efa18dcc8652501df11da9583e9"}, - {file = "grpcio-1.58.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:458899d2ebd55d5ca2350fd3826dfd8fcb11fe0f79828ae75e2b1e6051d50a29"}, - {file = "grpcio-1.58.0-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:bc7ffef430b80345729ff0a6825e9d96ac87efe39216e87ac58c6c4ef400de93"}, - {file = "grpcio-1.58.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:5b23d75e5173faa3d1296a7bedffb25afd2fddb607ef292dfc651490c7b53c3d"}, - {file = "grpcio-1.58.0-cp310-cp310-win32.whl", hash = "sha256:fad9295fe02455d4f158ad72c90ef8b4bcaadfdb5efb5795f7ab0786ad67dd58"}, - {file = "grpcio-1.58.0-cp310-cp310-win_amd64.whl", hash = "sha256:bc325fed4d074367bebd465a20763586e5e1ed5b943e9d8bc7c162b1f44fd602"}, - {file = "grpcio-1.58.0-cp311-cp311-linux_armv7l.whl", hash = "sha256:652978551af02373a5a313e07bfef368f406b5929cf2d50fa7e4027f913dbdb4"}, - {file = "grpcio-1.58.0-cp311-cp311-macosx_10_10_universal2.whl", hash = "sha256:9f13a171281ebb4d7b1ba9f06574bce2455dcd3f2f6d1fbe0fd0d84615c74045"}, - {file = "grpcio-1.58.0-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:8774219e21b05f750eef8adc416e9431cf31b98f6ce9def288e4cea1548cbd22"}, - {file = "grpcio-1.58.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:09206106848462763f7f273ca93d2d2d4d26cab475089e0de830bb76be04e9e8"}, - {file = "grpcio-1.58.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:62831d5e251dd7561d9d9e83a0b8655084b2a1f8ea91e4bd6b3cedfefd32c9d2"}, - {file = "grpcio-1.58.0-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:212f38c6a156862098f6bdc9a79bf850760a751d259d8f8f249fc6d645105855"}, - {file = "grpcio-1.58.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:4b12754af201bb993e6e2efd7812085ddaaef21d0a6f0ff128b97de1ef55aa4a"}, - {file = "grpcio-1.58.0-cp311-cp311-win32.whl", hash = "sha256:3886b4d56bd4afeac518dbc05933926198aa967a7d1d237a318e6fbc47141577"}, - {file = "grpcio-1.58.0-cp311-cp311-win_amd64.whl", hash = "sha256:002f228d197fea12797a14e152447044e14fb4fdb2eb5d6cfa496f29ddbf79ef"}, - {file = "grpcio-1.58.0-cp37-cp37m-linux_armv7l.whl", hash = "sha256:b5e8db0aff0a4819946215f156bd722b6f6c8320eb8419567ffc74850c9fd205"}, - {file = "grpcio-1.58.0-cp37-cp37m-macosx_10_10_universal2.whl", hash = "sha256:201e550b7e2ede113b63e718e7ece93cef5b0fbf3c45e8fe4541a5a4305acd15"}, - {file = "grpcio-1.58.0-cp37-cp37m-manylinux_2_17_aarch64.whl", hash = "sha256:d79b660681eb9bc66cc7cbf78d1b1b9e335ee56f6ea1755d34a31108b80bd3c8"}, - {file = "grpcio-1.58.0-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:2ef8d4a76d2c7d8065aba829f8d0bc0055495c998dce1964ca5b302d02514fb3"}, - {file = "grpcio-1.58.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6cba491c638c76d3dc6c191d9c75041ca5b8f5c6de4b8327ecdcab527f130bb4"}, - {file = "grpcio-1.58.0-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:6801ff6652ecd2aae08ef994a3e49ff53de29e69e9cd0fd604a79ae4e545a95c"}, - {file = "grpcio-1.58.0-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:24edec346e69e672daf12b2c88e95c6f737f3792d08866101d8c5f34370c54fd"}, - {file = "grpcio-1.58.0-cp37-cp37m-win_amd64.whl", hash = "sha256:7e473a7abad9af48e3ab5f3b5d237d18208024d28ead65a459bd720401bd2f8f"}, - {file = "grpcio-1.58.0-cp38-cp38-linux_armv7l.whl", hash = "sha256:4891bbb4bba58acd1d620759b3be11245bfe715eb67a4864c8937b855b7ed7fa"}, - {file = "grpcio-1.58.0-cp38-cp38-macosx_10_10_universal2.whl", hash = "sha256:e9f995a8a421405958ff30599b4d0eec244f28edc760de82f0412c71c61763d2"}, - {file = "grpcio-1.58.0-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:2f85f87e2f087d9f632c085b37440a3169fda9cdde80cb84057c2fc292f8cbdf"}, - {file = "grpcio-1.58.0-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:eb6b92036ff312d5b4182fa72e8735d17aceca74d0d908a7f08e375456f03e07"}, - {file = "grpcio-1.58.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d81c2b2b24c32139dd2536972f1060678c6b9fbd106842a9fcdecf07b233eccd"}, - {file = "grpcio-1.58.0-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:fbcecb6aedd5c1891db1d70efbfbdc126c986645b5dd616a045c07d6bd2dfa86"}, - {file = "grpcio-1.58.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:92ae871a902cf19833328bd6498ec007b265aabf2fda845ab5bd10abcaf4c8c6"}, - {file = "grpcio-1.58.0-cp38-cp38-win32.whl", hash = "sha256:dc72e04620d49d3007771c0e0348deb23ca341c0245d610605dddb4ac65a37cb"}, - {file = "grpcio-1.58.0-cp38-cp38-win_amd64.whl", hash = "sha256:1c1c5238c6072470c7f1614bf7c774ffde6b346a100521de9ce791d1e4453afe"}, - {file = "grpcio-1.58.0-cp39-cp39-linux_armv7l.whl", hash = "sha256:fe643af248442221db027da43ed43e53b73e11f40c9043738de9a2b4b6ca7697"}, - {file = "grpcio-1.58.0-cp39-cp39-macosx_10_10_universal2.whl", hash = "sha256:128eb1f8e70676d05b1b0c8e6600320fc222b3f8c985a92224248b1367122188"}, - {file = "grpcio-1.58.0-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:039003a5e0ae7d41c86c768ef8b3ee2c558aa0a23cf04bf3c23567f37befa092"}, - {file = "grpcio-1.58.0-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:8f061722cad3f9aabb3fbb27f3484ec9d4667b7328d1a7800c3c691a98f16bb0"}, - {file = "grpcio-1.58.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba0af11938acf8cd4cf815c46156bcde36fa5850518120920d52620cc3ec1830"}, - {file = "grpcio-1.58.0-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:d4cef77ad2fed42b1ba9143465856d7e737279854e444925d5ba45fc1f3ba727"}, - {file = "grpcio-1.58.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:24765a627eb4d9288ace32d5104161c3654128fe27f2808ecd6e9b0cfa7fc8b9"}, - {file = "grpcio-1.58.0-cp39-cp39-win32.whl", hash = "sha256:f0241f7eb0d2303a545136c59bc565a35c4fc3b924ccbd69cb482f4828d6f31c"}, - {file = "grpcio-1.58.0-cp39-cp39-win_amd64.whl", hash = "sha256:dcfba7befe3a55dab6fe1eb7fc9359dc0c7f7272b30a70ae0af5d5b063842f28"}, - {file = "grpcio-1.58.0.tar.gz", hash = "sha256:532410c51ccd851b706d1fbc00a87be0f5312bd6f8e5dbf89d4e99c7f79d7499"}, + {file = "grpcio-1.59.0-cp310-cp310-linux_armv7l.whl", hash = "sha256:225e5fa61c35eeaebb4e7491cd2d768cd8eb6ed00f2664fa83a58f29418b39fd"}, + {file = "grpcio-1.59.0-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:b95ec8ecc4f703f5caaa8d96e93e40c7f589bad299a2617bdb8becbcce525539"}, + {file = "grpcio-1.59.0-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:1a839ba86764cc48226f50b924216000c79779c563a301586a107bda9cbe9dcf"}, + {file = "grpcio-1.59.0-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f6cfe44a5d7c7d5f1017a7da1c8160304091ca5dc64a0f85bca0d63008c3137a"}, + {file = "grpcio-1.59.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d0fcf53df684fcc0154b1e61f6b4a8c4cf5f49d98a63511e3f30966feff39cd0"}, + {file = "grpcio-1.59.0-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:fa66cac32861500f280bb60fe7d5b3e22d68c51e18e65367e38f8669b78cea3b"}, + {file = "grpcio-1.59.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:8cd2d38c2d52f607d75a74143113174c36d8a416d9472415eab834f837580cf7"}, + {file = "grpcio-1.59.0-cp310-cp310-win32.whl", hash = "sha256:228b91ce454876d7eed74041aff24a8f04c0306b7250a2da99d35dd25e2a1211"}, + {file = "grpcio-1.59.0-cp310-cp310-win_amd64.whl", hash = "sha256:ca87ee6183421b7cea3544190061f6c1c3dfc959e0b57a5286b108511fd34ff4"}, + {file = "grpcio-1.59.0-cp311-cp311-linux_armv7l.whl", hash = "sha256:c173a87d622ea074ce79be33b952f0b424fa92182063c3bda8625c11d3585d09"}, + {file = "grpcio-1.59.0-cp311-cp311-macosx_10_10_universal2.whl", hash = "sha256:ec78aebb9b6771d6a1de7b6ca2f779a2f6113b9108d486e904bde323d51f5589"}, + {file = "grpcio-1.59.0-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:0b84445fa94d59e6806c10266b977f92fa997db3585f125d6b751af02ff8b9fe"}, + {file = "grpcio-1.59.0-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c251d22de8f9f5cca9ee47e4bade7c5c853e6e40743f47f5cc02288ee7a87252"}, + {file = "grpcio-1.59.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:956f0b7cb465a65de1bd90d5a7475b4dc55089b25042fe0f6c870707e9aabb1d"}, + {file = "grpcio-1.59.0-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:38da5310ef84e16d638ad89550b5b9424df508fd5c7b968b90eb9629ca9be4b9"}, + {file = "grpcio-1.59.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:63982150a7d598281fa1d7ffead6096e543ff8be189d3235dd2b5604f2c553e5"}, + {file = "grpcio-1.59.0-cp311-cp311-win32.whl", hash = "sha256:50eff97397e29eeee5df106ea1afce3ee134d567aa2c8e04fabab05c79d791a7"}, + {file = "grpcio-1.59.0-cp311-cp311-win_amd64.whl", hash = "sha256:15f03bd714f987d48ae57fe092cf81960ae36da4e520e729392a59a75cda4f29"}, + {file = "grpcio-1.59.0-cp312-cp312-linux_armv7l.whl", hash = "sha256:f1feb034321ae2f718172d86b8276c03599846dc7bb1792ae370af02718f91c5"}, + {file = "grpcio-1.59.0-cp312-cp312-macosx_10_10_universal2.whl", hash = "sha256:d09bd2a4e9f5a44d36bb8684f284835c14d30c22d8ec92ce796655af12163588"}, + {file = "grpcio-1.59.0-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:2f120d27051e4c59db2f267b71b833796770d3ea36ca712befa8c5fff5da6ebd"}, + {file = "grpcio-1.59.0-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ba0ca727a173ee093f49ead932c051af463258b4b493b956a2c099696f38aa66"}, + {file = "grpcio-1.59.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5711c51e204dc52065f4a3327dca46e69636a0b76d3e98c2c28c4ccef9b04c52"}, + {file = "grpcio-1.59.0-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:d74f7d2d7c242a6af9d4d069552ec3669965b74fed6b92946e0e13b4168374f9"}, + {file = "grpcio-1.59.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:3859917de234a0a2a52132489c4425a73669de9c458b01c9a83687f1f31b5b10"}, + {file = "grpcio-1.59.0-cp312-cp312-win32.whl", hash = "sha256:de2599985b7c1b4ce7526e15c969d66b93687571aa008ca749d6235d056b7205"}, + {file = "grpcio-1.59.0-cp312-cp312-win_amd64.whl", hash = "sha256:598f3530231cf10ae03f4ab92d48c3be1fee0c52213a1d5958df1a90957e6a88"}, + {file = "grpcio-1.59.0-cp37-cp37m-linux_armv7l.whl", hash = "sha256:b34c7a4c31841a2ea27246a05eed8a80c319bfc0d3e644412ec9ce437105ff6c"}, + {file = "grpcio-1.59.0-cp37-cp37m-macosx_10_10_universal2.whl", hash = "sha256:c4dfdb49f4997dc664f30116af2d34751b91aa031f8c8ee251ce4dcfc11277b0"}, + {file = "grpcio-1.59.0-cp37-cp37m-manylinux_2_17_aarch64.whl", hash = "sha256:61bc72a00ecc2b79d9695220b4d02e8ba53b702b42411397e831c9b0589f08a3"}, + {file = "grpcio-1.59.0-cp37-cp37m-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:f367e4b524cb319e50acbdea57bb63c3b717c5d561974ace0b065a648bb3bad3"}, + {file = "grpcio-1.59.0-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:849c47ef42424c86af069a9c5e691a765e304079755d5c29eff511263fad9c2a"}, + {file = "grpcio-1.59.0-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:c0488c2b0528e6072010182075615620071371701733c63ab5be49140ed8f7f0"}, + {file = "grpcio-1.59.0-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:611d9aa0017fa386809bddcb76653a5ab18c264faf4d9ff35cb904d44745f575"}, + {file = "grpcio-1.59.0-cp37-cp37m-win_amd64.whl", hash = "sha256:e5378785dce2b91eb2e5b857ec7602305a3b5cf78311767146464bfa365fc897"}, + {file = "grpcio-1.59.0-cp38-cp38-linux_armv7l.whl", hash = "sha256:fe976910de34d21057bcb53b2c5e667843588b48bf11339da2a75f5c4c5b4055"}, + {file = "grpcio-1.59.0-cp38-cp38-macosx_10_10_universal2.whl", hash = "sha256:c041a91712bf23b2a910f61e16565a05869e505dc5a5c025d429ca6de5de842c"}, + {file = "grpcio-1.59.0-cp38-cp38-manylinux_2_17_aarch64.whl", hash = "sha256:0ae444221b2c16d8211b55326f8ba173ba8f8c76349bfc1768198ba592b58f74"}, + {file = "grpcio-1.59.0-cp38-cp38-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ceb1e68135788c3fce2211de86a7597591f0b9a0d2bb80e8401fd1d915991bac"}, + {file = "grpcio-1.59.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c4b1cc3a9dc1924d2eb26eec8792fedd4b3fcd10111e26c1d551f2e4eda79ce"}, + {file = "grpcio-1.59.0-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:871371ce0c0055d3db2a86fdebd1e1d647cf21a8912acc30052660297a5a6901"}, + {file = "grpcio-1.59.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:93e9cb546e610829e462147ce724a9cb108e61647a3454500438a6deef610be1"}, + {file = "grpcio-1.59.0-cp38-cp38-win32.whl", hash = "sha256:f21917aa50b40842b51aff2de6ebf9e2f6af3fe0971c31960ad6a3a2b24988f4"}, + {file = "grpcio-1.59.0-cp38-cp38-win_amd64.whl", hash = "sha256:14890da86a0c0e9dc1ea8e90101d7a3e0e7b1e71f4487fab36e2bfd2ecadd13c"}, + {file = "grpcio-1.59.0-cp39-cp39-linux_armv7l.whl", hash = "sha256:34341d9e81a4b669a5f5dca3b2a760b6798e95cdda2b173e65d29d0b16692857"}, + {file = "grpcio-1.59.0-cp39-cp39-macosx_10_10_universal2.whl", hash = "sha256:986de4aa75646e963466b386a8c5055c8b23a26a36a6c99052385d6fe8aaf180"}, + {file = "grpcio-1.59.0-cp39-cp39-manylinux_2_17_aarch64.whl", hash = "sha256:aca8a24fef80bef73f83eb8153f5f5a0134d9539b4c436a716256b311dda90a6"}, + {file = "grpcio-1.59.0-cp39-cp39-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:936b2e04663660c600d5173bc2cc84e15adbad9c8f71946eb833b0afc205b996"}, + {file = "grpcio-1.59.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fc8bf2e7bc725e76c0c11e474634a08c8f24bcf7426c0c6d60c8f9c6e70e4d4a"}, + {file = "grpcio-1.59.0-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:81d86a096ccd24a57fa5772a544c9e566218bc4de49e8c909882dae9d73392df"}, + {file = "grpcio-1.59.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:2ea95cd6abbe20138b8df965b4a8674ec312aaef3147c0f46a0bac661f09e8d0"}, + {file = "grpcio-1.59.0-cp39-cp39-win32.whl", hash = "sha256:3b8ff795d35a93d1df6531f31c1502673d1cebeeba93d0f9bd74617381507e3f"}, + {file = "grpcio-1.59.0-cp39-cp39-win_amd64.whl", hash = "sha256:38823bd088c69f59966f594d087d3a929d1ef310506bee9e3648317660d65b81"}, + {file = "grpcio-1.59.0.tar.gz", hash = "sha256:acf70a63cf09dd494000007b798aff88a436e1c03b394995ce450be437b8e54f"}, ] [package.extras] -protobuf = ["grpcio-tools (>=1.58.0)"] +protobuf = ["grpcio-tools (>=1.59.0)"] [[package]] name = "h5py" @@ -2140,4 +2164,4 @@ tensorflow = ["tensorflow", "tensorflow-macos"] [metadata] lock-version = "2.0" python-versions = ">=3.8,<3.12" -content-hash = "8a05aa605ce557241dd71fed9a13453959593d6727880d0b5023feec61f217b1" +content-hash = "63464c3d949dae7ab8e6f4866faa9d6c5ddc8f0b467974fd0033138e807ad5ad" diff --git a/pyproject.toml b/pyproject.toml index 156fc329..ed13ec1b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,8 @@ dill = "^0.3.7" pandas = "^2.0.3" requests = "^2.31.0" pytest-cov = "^4.1.0" +coverage = "^7.3.1" +codecov = "^2.1.13" [tool.poetry.group.dev.dependencies] dunamai = "^1.18.0" From b48b380b3d31e0eba5b734773a5bdb0bcb23f707 Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 15:09:04 -0700 Subject: [PATCH 20/21] uncomment tests --- .github/workflows/coverage.yml | 1 + tests/test_modelscan.py | 1432 ++++++++++++++++---------------- tests/test_utils.py | 414 ++++----- 3 files changed, 924 insertions(+), 923 deletions(-) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index 15d2fbcd..f0c31e67 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -36,6 +36,7 @@ jobs: make install-test - name: Run Coverage run: | + poetry add coverage codecov poetry run coverage run -m pytest - name: Upload coverage reports to Codecov uses: codecov/codecov-action@v3 diff --git a/tests/test_modelscan.py b/tests/test_modelscan.py index 905a80ca..6a37f5d2 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 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 Malicious3: -# def __reduce__(self) -> Any: -# return http.client.HTTPSConnection, ("github.com",) +@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 + ) -# 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) - -# # 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: + # 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( -# "builtins", "eval", f"{zip_file_path}/test.zip:data.pkl" +# "__builtin__", +# "eval", +# "https://huggingface.co/ykilcher/totally-harmless-model/resolve/main/pytorch_model.bin:archive/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")) +# ms.scan_huggingface_model("ykilcher/totally-harmless-model") # 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 +# 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_utils.py b/tests/test_utils.py index 691520a5..34b3eeb5 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 5ad3809eb38fd1b2f101d3dfd5b8af44877609cc Mon Sep 17 00:00:00 2001 From: Sam Washko Date: Fri, 29 Sep 2023 15:11:55 -0700 Subject: [PATCH 21/21] pip --- .github/workflows/coverage.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/coverage.yml b/.github/workflows/coverage.yml index f0c31e67..ce6b4276 100644 --- a/.github/workflows/coverage.yml +++ b/.github/workflows/coverage.yml @@ -36,6 +36,7 @@ jobs: make install-test - name: Run Coverage run: | + pip install coverage poetry add coverage codecov poetry run coverage run -m pytest - name: Upload coverage reports to Codecov