OsOperations::set_env is corrected

src/local_ops.py

@@ -398,9 +398,19 @@ def is_executable(self, file):
         assert stat.S_IXUSR != 0
         return (os.stat(file).st_mode & stat.S_IXUSR) == stat.S_IXUSR
 
-    def set_env(self, var_name, var_val):
-        # Check if the directory is already in PATH
-        os.environ[var_name] = var_val
+    def set_env(
+        self,
+        var_name: str,
+        var_val: typing.Optional[str],
+    ) -> None:
+        assert type(var_name) is str
+        assert var_val is None or type(var_val) is str
+        assert var_name != ""
+
+        if var_val is None:
+            os.environ.pop(var_name, None)
+        else:
+            os.environ[var_name] = var_val
         return
 
     def get_name(self):

src/os_ops.py

@@ -135,8 +135,14 @@ def is_executable(self, file):
         # Check if the file is executable
         raise NotImplementedError()
 
-    def set_env(self, var_name, var_val):
-        # Check if the directory is already in PATH
+    def set_env(
+        self,
+        var_name: str,
+        var_val: typing.Optional[str],
+    ) -> None:
+        assert type(var_name) is str
+        assert var_val is None or type(var_val) is str
+        assert var_name != ""
         raise NotImplementedError()
 
     def get_user(self):

src/remote_ops.py

@@ -13,6 +13,7 @@
 import time
 import datetime
 import shlex
+import threading
 
 from .exceptions import ExecUtilException
 from .exceptions import InvalidOperationException
@@ -50,6 +51,8 @@ def cmdline(self):
 class RemoteOperations(OsOperations):
     _C_EOL = "\n"
 
+    T_ENVS = typing.Dict[str, typing.Optional[str]]
+
     #
     # Target system is Linux only.
     #
@@ -61,6 +64,7 @@ class RemoteOperations(OsOperations):
     _ssh_key: typing.Optional[str]
     _username: typing.Optional[str]
     _ssh_cmd: typing.List[str]
+    _remote_env: T_ENVS
 
     def __init__(self, conn_params: ConnectionParams):
         if conn_params is None:
@@ -98,6 +102,9 @@ def __init__(self, conn_params: ConnectionParams):
             self._ssh_cmd += [conn_params.username + "@" + self._host]
         else:
             self._ssh_cmd += [self._host]
+
+        self._remote_env = dict()
+        self._remote_env_guard = threading.Lock()
         return
 
     @property
@@ -155,7 +162,7 @@ def exec_command(
         get_process=None,
         timeout=None,
         ignore_errors=False,
-        exec_env: typing.Optional[dict] = None,
+        exec_env: typing.Optional[T_ENVS] = None,
         cwd: typing.Optional[str] = None
     ) -> OsOperations.T_EXEC_COMMAND_RESULT:
         """
@@ -187,7 +194,22 @@ def exec_command(
             assert type(cwd) is str
             cmds.append(__class__._build_cmdline(["cd", cwd]))
 
-        cmds.append(__class__._build_cmdline(cmd, exec_env))
+        assert self._remote_env_guard is not None
+        assert type(self._remote_env) is dict
+
+        exec_env2: typing.Optional[__class__.T_ENVS] = None
+
+        with self._remote_env_guard:
+            if len(self._remote_env) > 0:
+                exec_env2 = self._remote_env.copy()
+                assert type(exec_env2) is dict
+
+        if exec_env2 is None:
+            exec_env2 = exec_env
+        elif exec_env is not None:
+            exec_env2.update(exec_env)
+
+        cmds.append(__class__._build_cmdline(cmd, exec_env2))
 
         assert len(cmds) >= 1
 
@@ -373,14 +395,24 @@ def is_executable(self, file):
             out=output
         )
 
-    def set_env(self, var_name: str, var_val: str):
+    def set_env(
+        self,
+        var_name: str,
+        var_val: typing.Optional[str],
+    ) -> None:
         """
         Set the value of an environment variable.
         Args:
         - var_name (str): The name of the environment variable.
-        - var_val (str): The value to be set for the environment variable.
+        - var_val (str|None): The value to be set for the environment variable.
         """
-        return self.exec_command("export {}={}".format(var_name, var_val))
+        assert type(var_name) is str
+        assert var_val is None or type(var_val) is str
+        assert var_name != ""
+
+        with self._remote_env_guard:
+            self._remote_env[var_name] = var_val
+        return
 
     def get_name(self):
         cmd = 'python3 -c "import os; print(os.name)"'

tests/test_os_ops_common.py

@@ -3431,6 +3431,118 @@ def test_environ(
 
         return
 
+    def test_set_env__persistence(
+        self,
+        os_ops_descr: OsOpsDescr,
+    ):
+        assert type(os_ops_descr) is OsOpsDescr
+        os_ops = os_ops_descr.os_ops
+        assert isinstance(os_ops, OsOperations)
+
+        var_name = "TEST_SET_ENV_MAGIC_VAR"
+        var_value = "Cokanum_Detected_123"
+
+        # Step 1: Set the variable
+        os_ops.set_env(var_name, var_value)
+
+        # Step 2: In a SEPARATE command, we try to read it using our new environ()
+        # The old remote_ops is guaranteed to return None and crash!
+        fetched_value = os_ops.environ(var_name)
+
+        assert fetched_value == var_value, "Env variable persistence failed! Got: {}".format(fetched_value)
+        logging.info("SUCCESS. Environment variable persistence verified across commands.")
+
+        printenv = os_ops.find_executable("printenv")
+        assert type(printenv) is str
+        assert printenv != ""
+
+        cmd1 = [printenv, var_name]
+
+        exec_r = os_ops.exec_command(
+            cmd1,
+            encoding="utf-8",
+        )
+        assert type(exec_r) is str
+        exec_r = exec_r.rstrip()
+        assert fetched_value == var_value
+
+        exec_r = os_ops.exec_command(
+            cmd1,
+            encoding="utf-8",
+            exec_env={
+                var_name: "ABC",
+            }
+        )
+        assert type(exec_r) is str
+        exec_r = exec_r.rstrip()
+        assert exec_r == "ABC"
+
+        exec_r = os_ops.exec_command(
+            cmd1,
+            encoding="utf-8",
+            exec_env={
+                var_name: None,
+            },
+            ignore_errors=True,
+            verbose=True,
+        )
+        assert type(exec_r) is tuple
+        assert len(exec_r) == 3
+        assert exec_r[0] == 1
+
+        # ----------------
+        x = os_ops.environ("PATH")
+
+        try:
+            os_ops.set_env("PATH", None)
+
+            cmd2 = [printenv, "PATH"]
+
+            exec_r = os_ops.exec_command(
+                cmd2,
+                encoding="utf-8",
+                ignore_errors=True,
+                verbose=True,
+            )
+            assert type(exec_r) is tuple
+            assert len(exec_r) == 3
+            assert exec_r[0] == 1
+        finally:
+            os_ops.set_env("PATH", x)
+            assert os_ops.environ("PATH") == x
+
+        return
+
+    def test_set_env__evil(
+        self,
+        os_ops_descr: OsOpsDescr,
+    ):
+        assert type(os_ops_descr) is OsOpsDescr
+        os_ops = os_ops_descr.os_ops
+        assert isinstance(os_ops, OsOperations)
+
+        # --- INJECTION TEST VIA VARIABLE VALUE ---
+        # We stuff a hellish mixture of quotes,
+        # ampersands, and destructive shell commands into the variable value.
+        evil_value = "clean_val' && echo 'HACKED' && rm -rf / ; \" double_q"
+        evil_var = "TEST_EVIL_EXPORT_VAR"
+
+        # Step 1: Store this crazy value in the state
+        os_ops.set_env(evil_var, evil_value)
+
+        # Step 2: Call any harmless command (e.g., pwd or printenv)
+        # If quoting inside exec_command fails, the shell will execute the "echo 'HACKED'" chunk
+        # or crash with a quoting syntax error.
+        try:
+            fetched_evil = os_ops.environ(evil_var)
+            # printenv should return the string EXACTLY, without distortion and code execution
+            assert fetched_evil == evil_value, f"Evil env corruption! Got: {fetched_evil}"
+            logging.info("SUCCESS. Remote export variable value is completely bulletproof against shell injections.")
+        finally:
+            # Be sure to clean up after yourself
+            os_ops.set_env(evil_var, None)
+        return
+
     @staticmethod
     def helper__bug_check__unknown_os_ops_type(
         os_ops: OsOperations,