Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 40 additions & 2 deletions src/madengine/core/console.py
Original file line number Diff line number Diff line change
Expand Up @@ -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*)"""
)
Comment on lines +21 to +23

# 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
Comment thread
mkuznet1 marked this conversation as resolved.


class Console:
"""Class to run console commands.

Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -199,7 +237,7 @@ def sh(
if not secret:
raise RuntimeError(
"Subprocess '"
+ command
+ redact_secrets(command)
+ "' failed with exit code "
+ str(proc.returncode)
)
Expand Down
4 changes: 2 additions & 2 deletions src/madengine/execution/container_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
63 changes: 63 additions & 0 deletions tests/integration/test_console_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,10 @@
# project modules
from madengine.core import console

# Fake, non-real placeholder token. Split + f-string assembly keeps the
# literal "KEY=<token>" form out of source so secret scanners don't flag it.
_FAKE_HF = "hf_" + "abcdef123456"


class TestConsole:
"""Test the console module.
Expand Down Expand Up @@ -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

Comment on lines +72 to +77
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