AIShield.Watchtower 🔍: Fortifying AI/ML Model and Notebook Security

In today's rapidly advancing landscape of machine learning and artificial intelligence (AI), ensuring the security of AI models has become an imperative. These models serve as the driving force behind decision-making in a myriad of applications. As such, safeguarding their integrity and protection against potential attacks remains a top priority. Unfortunately, the critical components of this AI ecosystem, namely Jupyter notebooks and models, are often overlooked in routine security assessments, leaving them vulnerable and appealing targets for potential attackers. Initial tests with open-source tools painted a grim picture, underscoring severe limitations and an urgent need for a bespoke solution. As AI/ML models continue to reshape industries and drive innovation, the significance of model security is paramount. However, amidst the excitement of AI's potential, the security of Jupyter notebooks and AI models has been consistently relegated to the shadows.

A compelling illustration of this security blind spot revolves around the widespread usage of Jupyter notebooks. Within these digital pages lie the blueprints, codes, and algorithms that breathe life into many of today's AI-driven products and services. However, like any other codebase, Jupyter notebooks are not immune to vulnerabilities that could inadvertently find their way into the final product. These vulnerabilities, if left unaddressed, can become gateways for cyber-attacks, risking not only the integrity of AI-driven products but also the security of user data.

AIShield Watchtower is designed to automate model and notebook discoveries and conduct comprehensive vulnerability scan. Its capabilities go beyond merely identifying all the models and notebooks within your repository. It assesses risks, such as hard-coded secrets, PIIs, outdated/unsafe libraries, model serialization attacks, custom unsafe operations etc.

AIShield Watchtower stands out with its capability to categorize scans into four distinct risk levels: "low," "medium," "high," and "critical." This classification equips organizations with the ability to tailor their security efforts to the level of risk detected. Its adaptive approach and meticulous risk categorization significantly bolster security efforts, fortifying them effectively. Watchtower's alignment with industry standards such as OWASP, MITRE, CWE, NIST AI RMF MAP functions further enhances its market standing by providing advanced security solutions.

Table of Contents

• Usage
  • Prerequisite
  • Inspect Jupyter Notebooks and Models Using
    • CLI
    • UI
    • UI-Docker
• Reports
• Features
• Benefits
• Limitation
• Contribution
• License

Usage

AIShield Watchtower can be used to inspect vulnerabilities in Jupyter notebooks, requirements.txt and AI/ML Models (Ref:Features).

For using AIShield Watchtower, clone Watchtower repo. Install prerequisites and scan your notebooks and AI/ML models. Some starting sample test files are available within the Watchtower repo.

Prerequisites

• For running Watchtower in CLI or UI version, python3 and pip should be installed in the host system.
• For running UI-Docker version, docker and docker-compose should be installed in the host system.

⚠️ Warning: As a security precaution when using the webapp or the Docker version, it is recommended to scan unknown/untrusted artifacts in a an isolated sandbox environment (network/docker/compute etc).

Cloning Watchtower repo

git clone https://github.com/bosch-aisecurity-aishield/watchtower.git
cd watchtower

Install Watchtower related dependency libraries using following commands:

Run this command to install dependency in Linux(Ubuntu) system:

./install.sh

Run this command to install dependency in windows system:

./install.bat

Inspect Jupyter Notebooks and ML/DL Models

Change directory to src.

cd src

Inspection of Jupyter Notebooks and ML/DL models can be done by any of the below methods:

CLI

To View the available options in CLI mode

python watchtower.py -h 

To scan the model files (.h5, .pb and .keras)

To scan the model files in .h5, .pb and .keras format, use the --scan_tf_model argument together with any of the Repo. This is currently supported for Linux OS Systems. For Example for github:

python watchtower.py --repo_type=<Repo-Type> --repo_url=<Repo-url> --branch_name=<Enter Branch Name> --scan_tf_model

To inspect artifacts present in Public Github Repo

python watchtower.py --repo_type=github --repo_url=<Enter Repo Url> --branch_name=<Enter Branch Name> --depth=<Enter a number>

NOTE: branch_name and depth parameters are optional. Default value of branch_name is main and default depth value is 1

To inspect artifacts present in Huggingface Repo

python watchtower.py --repo_type=huggingface --repo_url=<Enter Hugging Face Url>

To scan the artifacts present in AWS S3 bucket

python watchtower.py --repo_type=s3 --bucket_name="<Enter Bucket name>" --region="<Enter region of s3 bucket>" --aws_access_key_id="<Enter aws access key>" --aws_secret_access_key="<Enter aws secret key>"

To scan the artifacts present in local system

#Select Repo_type = file for scanning individual file
python watchtower.py --repo_type=file --path=<Enter path of File>
#Select Repo_type = file for scanning individual file
python watchtower.py --repo_type=folder --path=<Enter path of Folder>

UI

For using Watchtower UI, execute the following command:

python watchtower_webapp.py

UI-Docker

For using Watchtower UI with Docker, build docker image for watchtower and run watchtower image

cd watchtower
docker-compose build
docker-compose up

open browser and paste: http://localhost:5015/watchtower-aishield

On successful completion of scan, Watchtower vulnerability reports will be available in reports folder in Watchtower root folder.

For stopping and removing Watchtower image execute following

docker-compose down

Reports

On successful completion of the Watchtower scan, three reports will be generated in the following path :

• For CLI Mode - all three reports will be available inside the Watchtower src folder. Users may refer last line of the summary report in the console for the complete path of the reports location

• For UI Mode - all three will be reports will be available inside the Watchtower src folder. Users may refer to the success message on the UI to get the path of the reports location

• For UI-Docker Mode - all three reports will be available inside the Watchtower reports folder. Users may refer to the success message on the UI to get the path of the reports location

1. Summary Report - In summary report will provide information on number of model files and notebook files detected, Number of vulnerabilities detected and Count of those vulnerabilities mapped to Critical, High, Medium and low. Sample snippet of Summary Report:

{
       "Repository Type": "github",
       "Repository URL": "https://github.com/bosch-aisecurity-aishield/watchtower.git",
       "Total Number of Model Found": 17,
       "Total Number of Notebooks & Requirement files Found": 6,
       "Total Number of Model Scanned": 17,
       "Total Number of Notebooks & Requirement files Scanned": 6,
       "Total Notebook & Requirement files Vulnerabilities Found": {
              "Critical": 0,
              "High": 40,
              "Medium": 17,
              "Low": 1
       },
       "Total Model Vulnerabilities Found": {
              "Critical": 0,
              "High": 4,
              "Medium": 2,
              "Low": 8
       }
}

2. Severity Mapping Report - In this report, details of High, Medium and Low mapping to Model or Notebooks is reported. Sample snippet of Severity Mapping Report:

  {
         "type": "Hex High Entropy String",
         "filename": "repo_dir_1696827949/sample_test_files/sample_notebook_files/classification_notebook.py",
         "hashed_secret": "8d1e60a0b91ca2071dc4027b6f227990fb599d27",
         "is_verified": false,
         "line_number": 47,
         "vulnerability_severity": "High"
      },
      {
         "type": "Secret Keyword",
         "filename": "repo_dir_1696827949/sample_test_files/sample_notebook_files/classification_notebook.py",
         "hashed_secret": "8d1e60a0b91ca2071dc4027b6f227990fb599d27",
         "is_verified": false,
         "line_number": 47,
         "vulnerability_severity": "High"
   }

3. Detailed Report - In this report, Watchtower users will be able to find all logs generated during the vulnerability scan.

Features

• Model and Notebook Detection: Automatically recognizes AI/ML models and Notebooks within a provided repository.
• Scanning:Executes thorough scans of the models