From 0226c8f55bf77a8e82469a41d77d53e7a082ce2a Mon Sep 17 00:00:00 2001
From: Andrew Konstantinov <105389353+execveat@users.noreply.github.com>
Date: Tue, 19 Dec 2023 12:09:14 +0100
Subject: [PATCH] TorchServe testbed

---
 torchserve/exposed_ui/README.md           | 118 ++++++++++++++++++++++
 torchserve/exposed_ui/docker-compose.yaml |  76 ++++++++++++++
 torchserve/exposed_ui/tsunami/Dockerfile  |  14 +++
 torchserve/exposed_ui/tsunami/tsunami.sh  |  68 +++++++++++++
 4 files changed, 276 insertions(+)
 create mode 100644 torchserve/exposed_ui/README.md
 create mode 100644 torchserve/exposed_ui/docker-compose.yaml
 create mode 100644 torchserve/exposed_ui/tsunami/Dockerfile
 create mode 100755 torchserve/exposed_ui/tsunami/tsunami.sh

diff --git a/torchserve/exposed_ui/README.md b/torchserve/exposed_ui/README.md
new file mode 100644
index 00000000..e54bf311
--- /dev/null
+++ b/torchserve/exposed_ui/README.md
@@ -0,0 +1,118 @@
+# TorchServe Vulnerability Testbed
+## Overview
+This testbed provides a Docker Compose setup to test the TorchServe Management API Detection Plugin against various versions of TorchServe, demonstrating different vulnerability scenarios. It includes four TorchServe containers with varying configurations to simulate different security postures.
+
+## Prerequisites
+Before starting, ensure Docker and Docker Compose are installed on your system.
+
+## TorchServe Containers
+The testbed includes the following TorchServe containers:
+
+- **torchserve-081:** Version 0.8.1, vulnerable to the ShellTorch attack.
+- **torchserve-082:** Version 0.8.2, patched for ShellTorch but still vulnerable due to default `allowed_urls`.
+- **torchserve-safe:** Version 0.8.2 with `allowed_urls` correctly restricted, representing a secure setup.
+- **torchserve-latest:** The latest version (currently 0.9.0), with default settings and vulnerabilities similar to torchserve-082.
+
+## Setup Instructions
+
+### Building Tsunami Docker Image
+Clone the Tsunami repository and build the Docker image:
+
+```bash
+git clone git@github.com:google/tsunami-security-scanner.git
+cd tsunami-security-scanner
+docker build -t tsunami .
+```
+
+### Custom Tsunami Plugins
+1. Adding custom plugins
+    - Compile your custom plugins into JAR files.
+    - Place your custom plugin JAR files into the `tsunami/custom_plugins/` directory.
+2. Rebuilding the Docker image
+    - Each time add or update plugins, you need to rebuild the Docker image.
+    - Run the following command to rebuild the Docker image:
+        ```bash
+        docker compose build tsunami
+        ```
+    - `docker compose up --build` will also work
+3. Configure which plugins to run
+    - `USE_CUSTOM_PLUGINS`: set this environment variable to `true` to enable custom plugins.
+    - `USE_DEFAULT_PLUGINS`: set this environment variable to `false` to disable default plugins.
+    - both can be enabled at the same time, or just one of them.
+
+### Starting TorchServe Services
+To start all the TorchServe services in detached mode:
+
+```bash
+docker compose up -d
+```
+
+To start a specific TorchServe service, for example, `torchserve-081`:
+
+```bash
+docker compose up -d torchserve-081
+```
+
+### Stopping Services
+To stop all services without removing containers:
+
+```bash
+docker compose stop
+```
+
+To stop and remove containers:
+
+```bash
+docker compose down
+```
+
+To stop, remove containers, and volumes:
+
+```bash
+docker compose down -v
+```
+
+### Accessing Logs
+To follow the log output of the containers:
+
+```bash
+docker compose logs -f
+```
+
+## Using Tsunami for Testing
+Getting a Shell in Tsunami Container
+
+```bash
+docker compose run --rm --entrypoint /bin/bash tsunami
+```
+
+Running Tsunami Commands:
+
+```bash
+# Regular scan
+docker compose run --rm tsunami --uri-target=http://torchserver-081:8081
+# Custom options to force Static mode
+docker compose run --rm tsunami --uri-target=http://torchserve-safe:8081 \
+    --torchserve-management-api-mode=static \
+    --torchserve-management-api-model-static-url=http://e4d14c157244:8000/model.mar
+```
+
+Note that the `--name=tsunami` option is required for Local mode scan to work, as otherwise the container name will be randomly generated and not match the hostname in the scan results:
+
+```bash
+docker compose run --rm --name=tsunami tsunami --uri-target=http://torchserve-081:8081 \
+    --torchserve-management-api-mode=local --torchserve-management-api-local-bind-host=tsunami \
+    --torchserve-management-api-local-bind-port=1234 \
+    --torchserve-management-api-local-accessible-url=http://tsunami:1234
+```
+
+To only output JSON and nicely format the output with `jq`:
+
+```bash
+docker compose run --rm tsunami --uri-target=http://torchserver-081:8081 --json | jq
+```
+For a concise output:
+
+```bash
+docker compose run --rm tsunami --uri-target=http://torchserver-081:8081 --short
+```
diff --git a/torchserve/exposed_ui/docker-compose.yaml b/torchserve/exposed_ui/docker-compose.yaml
new file mode 100644
index 00000000..316595c2
--- /dev/null
+++ b/torchserve/exposed_ui/docker-compose.yaml
@@ -0,0 +1,76 @@
+# TorchServe testbed for detecting exposed management interface
+#
+# The risks of exposing TorchServe management interface to the Internet
+# were initially demonstrated in the "ShellTorch" attack
+# (https://www.oligo.security/shelltorch) that allowed arbitrary code
+# execution by chaining CVE-2023-43654 and CVE-2022-1471.
+#
+# Although TorchServe versions 0.8.2 and later are not vulnerable to
+# insecure deserialization (CVE-2022-1471), they still allow arbitrary
+# code execution via the management interface, if the `allowed_urls`
+# configuration option is not sufficiently restricted.
+#
+# This testbed builds three TorchServe containers:
+#
+#   - torchserve-081: 0.8.1, vulnerable to ShellTorch attack
+#   - torchserve-082: 0.8.2, with ShellTorch mitigations, but still vulnerable
+#   - torchserve-safe: 0.8.2, with `allowed_urls` correctly restricted
+#   - torchserve-latest: latest version, with default settings (currently 0.9.0, vulnerable)
+
+# Shared configuration for TorchServe containers
+x-torchserve-common: &torchserve-common
+  platform: linux/amd64 # TorchServe doesn't support ARM at the moment
+  #ports:
+  #  - "8080"
+  #  - "8081"
+  #  - "8082"
+  #  - "7070"
+  #  - "7071"
+  networks:
+    - torchserve-network
+
+services:
+  # 0.8.1 is the last version that is vulnerable to ShellTorch attack
+  torchserve-081:
+    container_name: torchserve-081
+    <<: *torchserve-common  # Inherits common settings
+    image: pytorch/torchserve:0.8.1-cpu
+
+  # 0.8.2 with default `allowed_urls` configuration - still vulnerable
+  torchserve-082:
+    container_name: torchserve-082
+    <<: *torchserve-common
+    image: pytorch/torchserve:0.8.2-cpu
+
+  # 0.8.2 with restricted `allowed_urls` configuration - safe
+  torchserve-safe:
+    container_name: torchserve-safe
+    <<: *torchserve-common
+    image: pytorch/torchserve:0.8.2-cpu
+    entrypoint: >
+      /bin/sh -c '
+        (cat config.properties; echo "allowed_urls = http://localhost/*") > safe_config.properties;
+        torchserve --start --ts-config safe_config.properties;
+        tail -f /dev/null'
+
+  # Latest version with default `allowed_urls` configuration (atm 0.9.0, still vulnerable)
+  torchserve-latest:
+    container_name: torchserve-latest
+    <<: *torchserve-common  # Inherits common settings
+    image: pytorch/torchserve:latest-cpu  # Latest version tag
+
+  # Tsunami container
+  tsunami:
+    container_name: tsunami
+    build: ./tsunami/
+    networks:
+      - torchserve-network
+    environment:
+      - USE_CUSTOM_PLUGINS  # Include custom plugins; set to 'true' to include
+      - USE_DEFAULT_PLUGINS # Include default plugins; set to 'false' to exclude
+    logging:
+      driver: none
+
+networks:
+  torchserve-network:
+    driver: bridge
diff --git a/torchserve/exposed_ui/tsunami/Dockerfile b/torchserve/exposed_ui/tsunami/Dockerfile
new file mode 100644
index 00000000..50d257b2
--- /dev/null
+++ b/torchserve/exposed_ui/tsunami/Dockerfile
@@ -0,0 +1,14 @@
+# Start from the existing Tsunami image
+FROM tsunami
+
+RUN apt update && apt upgrade -y && apt install -y \
+    python3 nmap ncat inetutils-ping curl wget less jq vim zip jq
+
+# Copy custom plugins into the container
+COPY custom_plugins/ /usr/tsunami/custom_plugins
+
+# Copy a custom entrypoint script into the container
+COPY tsunami.sh /tsunami.sh
+
+# Set the custom entrypoint script as the entrypoint
+ENTRYPOINT ["/tsunami.sh"]
diff --git a/torchserve/exposed_ui/tsunami/tsunami.sh b/torchserve/exposed_ui/tsunami/tsunami.sh
new file mode 100755
index 00000000..2f75d07d
--- /dev/null
+++ b/torchserve/exposed_ui/tsunami/tsunami.sh
@@ -0,0 +1,68 @@
+#!/bin/bash
+
+# No arguments, exit
+if [ -z "$*" ]; then
+  echo "No command specified; exiting Tsunami CLI"
+  exit;
+fi
+
+# If the --json flag is specified (any position), set the environment variable and remove the flag
+args=()
+for arg in "$@"; do
+    if [ "$arg" = "--json" ]; then
+        export OUTPUT_JSON="json"
+    elif [ "$arg" = "--short" ]; then
+        export OUTPUT_JSON="short"
+    else
+        args+=("$arg")
+    fi
+done
+
+setup_classpath() {
+    classpath="tsunami.jar"
+    if [ "${USE_CUSTOM_PLUGINS}" = "true" ]; then
+        classpath="${classpath}:custom_plugins/*.jar"
+    fi
+
+    if [ -z "${USE_DEFAULT_PLUGINS}" ] || [ "${USE_DEFAULT_PLUGINS}" = "true" ]; then
+        classpath="${classpath}:plugins/*.jar"
+    fi
+    echo "${classpath}"
+}
+
+# If --json was specified, only output JSON
+if [[ -n "${OUTPUT_JSON}" ]]; then
+    echo "Running Tsunami in JSON output mode ($(pwd))" >&2
+    echo "Classpath: $(setup_classpath)" >&2
+
+    java -cp "$(setup_classpath)" -Dtsunami-config.location=tsunami.yaml \
+        com.google.tsunami.main.cli.TsunamiCli \
+        --scan-results-local-output-format=JSON \
+        --scan-results-local-output-filename=output.json \
+        "${args[@]}" >./original_output.txt 2>&1 &
+    tsunami_pid=$!
+
+    # Wait a few seonds and check if Tsunami is still running
+    sleep 10
+    if ! kill -0 $tsunami_pid 2>/dev/null; then
+        echo "Tsunami exited early; here's the output:" >&2
+        cat original_output.txt >&2
+        exit 1
+    fi
+
+    wait $tsunami_pid; sleep 2
+    echo "Tsunami exited" >&2
+    if [[ -f output.json ]]; then
+        if [[ ${OUTPUT_JSON} == "short" ]]; then
+            cat output.json | jq '.scanFindings[] | {targetInfo, vulnerability, networkService: .networkService | {networkEndpoint, transportProtocol, serviceName}}'
+        else
+            cat output.json
+        fi
+    else
+        echo "No output.json file found" >&2
+    fi
+else
+    # Otherwise just run the command with original output style
+    java -cp "$(setup_classpath)" -Dtsunami-config.location=tsunami.yaml \
+        com.google.tsunami.main.cli.TsunamiCli "${args[@]}" 2>&1
+fi