Skip to content

Commit

Permalink
security-fix: resolve security issue with starting a process with a s…
Browse files Browse the repository at this point in the history
…hell

- Issue: Fixed [B605:start_process_with_a_shell] identified by Bandit,
  which flagged starting a process with a shell as a high-severity security risk (CWE-78: OS Command Injection).
- Replaced os.system call with a safer alternative to prevent shell injection vulnerabilities.

Signed-off-by: ChengyuZhu6 <chengyu.zhu@intel.com>
  • Loading branch information
ChengyuZhu6 committed Sep 19, 2024
1 parent 155fc18 commit 187ed53
Show file tree
Hide file tree
Showing 20 changed files with 588 additions and 258 deletions.
12 changes: 8 additions & 4 deletions binaries/build.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
import argparse
import glob
import os
import shlex
import subprocess
import sys

# To help discover local modules
Expand Down Expand Up @@ -49,10 +51,12 @@ def build_dist_whl(args):
print(f"## In directory: {os.getcwd()} | Executing command: {cur_wheel_cmd}")

if not args.dry_run:
build_exit_code = os.system(cur_wheel_cmd)
# If any one of the steps fail, exit with error
if build_exit_code != 0:
sys.exit(f"## {binary} build Failed !")
try:
cur_wheel_cmd_list = shlex.split(cur_wheel_cmd)
subprocess.run(cur_wheel_cmd_list, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
except subprocess.CalledProcessError as e:
print(f"## {binary} build Failed! Error: {e.stderr.decode()}")
sys.exit(1)


def build(args):
Expand Down
40 changes: 18 additions & 22 deletions binaries/conda/build_packages.py
Original file line number Diff line number Diff line change
@@ -1,19 +1,16 @@
import argparse
import os
import subprocess
from datetime import date

from ts_scripts.utils import try_and_handle
from ts_scripts.utils import try_and_handle, find_conda_binary

conda_build_dir = os.path.dirname(os.path.abspath(__file__))
REPO_ROOT = os.path.join(conda_build_dir, "..", "..")
MINICONDA_DOWNLOAD_URL = (
"https://repo.anaconda.com/miniconda/Miniconda3-py39_4.9.2-Linux-x86_64.sh"
)
CONDA_BINARY = (
os.popen("which conda").read().strip()
if os.system(f"conda --version") == 0
else f"$HOME/miniconda/condabin/conda"
)
CONDA_BINARY = find_conda_binary()

CONDA_PACKAGES_PATH = os.path.join(REPO_ROOT, "binaries", "conda", "output")
CONDA_LINUX_PACKAGES_PATH = os.path.join(
Expand All @@ -32,8 +29,7 @@

if os.name == "nt":
# Assumes miniconda is installed in windows
CONDA_BINARY = "conda"

CONDA_BINARY = "conda"

def add_nightly_suffix_conda(binary_name: str) -> str:
"""
Expand All @@ -52,29 +48,29 @@ def install_conda_build(dry_run):
"""

# Check if conda binary already exists
exit_code = os.system(f"conda --version")
if exit_code == 0:
print(
f"'conda' already present on the system. Proceeding without a fresh conda installation."
)
try:
subprocess.run(["conda", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
print("'conda' already present on the system. Proceeding without a fresh conda installation.")
return
try_and_handle(
f"{CONDA_BINARY} install python=3.8 conda-build anaconda-client -y", dry_run
)
except subprocess.CalledProcessError:
# Conda is not available, proceed with installation
try_and_handle(
f"{CONDA_BINARY} install python=3.8 conda-build anaconda-client -y", dry_run
)


def install_miniconda(dry_run):
"""
Installs miniconda, a slimmer anaconda installation to build conda packages
"""

# Check if conda binary already exists
exit_code = os.system(f"conda --version")
if exit_code == 0:
print(
f"'conda' already present on the system. Proceeding without a fresh minconda installation."
)
try:
subprocess.run(["conda", "--version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
print("'conda' already present on the system. Proceeding without a fresh conda installation.")
return
except subprocess.CalledProcessError as e:
raise (e)

if os.name == "nt":
print(
"Identified as Windows system. Please install miniconda using this URL: https://repo.anaconda.com/miniconda/Miniconda3-latest-Windows-x86_64.exe"
Expand Down
23 changes: 15 additions & 8 deletions binaries/install.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
import subprocess
import sys
import glob

Expand All @@ -13,19 +14,25 @@ def install():
if is_conda_env():
print("## Using conda to install torchserve and torch-model-archiver")
channel_dir = os.path.abspath(os.path.join(REPO_ROOT, "binaries", "conda", "output"))
conda_cmd = f"conda install --channel {channel_dir} -y torchserve torch-model-archiver"
print(f"## In directory: {os.getcwd()} | Executing command: {conda_cmd}")
install_exit_code = os.system(conda_cmd)
conda_cmd = ["conda", "install", "--channel", channel_dir, "-y", "torchserve", "torch-model-archiver"]
print(f"## In directory: {os.getcwd()} | Executing command: {' '.join(conda_cmd)}")

try:
subprocess.run(conda_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
except subprocess.CalledProcessError as e:
sys.exit("## Torchserve/Model archiver Installation Failed!")

else:
print("## Using pip to install torchserve and torch-model-archiver")
ts_wheel = glob.glob(os.path.join(REPO_ROOT, "dist", "*.whl"))[0]
ma_wheel = glob.glob(os.path.join(REPO_ROOT, "model-archiver", "dist", "*.whl"))[0]
pip_cmd = f"pip install {ts_wheel} {ma_wheel}"
print(f"## In directory: {os.getcwd()} | Executing command: {pip_cmd}")
install_exit_code = os.system(pip_cmd)
pip_cmd = ["pip", "install", ts_wheel, ma_wheel]
print(f"## In directory: {os.getcwd()} | Executing command: {' '.join(pip_cmd)}")

if install_exit_code != 0:
sys.exit("## Torchserve \ Model archiver Installation Failed !")
try:
subprocess.run(pip_cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
except subprocess.CalledProcessError as e:
sys.exit("## Torchserve/Model archiver Installation Failed!")


if __name__ == "__main__":
Expand Down
17 changes: 12 additions & 5 deletions binaries/upload.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import glob
import os
import sys
import subprocess

# To help discover local modules
REPO_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..")
Expand Down Expand Up @@ -44,11 +45,17 @@ def upload_conda_packages(args, PACKAGES, CONDA_PACKAGES_PATH):
"tar.bz2"
):
print(f"Uploading to anaconda package: {name}")
anaconda_upload_command = f"anaconda upload {file_path} --force"
exit_code = os.system(anaconda_upload_command)
if exit_code != 0:
print(f"Anaconda package upload failed for package {name}")
return exit_code
anaconda_upload_command = ["anaconda", "upload", file_path, "--force"]

try:
subprocess.run(
anaconda_upload_command,
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
except subprocess.CalledProcessError as e:
return e.returncode
print(f"All packages uploaded to anaconda successfully")


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import argparse
import os
import shutil
import subprocess

MODEL_PTH_FILE = "resnet18-f37072fd.pth"
MODEL_STORE = "model_store"
Expand All @@ -15,7 +16,7 @@ def download_pth_file(output_file: str) -> None:
if not os.path.exists(output_file):
cmd = ["wget", " https://download.pytorch.org/models/resnet18-f37072fd.pth"]
print("Downloading resnet-18 pth file")
os.system(" ".join(cmd))
subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)


def create_mar():
Expand All @@ -42,9 +43,8 @@ def create_mar():
"--handler image_classifier",
"--force",
]

print(f"Archiving resnet-18 model into {MAR_FILE}")
os.system(" ".join(cmd))
subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)


def move_mar_file():
Expand Down
5 changes: 2 additions & 3 deletions examples/torchrec_dlrm/create_dlrm_mar.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,8 +2,7 @@
This script creates a DLRM model and packs it into a TorchServe mar file
"""

import os

import subprocess
import torch
from dlrm_factory import DLRMFactory

Expand Down Expand Up @@ -32,7 +31,7 @@ def main():
]

print("Archiving model into dlrm.mar")
os.system(" ".join(cmd))
subprocess.run(cmd, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
print("Done")


Expand Down
12 changes: 8 additions & 4 deletions test/pytest/test_distributed_inference_handler.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import os
import sys

import subprocess
import pytest
import test_utils

Expand Down Expand Up @@ -37,10 +37,14 @@ def test_large_model_inference():
)

try:
command = f"newman run -e {POSTMAN_ENV_FILE} {POSTMAN_COLLECTION_INFERENCE} -d {POSTMAN_LARGE_MODEL_INFERENCE_DATA_FILE} -r cli,htmlextra --reporter-htmlextra-export {ARTIFACTS_INFERENCE_DIR}/{REPORT_FILE} --verbose"
result = os.system(command)
command = [
"newman", "run", "-e", POSTMAN_ENV_FILE, POSTMAN_COLLECTION_INFERENCE,
"-d", POSTMAN_LARGE_MODEL_INFERENCE_DATA_FILE, "-r", "cli,htmlextra",
"--reporter-htmlextra-export", f"{ARTIFACTS_INFERENCE_DIR}/{REPORT_FILE}", "--verbose"
]
result = subprocess.run(command, check=True)
assert (
result == 0
result.returncode == 0
), "Error: Distributed inference failed, the exit code is not zero"
finally:
test_utils.stop_torchserve()
Expand Down
24 changes: 15 additions & 9 deletions test/pytest/test_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
import json
import logging
import os

import subprocess
import numpy as np
import pytest
import requests
Expand Down Expand Up @@ -314,10 +314,14 @@ def test_huggingface_bert_batch_inference():

# Make 2 curl requests in parallel with &
# curl --header \"X-Forwarded-For: 1.2.3.4\" won't work since you can't access local host anymore
response = os.popen(
f"curl http://127.0.0.1:8080/predictions/BERTSeqClassification -T {input_text} & curl http://127.0.0.1:8080/predictions/BERTSeqClassification -T {input_text}"
)
response = response.read()
curl_command = [
"curl", "http://127.0.0.1:8080/predictions/BERTSeqClassification", "-T", input_text
]
process1 = subprocess.Popen(curl_command, stdout=subprocess.PIPE)
process2 = subprocess.Popen(curl_command, stdout=subprocess.PIPE)
response1, _ = process1.communicate()
response2, _ = process2.communicate()
response = response1.decode() + response2.decode()

## Assert that 2 responses are returned from the same batch
assert response == "Not AcceptedNot Accepted"
Expand Down Expand Up @@ -360,7 +364,7 @@ def test_MMF_activity_recognition_model_register_and_inference_on_valid_model():

def test_huggingface_bert_model_parallel_inference():
number_of_gpus = torch.cuda.device_count()
check = os.popen(f"curl http://localhost:8081/models")
check = subprocess.run(["curl", "http://localhost:8081/models"], capture_output=True, text=True)
print(check)
if number_of_gpus > 1:
batch_size = 1
Expand All @@ -385,10 +389,12 @@ def test_huggingface_bert_model_parallel_inference():
"sample_text_captum_input.txt",
)

response = os.popen(
f"curl http://127.0.0.1:8080/predictions/Textgeneration -T {input_text}"
response = subprocess.run(
["curl", "http://127.0.0.1:8080/predictions/Textgeneration", "-T", input_text],
capture_output=True,
text=True
)
response = response.read()
response = response.stdout

assert (
"Bloomberg has decided to publish a new report on the global economy"
Expand Down
33 changes: 16 additions & 17 deletions test/pytest/test_sm_mme_requirements.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import os
import pathlib

import subprocess
import pytest
import requests
import test_utils
Expand Down Expand Up @@ -103,25 +103,24 @@ def test_oom_on_invoke():

# Make 8 curl requests in parallel with &
# Send multiple requests to make sure to hit OOM
processes = []
for i in range(10):
response = os.popen(
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} && "
f"curl http://127.0.0.1:8080/models/BERTSeqClassification/invoke -T {input_text} "
)
response = response.read()
for _ in range(8):
process = subprocess.Popen(
["curl", "http://127.0.0.1:8080/models/BERTSeqClassification/invoke", "-T", input_text],
stdout=subprocess.PIPE,
stderr=subprocess.PIPE
)
processes.append(process)
responses = []
for process in processes:
stdout, stderr = process.communicate()
responses.append(stdout.decode())

# If OOM is hit, we expect code 507 to be present in the response string
lines = response.split("\n")
output = ""
for line in lines:
if "code" in line:
line = line.strip()
output = line
for response in responses:
if '"code": 507,' in response:
output = '"code": 507,'
break
assert output == '"code": 507,', "OOM Error expected"
Loading

0 comments on commit 187ed53

Please sign in to comment.