Skip to content

Commit

Permalink
refactor(pickle): 🚧 dry up pickle inject code that is shared between …
Browse files Browse the repository at this point in the history
…unit tests and notebooks
  • Loading branch information
CandiedCode committed Apr 24, 2024
1 parent 6e2c28e commit 492bf57
Show file tree
Hide file tree
Showing 10 changed files with 286 additions and 418 deletions.
18 changes: 9 additions & 9 deletions notebooks/pytorch_sentiment_analysis.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -60,26 +60,26 @@
"metadata": {},
"outputs": [
{
"name": "stdout",
"name": "stderr",
"output_type": "stream",
"text": [
"env: TOKENIZERS_PARALLELISM=false\n"
"/home/vscode/.cache/pypoetry/virtualenvs/modelscan-d-AHl6rn-py3.9/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
]
},
{
"name": "stderr",
"name": "stdout",
"output_type": "stream",
"text": [
"/home/vscode/.cache/pypoetry/virtualenvs/modelscan-d-AHl6rn-py3.9/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
" from .autonotebook import tqdm as notebook_tqdm\n"
"env: TOKENIZERS_PARALLELISM=false\n"
]
}
],
"source": [
"import torch\n",
"import os\n",
"from utils.pytorch_sentiment_model import download_model, predict_sentiment\n",
"from utils.pickle_codeinjection import PickleInject, get_payload\n",
"from tests.pickle_utils.codeinjection import PickleInject, get_inject_payload\n",
"\n",
"%env TOKENIZERS_PARALLELISM=false"
]
Expand Down Expand Up @@ -234,7 +234,7 @@
"\n",
"unsafe_model_path = os.path.join(model_directory, \"unsafe_model.pt\")\n",
"\n",
"payload = get_payload(command, malicious_code)\n",
"payload = get_inject_payload(command, malicious_code)\n",
"torch.save(\n",
" torch.load(safe_model_path),\n",
" f=unsafe_model_path,\n",
Expand Down Expand Up @@ -392,7 +392,7 @@
"\u001b[32m\"CRITICAL\"\u001b[0m: \u001b[1;36m1\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m\"total_issues\"\u001b[0m: \u001b[1;36m1\u001b[0m, \u001b[32m\"input_path\"\u001b[0m: \n",
"\u001b[32m\"./PyTorchModels/unsafe_model.pt\"\u001b[0m, \u001b[32m\"absolute_path\"\u001b[0m: \n",
"\u001b[32m\"/workspaces/modelscan/notebooks/PyTorchModels\"\u001b[0m, \u001b[32m\"modelscan_version\"\u001b[0m: \u001b[32m\"0.0.0\"\u001b[0m, \n",
"\u001b[32m\"timestamp\"\u001b[0m: \u001b[32m\"2024-04-21T10:49:44.690078\"\u001b[0m, \u001b[32m\"scanned\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"total_scanned\"\u001b[0m: \u001b[1;36m1\u001b[0m, \n",
"\u001b[32m\"timestamp\"\u001b[0m: \u001b[32m\"2024-04-24T00:29:46.665913\"\u001b[0m, \u001b[32m\"scanned\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"total_scanned\"\u001b[0m: \u001b[1;36m1\u001b[0m, \n",
"\u001b[32m\"scanned_files\"\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m\"unsafe_model.pt:unsafe_model/data.pkl\"\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m\"issues\"\u001b[0m: \n",
"\u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m\"description\"\u001b[0m: \u001b[32m\"Use of unsafe operator 'system' from module 'posix'\"\u001b[0m, \n",
"\u001b[32m\"operator\"\u001b[0m: \u001b[32m\"system\"\u001b[0m, \u001b[32m\"module\"\u001b[0m: \u001b[32m\"posix\"\u001b[0m, \u001b[32m\"source\"\u001b[0m: \n",
Expand Down Expand Up @@ -423,7 +423,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.9.19"
"version": "3.10.14"
},
"vscode": {
"interpreter": {
Expand Down
131 changes: 0 additions & 131 deletions notebooks/utils/pickle_codeinjection.py

This file was deleted.

6 changes: 3 additions & 3 deletions notebooks/xgboost_diabetes_classification.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,7 @@
"from pathlib import Path\n",
"import os\n",
"import numpy as np\n",
"from utils.pickle_codeinjection import generate_unsafe_file\n",
"from tests.pickle_utils.codeinjection import generate_unsafe_pickle_file\n",
"from utils.xgboost_diabetes_model import train_model, get_predictions"
]
},
Expand Down Expand Up @@ -202,7 +202,7 @@
" safe_model_pickle = pickle.load(fo)\n",
"\n",
"unsafe_model_path = os.path.join(model_directory, \"unsafe_model.pkl\")\n",
"generate_unsafe_file(model, command, malicious_code, unsafe_model_path)"
"generate_unsafe_pickle_file(model, command, malicious_code, unsafe_model_path)"
]
},
{
Expand Down Expand Up @@ -344,7 +344,7 @@
"\u001b[32m\"CRITICAL\"\u001b[0m: \u001b[1;36m1\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m\"total_issues\"\u001b[0m: \u001b[1;36m1\u001b[0m, \u001b[32m\"input_path\"\u001b[0m: \n",
"\u001b[32m\"XGBoostModels/unsafe_model.pkl\"\u001b[0m, \u001b[32m\"absolute_path\"\u001b[0m: \n",
"\u001b[32m\"/workspaces/modelscan/notebooks/XGBoostModels\"\u001b[0m, \u001b[32m\"modelscan_version\"\u001b[0m: \u001b[32m\"0.0.0\"\u001b[0m, \n",
"\u001b[32m\"timestamp\"\u001b[0m: \u001b[32m\"2024-04-21T12:13:42.698872\"\u001b[0m, \u001b[32m\"scanned\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"total_scanned\"\u001b[0m: \u001b[1;36m1\u001b[0m, \n",
"\u001b[32m\"timestamp\"\u001b[0m: \u001b[32m\"2024-04-24T00:38:47.509270\"\u001b[0m, \u001b[32m\"scanned\"\u001b[0m: \u001b[1m{\u001b[0m\u001b[32m\"total_scanned\"\u001b[0m: \u001b[1;36m1\u001b[0m, \n",
"\u001b[32m\"scanned_files\"\u001b[0m: \u001b[1m[\u001b[0m\u001b[32m\"unsafe_model.pkl\"\u001b[0m\u001b[1m]\u001b[0m\u001b[1m}\u001b[0m\u001b[1m}\u001b[0m, \u001b[32m\"issues\"\u001b[0m: \u001b[1m[\u001b[0m\u001b[1m{\u001b[0m\u001b[32m\"description\"\u001b[0m: \u001b[32m\"Use of \u001b[0m\n",
"\u001b[32munsafe operator 'system' from module 'posix'\"\u001b[0m, \u001b[32m\"operator\"\u001b[0m: \u001b[32m\"system\"\u001b[0m, \u001b[32m\"module\"\u001b[0m: \n",
"\u001b[32m\"posix\"\u001b[0m, \u001b[32m\"source\"\u001b[0m: \u001b[32m\"unsafe_model.pkl\"\u001b[0m, \u001b[32m\"scanner\"\u001b[0m: \n",
Expand Down
11 changes: 6 additions & 5 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ h5py = ["h5py"]
[tool.poetry.group.test.dependencies]
aiohttp = "^3.8.5"
bandit = { version = "1.7.8", extras = ["toml"] }
dill = "^0.3.7"
dill = "^0.3.8"
mypy = "^1.4.1"
nbmake = "^1.5.3"
pytest = "^7.4.0"
Expand Down
Empty file added tests/__init__.py
Empty file.
Empty file added tests/pickle_utils/__init__.py
Empty file.
Loading

0 comments on commit 492bf57

Please sign in to comment.