File size: 1,519 Bytes
dcd82ce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | import json
import os
from pathlib import Path
import tensorflow as tf
from modelscan.modelscan import ModelScan
ROOT = Path(__file__).resolve().parent
SECRET_LINE = "TOKEN_FROM_MODEL_CHOSEN_FILE=local-fixture-secret"
def scan_model() -> int:
scanner = ModelScan()
scanner.scan(ROOT)
return len(scanner.issues.all_issues)
def run_model() -> str:
loader_cwd = ROOT / "loader_cwd"
victim_workspace = ROOT / "victim_workspace"
loader_cwd.mkdir(exist_ok=True)
victim_workspace.mkdir(exist_ok=True)
(victim_workspace / "secret.txt").write_text(SECRET_LINE + "\n")
old_cwd = Path.cwd()
os.chdir(loader_cwd)
try:
model = tf.saved_model.load(str(ROOT))
output = model.signatures["serving_default"](x=tf.constant(1.0))[
"leaked_line"
]
return output.numpy().decode("utf-8", errors="replace")
finally:
os.chdir(old_cwd)
def main() -> None:
leaked_line = run_model()
print(
json.dumps(
{
"tensorflow_version": tf.__version__,
"visible_gpu_count": len(tf.config.list_physical_devices("GPU")),
"cuda_visible_devices": os.environ.get("CUDA_VISIBLE_DEVICES"),
"modelscan_issue_count": scan_model(),
"leaked_line": leaked_line,
"leaked_expected_line": leaked_line == SECRET_LINE,
},
indent=2,
sort_keys=True,
)
)
if __name__ == "__main__":
main()
|