multi-agent-mux/tests/test_issue12_registry_secrets.py

#!/usr/bin/env python3
"""Issue #12 — the broker password must not be persisted in job records.

Every sub-agent can read `.mam/jobs/<id>.json`, so `MQTT_PASSWORD` copied
there ends up in agent transcripts. The broker block now carries only
host/port/tls/username; the password is layered back from the environment
(`.mam.env`) at connect time. Records written before this rule are cleaned on
their next write and by `registry.py scrub-secrets`.
"""
from __future__ import annotations

import json
import os
import subprocess
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
SKILLS = REPO_ROOT / ".agents" / "skills"
REGISTRY_PY = SKILLS / "multi-agent-mux-delegate-job" / "scripts" / "registry.py"

sys.path.insert(0, str(SKILLS))
from lib_py.delegation import mqtt_common, registry  # noqa: E402


def _secure_env(monkeypatch):
    monkeypatch.setenv("MQTT_BROKER", "broker.example.test")
    monkeypatch.setenv("MQTT_USERNAME", "mam-user")
    monkeypatch.setenv("MQTT_PASSWORD", "s3cret-broker-pass")


def test_registry_block_carries_no_password(monkeypatch):
    _secure_env(monkeypatch)
    block = mqtt_common.broker_config_from_env().to_registry_block()
    assert "password" not in block
    assert block["username"] == "mam-user" and block["host"] == "broker.example.test"


def test_registered_record_has_no_password_but_connects_with_it(mam_sandbox, monkeypatch):
    _secure_env(monkeypatch)
    rd = str(mam_sandbox / ".mam" / "jobs")
    job_id = registry.register_job("issue 12", "claude-code", "herdr:claude", registry_dir=rd)

    record = json.loads((Path(rd) / f"{job_id}.json").read_text())
    assert "password" not in record["broker"], record["broker"]
    assert "s3cret-broker-pass" not in json.dumps(record)
    # The secure broker still gets its per-job HMAC key (needed by the agent's
    # own publish_event to sign; documented as remaining in the record).
    assert record["auth_token"]

    cfg = mqtt_common.broker_config_from_job(record)
    assert cfg.password == "s3cret-broker-pass", "password must come back from the environment"
    assert cfg.username == "mam-user"


def test_legacy_record_is_scrubbed_on_next_write(mam_sandbox, monkeypatch):
    _secure_env(monkeypatch)
    rd = str(mam_sandbox / ".mam" / "jobs")
    job_id = registry.register_job("legacy", "claude-code", "herdr:claude", registry_dir=rd)
    path = Path(rd) / f"{job_id}.json"
    record = json.loads(path.read_text())
    record["broker"]["password"] = "leaked-earlier"
    path.write_text(json.dumps(record))

    registry.update_status(job_id, rd, "running")

    after = json.loads(path.read_text())
    assert after["status"] == "running"
    assert "password" not in after["broker"]


def test_scrub_secrets_cli_migrates_a_registry(mam_sandbox, monkeypatch):
    _secure_env(monkeypatch)
    rd = str(mam_sandbox / ".mam" / "jobs")
    dirty = registry.register_job("dirty", "claude-code", "herdr:claude", registry_dir=rd)
    clean = registry.register_job("clean", "claude-code", "herdr:claude", registry_dir=rd)
    dirty_path = Path(rd) / f"{dirty}.json"
    rec = json.loads(dirty_path.read_text())
    rec["broker"]["password"] = "leaked-earlier"
    dirty_path.write_text(json.dumps(rec))
    clean_before = (Path(rd) / f"{clean}.json").read_text()
    # What the old to_registry_block wrote for every unauthenticated job: the key
    # with a null value. No secret, so the migration must neither report nor
    # rewrite it (a real PoC registry held ~900 of these and no password at all).
    empty = registry.register_job("empty-key", "claude-code", "herdr:claude", registry_dir=rd)
    empty_path = Path(rd) / f"{empty}.json"
    rec_empty = json.loads(empty_path.read_text())
    rec_empty["broker"]["password"] = None
    empty_path.write_text(json.dumps(rec_empty))
    empty_before = empty_path.read_text()

    env = dict(os.environ, PYTHONPATH=f"{SKILLS}:{os.environ.get('PYTHONPATH', '')}")
    dry = subprocess.run([sys.executable, str(REGISTRY_PY), "--registry-dir", rd, "scrub-secrets", "--dry-run"],
                         capture_output=True, text=True, env=env, cwd=str(mam_sandbox))
    assert dry.returncode == 0, dry.stderr
    assert f"would scrub 1 record(s)" in dry.stdout and dirty in dry.stdout
    assert empty not in dry.stdout, "an empty password key is not a stored secret"
    assert "leaked-earlier" in dirty_path.read_text(), "--dry-run must not write"

    res = subprocess.run([sys.executable, str(REGISTRY_PY), "--registry-dir", rd, "scrub-secrets"],
                         capture_output=True, text=True, env=env, cwd=str(mam_sandbox))
    assert res.returncode == 0, res.stderr
    assert "scrubbed 1 record(s)" in res.stdout
    scrubbed = json.loads(dirty_path.read_text())
    assert "password" not in scrubbed["broker"]
    assert scrubbed["prompt"] == "dirty" and scrubbed["auth_token"] == rec["auth_token"]
    assert (Path(rd) / f"{clean}.json").read_text() == clean_before, "untouched records are not rewritten"
    assert empty_path.read_text() == empty_before, "a null password key is left for the record's next write"