From 1ef996d98749b715046c24388bbcfa9bb34f4784 Mon Sep 17 00:00:00 2001 From: anupamme Date: Fri, 30 May 2025 03:03:24 +0000 Subject: [PATCH 1/2] fix: python.lang.security.audit.eval-detected.eval-detected-examples-benchmarks-TRA-src-model.py --- examples/benchmarks/TRA/src/model.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/examples/benchmarks/TRA/src/model.py b/examples/benchmarks/TRA/src/model.py index ebafd6a521..a760c5ee2a 100644 --- a/examples/benchmarks/TRA/src/model.py +++ b/examples/benchmarks/TRA/src/model.py @@ -1,3 +1,4 @@ +import ast # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. @@ -51,7 +52,7 @@ def __init__( self.logger = get_module_logger("TRA") self.logger.info("TRA Model...") - self.model = eval(model_type)(**model_config).to(device) + self.model = ast.literal_eval(model_type)(**model_config).to(device) if model_init_state: self.model.load_state_dict(torch.load(model_init_state, map_location="cpu")["model"]) if freeze_model: From e0f54a789b65b4f26b6d6faf2dd9d6c9109601a0 Mon Sep 17 00:00:00 2001 From: anupamme Date: Mon, 3 Nov 2025 07:52:46 +0530 Subject: [PATCH 2/2] fixing the ast.literal_eval case --- examples/benchmarks/TRA/src/model.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/examples/benchmarks/TRA/src/model.py b/examples/benchmarks/TRA/src/model.py index a760c5ee2a..95fa1dace3 100644 --- a/examples/benchmarks/TRA/src/model.py +++ b/examples/benchmarks/TRA/src/model.py @@ -1,4 +1,3 @@ -import ast # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. @@ -52,7 +51,21 @@ def __init__( self.logger = get_module_logger("TRA") self.logger.info("TRA Model...") - self.model = ast.literal_eval(model_type)(**model_config).to(device) + # Secure model registry - whitelist of allowed model classes + # This prevents arbitrary code execution while allowing dynamic model selection + model_registry = { + "LSTM": LSTM, + "Transformer": Transformer, + } + + if model_type not in model_registry: + raise ValueError( + f"Unknown model_type: '{model_type}'. " + f"Supported types: {list(model_registry.keys())}" + ) + + model_class = model_registry[model_type] + self.model = model_class(**model_config).to(device) if model_init_state: self.model.load_state_dict(torch.load(model_init_state, map_location="cpu")["model"]) if freeze_model: