feat(connectors): allow HiddenValue in _env (#1853)

Author: Jesse-11

src/pyinfra/api/arguments.py

@@ -16,13 +16,16 @@
 from pyinfra.api.exceptions import ArgumentTypeError
 from pyinfra.api.util import raise_if_bad_type
 from pyinfra.context import ctx_config
+from pyinfra.api.hiddenvalue import HiddenValue
 
 if TYPE_CHECKING:
     from pyinfra.api import Config, Host, State
 
 T = TypeVar("T")
 default_sentinel = object()
 
+EnvValue = str | HiddenValue
+
 
 class ArgumentMeta(Generic[T]):
     description: str
@@ -62,7 +65,7 @@ class ConnectorArguments(TypedDict, total=False):
     # Shell arguments
     _shell_executable: str
     _chdir: str
-    _env: Mapping[str, str]
+    _env: Mapping[str, EnvValue]
 
     # Connector control (outside of command generation)
     _success_exit_codes: Iterable[int]
@@ -79,10 +82,13 @@ class ConnectorArguments(TypedDict, total=False):
     _temp_dir: str
 
 
-def generate_env(config: Config, value: dict) -> dict:
-    env = {key: os.environ[key] for key in config.INHERIT_ENV if key in os.environ}
+def generate_env(config: Config, value: Mapping[str, EnvValue] | None) -> dict[str, EnvValue]:
+    env: dict[str, EnvValue] = {
+        key: os.environ[key] for key in config.INHERIT_ENV if key in os.environ
+    }
     env.update(config.ENV)
-    env.update(value)
+    if value is not None:
+        env.update(value)
     return env
 
 

src/pyinfra/api/arguments_typed.py

@@ -5,11 +5,15 @@
 
 from typing_extensions import ParamSpec, Protocol
 
+from pyinfra.api.hiddenvalue import HiddenValue
+
 if TYPE_CHECKING:
     from pyinfra.api.operation import OperationMeta
 
 P = ParamSpec("P")
 
+EnvValue = str | HiddenValue
+
 
 # Unfortunately we have to re-type out all of the global arguments here because
 # Python typing doesn't (yet) support merging kwargs. This acts as the operation
@@ -43,7 +47,7 @@ def __call__(
         # Shell arguments
         _shell_executable: None | str = None,
         _chdir: None | str = None,
-        _env: None | Mapping[str, str] = None,
+        _env: None | Mapping[str, EnvValue] = None,
         # Connector control
         _success_exit_codes: Iterable[int] = (0,),
         _timeout: None | int = None,

src/pyinfra/connectors/util.py

@@ -428,7 +428,7 @@ def make_unix_command(
             # Quote the whole `key=value` pair so arbitrary values cannot break
             # out into additional shell tokens. Invalid identifiers in `key` will
             # fail safely when the shell rejects the resulting `export` statement.
-            env_bits.append(QuoteString(f"{key}={value}"))
+            env_bits.append(QuoteString(StringCommand(key, value, _separator="=")))
         env_bits.append("&&")
         env_bits.append(command)
         command = StringCommand(*env_bits)

tests/test_connectors/test_util.py

@@ -1,7 +1,7 @@
 from unittest import TestCase
 from unittest.mock import MagicMock, patch
 
-from pyinfra.api import Config, State
+from pyinfra.api import Config, State, HiddenValue
 from pyinfra.connectors.util import (
     CommandOutput,
     OutputLine,
@@ -147,6 +147,40 @@ def test_command_env_injection_attempt(self):
             == "sh -c 'export '\"'\"'KEY=value\"; id; echo \"'\"'\"' && uptime'"
         )
 
+    def test_command_env_hidden_value(self):
+        command = make_unix_command(
+            "uptime",
+            _env={
+                "KEY": HiddenValue("super-secret"),
+            },
+        )
+
+        raw = command.get_raw_value()
+        masked = command.get_masked_value()
+        assert "super-secret" in raw
+        assert "super-secret" not in masked
+        assert "*MASKED*" in masked
+        assert "KEY=*MASKED*" in masked
+
+    def test_command_env_hidden_value_with_shell_chars(self):
+        command = make_unix_command(
+            "uptime",
+            _env={
+                "KEY": HiddenValue("abc; id; echo bad"),
+            },
+        )
+
+        raw = command.get_raw_value()
+        masked = command.get_masked_value()
+
+        assert "abc; id; echo bad" in raw
+        assert "abc; id; echo bad" not in masked
+        assert "*MASKED*" in masked
+        assert "KEY=*MASKED*" in masked
+
+        # The raw command should still quote the assignment as one shell token.
+        assert "'\"'\"'KEY=abc; id; echo bad'\"'\"'" in raw
+
     def test_command_chdir(self):
         command = make_unix_command("uptime", _chdir="/opt/somedir")
         assert command.get_raw_value() == "sh -c 'cd /opt/somedir && uptime'"