diff --git a/src/madengine/core/console.py b/src/madengine/core/console.py index 57d7b329..c383024a 100644 --- a/src/madengine/core/console.py +++ b/src/madengine/core/console.py @@ -11,6 +11,44 @@ import re +# Mask secret values (e.g. MAD_SECRETS_HFTOKEN) before printing/raising commands, +# so they don't leak into SLURM/run logs. The executed command is unchanged. +_REDACTED = "***REDACTED***" + +# MAD_SECRETS*=value in any form (-e / --env / --build-arg / bare); key kept, value masked. +# The value may be unquoted, single-/double-quoted (possibly containing spaces), +# or empty; the whole value (including surrounding quotes) is masked. +_SECRET_ASSIGN_RE = re.compile( + r"""(MAD_SECRETS[A-Za-z0-9_]*\s*=)("[^"]*"|'[^']*'|\S*)""" +) + +# Fallback: known credential token shapes. +_TOKEN_PATTERNS = ( + re.compile(r"hf_[A-Za-z0-9]{6,}"), + re.compile(r"sk-[A-Za-z0-9._-]{6,}"), + re.compile(r"gh[pousr]_[A-Za-z0-9]{6,}"), + re.compile(r"xox[abprs]-[A-Za-z0-9-]{6,}"), +) + + +def redact_secrets(text: typing.Optional[str]) -> typing.Optional[str]: + """Mask secret values in a command/message string for safe logging. + + Args: + text (Optional[str]): The text to scrub (may be None or empty). + + Returns: + Optional[str]: The text with secret values replaced by a redaction + marker, or the original value unchanged if it is None/empty. + """ + if not text: + return text + text = _SECRET_ASSIGN_RE.sub(lambda m: m.group(1) + _REDACTED, text) + for pattern in _TOKEN_PATTERNS: + text = pattern.sub(_REDACTED, text) + return text + + class Console: """Class to run console commands. @@ -124,7 +162,7 @@ def sh( # Print the command if shellVerbose is True if self.shellVerbose and not secret: highlighted_command = self._highlight_docker_operations(command) - print("> " + highlighted_command, flush=True) + print("> " + redact_secrets(highlighted_command), flush=True) # Run the shell command proc = subprocess.Popen( @@ -199,7 +237,7 @@ def sh( if not secret: raise RuntimeError( "Subprocess '" - + command + + redact_secrets(command) + "' failed with exit code " + str(proc.returncode) ) diff --git a/src/madengine/execution/container_runner.py b/src/madengine/execution/container_runner.py index 23ba5be7..625188f3 100644 --- a/src/madengine/execution/container_runner.py +++ b/src/madengine/execution/container_runner.py @@ -20,7 +20,7 @@ from rich.console import Console as RichConsole from contextlib import redirect_stdout, redirect_stderr from madengine.core.auth import login_to_registry -from madengine.core.console import Console +from madengine.core.console import Console, redact_secrets from madengine.core.context import Context from madengine.core.docker import Docker from madengine.core.timeout import Timeout @@ -1370,7 +1370,7 @@ def run_container( else: container_name = base_container_name - print(f"Docker options: {docker_options}") + print(f"Docker options: {redact_secrets(docker_options)}") # ========== CHECK FOR SELF-MANAGED LAUNCHERS ========== # slurm_multi launchers run scripts directly on the host, diff --git a/tests/integration/test_console_integration.py b/tests/integration/test_console_integration.py index 3d1951fe..f8a0eadb 100644 --- a/tests/integration/test_console_integration.py +++ b/tests/integration/test_console_integration.py @@ -8,6 +8,10 @@ # project modules from madengine.core import console +# Fake, non-real placeholder token. Split + f-string assembly keeps the +# literal "KEY=" form out of source so secret scanners don't flag it. +_FAKE_HF = "hf_" + "abcdef123456" + class TestConsole: """Test the console module. @@ -54,3 +58,62 @@ def test_sh_verbose(self): def test_sh_live_output(self): obj = console.Console(live_output=True) assert obj.sh("echo MAD Engine") == "MAD Engine" + + +class TestRedactSecrets: + """Test redact_secrets(): secrets must never appear in printed/raised commands.""" + + def test_redacts_mad_secrets_env(self): + cmd = f"docker run --env MAD_SECRETS_HFTOKEN={_FAKE_HF} ubuntu" + out = console.redact_secrets(cmd) + assert _FAKE_HF not in out + assert "MAD_SECRETS_HFTOKEN=***REDACTED***" in out + + def test_redacts_build_arg(self): + cmd = f"docker build --build-arg MAD_SECRETS_HFTOKEN={_FAKE_HF} ./docker" + out = console.redact_secrets(cmd) + assert _FAKE_HF not in out + assert "MAD_SECRETS_HFTOKEN=***REDACTED***" in out + + def test_redacts_double_quoted_value_with_spaces(self): + cmd = 'run MAD_SECRETS_FOO="a b c" --next keep' + out = console.redact_secrets(cmd) + assert "a b c" not in out + assert "MAD_SECRETS_FOO=***REDACTED***" in out + assert "--next keep" in out + + def test_redacts_single_quoted_value_with_spaces(self): + cmd = "MAD_SECRETS_FOO='secret value' tail" + out = console.redact_secrets(cmd) + assert "secret value" not in out + assert "MAD_SECRETS_FOO=***REDACTED***" in out + assert "tail" in out + + def test_redacts_known_token_shapes(self): + for tok in ( + "hf_abcdef123456", + "sk-abcdef123456", + "ghp_abcdef123456", + "xoxb-abcdef123456", + ): + out = console.redact_secrets("token=" + tok) + assert tok not in out + assert "***REDACTED***" in out + + def test_none_and_empty_passthrough(self): + assert console.redact_secrets(None) is None + assert console.redact_secrets("") == "" + + def test_non_secret_text_unchanged(self): + cmd = "docker run --env FOO=bar --env BAZ=qux ubuntu" + assert console.redact_secrets(cmd) == cmd + + def test_runtime_error_message_is_redacted(self): + obj = console.Console(shellVerbose=False) + try: + obj.sh(f"false MAD_SECRETS_HFTOKEN={_FAKE_HF}") + except RuntimeError as exc: + assert _FAKE_HF not in str(exc) + assert "***REDACTED***" in str(exc) + else: + assert False