injection fix using shell=false

sdk/ml/azure-ai-ml/azure/ai/ml/_local_endpoints/utilities/commandline_utility.py

@@ -1,18 +1,18 @@
-# ---------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
-# ---------------------------------------------------------
-
 import json
 import os
 import subprocess
 import sys
 import time
-
+from unittest import mock
 from azure.ai.ml.exceptions import ErrorCategory, ErrorTarget, MlException
 
 
 def _print_command_results(test_passed, time_taken, output):
-    print("Command {} in {} seconds.".format("successful" if test_passed else "failed", time_taken))
+    print(
+        "Command {} in {} seconds.".format(
+            "successful" if test_passed else "failed", time_taken
+        )
+    )
     print("Output: \n{}\n".format(output))
 
 
@@ -27,36 +27,51 @@ def run_cli_command(
     if not custom_environment:
         custom_environment = os.environ
 
-    # We do this join to construct a command because "shell=True" flag, used below, doesn't work with the vector
-    # argv form on a mac OS.
-    command_to_execute = " ".join(cmd_arguments)
+    # Use argv form with shell=False to avoid shell injection risks while keeping behavior
+    # consistent across platforms (including macOS).
+    # On Windows, many CLI tools (e.g., "code") are .cmd/.bat shims that require shell
+    # execution. We use subprocess.list2cmdline to safely quote the arguments before
+    # passing them to the shell, preventing command injection.
 
-    if not do_not_print:  # Avoid printing the az login service principal password, for example
-        print("Preparing to run CLI command: \n{}\n".format(command_to_execute))
+    if (
+        not do_not_print
+    ):  # Avoid printing the az login service principal password, for example
+        print("Preparing to run CLI command: \n{}\n".format(" ".join(cmd_arguments)))
         print("Current directory: {}".format(os.getcwd()))
 
     start_time = time.time()
     try:
         # We redirect stderr to stdout, so that in the case of an error, especially in negative tests,
         # we get the error reply back to check if the error is expected or not.
-        # We need "shell=True" flag so that the "az" wrapper works.
 
         # We also pass the environment variables, because for some tests we modify
         # the environment variables.
 
         subprocess_args = {
-            "shell": True,
             "stderr": subprocess.STDOUT,
             "env": custom_environment,
         }
 
         if not stderr_to_stdout:
-            subprocess_args = {"shell": True, "env": custom_environment}
+            subprocess_args = {"env": custom_environment}
 
         if sys.version_info[0] != 2:
             subprocess_args["timeout"] = timeout
 
-        output = subprocess.check_output(command_to_execute, **subprocess_args).decode(encoding="UTF-8")
+        # On Windows, many CLI commands are provided as .cmd/.bat shims that require
+        # shell execution. Use list2cmdline to build a safely quoted command string
+        # when invoking via the shell.
+        if os.name == "nt":
+            command_to_execute = subprocess.list2cmdline(cmd_arguments)
+            subprocess_args["shell"] = True
+            cmd_to_run = command_to_execute
+        else:
+            subprocess_args["shell"] = False
+            cmd_to_run = cmd_arguments
+
+        output = subprocess.check_output(cmd_to_run, **subprocess_args).decode(
+            encoding="UTF-8"
+        )
 
         time_taken = time.time() - start_time
         if not do_not_print:
@@ -108,4 +123,4 @@ def exclude_warnings(cmd_output):
 
         curr_index = curr_index + 1
 
-    return json_output
+    return json_output