jj.wang
828 lines
42 KiB
Python
828 lines
42 KiB
Python
"""direct-ssh (PAD-12): config validation, channel selection and fallback rules, the ssh command line, the Windows
|
|
bootstrap / payload, single-call launch and collect. Nothing here opens a network connection: the POSIX direct
|
|
channel is replaced by a local `bash -c <loader>` fed on stdin, so the real loader and scripts run."""
|
|
import base64
|
|
import hashlib
|
|
import io
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import tempfile
|
|
import unittest
|
|
from contextlib import redirect_stdout
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
from agentd import ROOT, cli, config
|
|
from agentd.config import Config
|
|
from agentd.executors import Run, make_executor
|
|
from agentd.transports import (CallLog, HostKeyError, RoutedTransport, TransportError, Unreached, make_transport,
|
|
parse_fetch)
|
|
from agentd.transports import direct_ssh
|
|
from agentd.transports.direct_ssh import (BOOTSTRAP, PS_COMMAND, SH_COMMAND, DirectChannel, KeyAgent, classify,
|
|
ps_payload)
|
|
from agentd.transports.warpgate_ssh import PosixWarpgateTransport, WarpgateSshTransport, WindowsWarpgateTransport
|
|
from tests.fakes import config_data
|
|
|
|
SCHEMA = json.loads((ROOT / "config" / "schema.json").read_text())
|
|
WIN_KEY = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIFzLX7QNHDFYwTgY4wNl0ihZryIyKgpp9HPpPDWbrTX3"
|
|
LIN_KEY = "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAICUoI6W7npX+h/kHcXX3iy2XAWbNusrOjg28Xed75wWr"
|
|
SECRET = "-----BEGIN OPENSSH PRIVATE KEY-----\nNOT-A-REAL-KEY-but-must-never-leak\n-----END OPENSSH PRIVATE KEY-----\n"
|
|
|
|
|
|
def direct_data(root="/srv/x"):
|
|
d = config_data()
|
|
d["vault"] = {"client": "/opt/vault/agent-access.py", "config": "https://vault.example/identity.json"}
|
|
d["agentd"]["ssh_dir"] = str(Path(root) / "ssh")
|
|
d["terminals"]["win10"].update(transport="direct-ssh", fallback="warpgate-ssh",
|
|
ssh={"host": "100.123.209.19", "user": "administrator", "port": 22, "host_key": WIN_KEY})
|
|
d["terminals"]["lin"] = {"transport": "direct-ssh", "target": "lin", "os": "linux", "launch": "nohup",
|
|
"repo": str(Path(root) / "repo"), "run_root": str(Path(root) / "runs"), "plane_user": "oregon",
|
|
"token_key": "TOK_LIN", "concurrency": 1, "executors": ["claude-code", "shell"],
|
|
"ssh": {"host": "100.127.58.78", "user": "ubuntu", "host_key": LIN_KEY},
|
|
"fallback": "warpgate-ssh"}
|
|
return d
|
|
|
|
|
|
class FakeAgent:
|
|
sock, pub, error = "/tmp/agentd-agent-x/agent.sock", "/tmp/agentd-agent-x/agentd.pub", None
|
|
|
|
|
|
class ConfigTest(unittest.TestCase):
|
|
def errs(self, d):
|
|
return config.validate(d)
|
|
|
|
def test_direct_config_is_valid_and_live_config_uses_it(self):
|
|
self.assertEqual(self.errs(direct_data()), [])
|
|
self.assertEqual(config.mini_validate(direct_data(), SCHEMA), [])
|
|
cfg = config.load(ROOT / "config" / "agentd.yaml")
|
|
for term in ("win10", "oregon"):
|
|
t = cfg.terminals[term]
|
|
self.assertEqual((t["transport"], t["fallback"], t["target"]), ("direct-ssh", "warpgate-ssh", term))
|
|
self.assertIsNone(config.host_key_error(t["ssh"]["host_key"]))
|
|
self.assertIn("--secret-env AGENTD_SSH_KEY=AGENTD_SSH_KEY", cfg.secret_env_args())
|
|
config.load(ROOT / "config" / "agentd.oregon.yaml") # rollback config still valid
|
|
|
|
def test_bad_direct_settings_have_clear_errors(self):
|
|
cases = [
|
|
(lambda d: d["terminals"]["lin"].pop("ssh"), "direct-ssh needs ssh:"),
|
|
(lambda d: d["terminals"]["lin"]["ssh"].update(host_key="ssh-rsa AAAAB3Nza"), "must be 'ssh-ed25519 <base64>'"),
|
|
(lambda d: d["terminals"]["lin"]["ssh"].update(host_key=WIN_KEY[:-4]), "not an ssh-ed25519 public key"),
|
|
(lambda d: d["terminals"]["lin"]["ssh"].update(host_key="ssh-ed25519 !!!"), "not an ssh-ed25519 public key"),
|
|
(lambda d: d["terminals"]["lin"].pop("target"), "needs the Warpgate 'target'"),
|
|
(lambda d: d["terminals"]["lin"]["ssh"].update(host="-oProxyCommand=x"), "host and user must be non-empty"),
|
|
(lambda d: d["terminals"]["lin"]["ssh"].update(port=70000), "> 65535"),
|
|
(lambda d: d["terminals"]["lin"].update(run_root="~/runs"), "must be an absolute path"),
|
|
(lambda d: d["terminals"]["win10"].update(launch="nohup"), "direct-ssh on windows needs launch=schtasks"),
|
|
(lambda d: d["terminals"]["win10"].update(target="win11"), "(fallback) only reach target 'win10'"),
|
|
(lambda d: d["vault"].pop("config"), "fallback warpgate-ssh needs vault.config"),
|
|
(lambda d: d["terminals"]["oregon"].update(fallback="warpgate-ssh"), "only applies to transport direct-ssh"),
|
|
(lambda d: d["agentd"].update(ssh_dir="/" + "x" * 60), "too long for the ControlMaster socket"),
|
|
]
|
|
for mutate, fragment in cases:
|
|
d = direct_data()
|
|
mutate(d)
|
|
errs = self.errs(d)
|
|
self.assertTrue(any(fragment in e for e in errs), (fragment, errs))
|
|
|
|
def test_schema_rejects_unknown_ssh_key_and_fallback_value(self):
|
|
d = direct_data()
|
|
d["terminals"]["lin"]["ssh"]["password"] = "x"
|
|
self.assertTrue(any("password" in e for e in config.mini_validate(d, SCHEMA)))
|
|
d = direct_data()
|
|
d["terminals"]["lin"]["fallback"] = "local"
|
|
self.assertTrue(config.mini_validate(d, SCHEMA))
|
|
|
|
def test_key_env_is_always_scrubbed_and_only_requested_when_used(self):
|
|
cfg = Config(config_data()) # no direct-ssh terminal
|
|
self.assertNotIn("--secret-env AGENTD_SSH_KEY=AGENTD_SSH_KEY", cfg.secret_env_args())
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": SECRET}):
|
|
self.assertNotIn("AGENTD_SSH_KEY", cfg.scrubbed_env())
|
|
d = direct_data()
|
|
d["agentd"]["ssh_key"] = "OTHER_KEY"
|
|
self.assertEqual(Config(d).secret_env_args()[-1], "--secret-env OTHER_KEY=OTHER_KEY")
|
|
|
|
|
|
class SelectionTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.cfg = Config(direct_data())
|
|
|
|
def test_direct_terminal_gets_direct_primary_and_warpgate_fallback(self):
|
|
t = make_transport(self.cfg, "win10")
|
|
self.assertIsInstance(t, RoutedTransport)
|
|
self.assertIsInstance(t.primary, WindowsWarpgateTransport)
|
|
self.assertIsInstance(t.primary.channel, DirectChannel)
|
|
self.assertEqual((t.fallback.channel.via, t.fallback.channel.stdin, t.fallback.target), ("fallback", False, "win10"))
|
|
self.assertEqual(t.os, "windows")
|
|
self.assertTrue(hasattr(t, "start_run"))
|
|
|
|
def test_forced_channel_has_no_fallback(self):
|
|
self.assertIsNone(make_transport(self.cfg, "lin", force="direct-ssh").fallback)
|
|
t = make_transport(self.cfg, "lin", force="warpgate-ssh")
|
|
self.assertIsInstance(t, PosixWarpgateTransport)
|
|
self.assertEqual(t.channel.via, "warpgate")
|
|
|
|
def test_rollback_to_warpgate_keeps_ssh_section(self):
|
|
d = direct_data()
|
|
d["terminals"]["lin"]["transport"] = "warpgate-ssh" # ssh/fallback stay, ignored
|
|
self.assertEqual(config.validate(d), [])
|
|
self.assertIsInstance(make_transport(Config(d), "lin"), PosixWarpgateTransport)
|
|
|
|
|
|
class ClassifyTest(unittest.TestCase):
|
|
def test_classification(self):
|
|
self.assertEqual(classify(0, ""), "ok")
|
|
self.assertEqual(classify(1, "ssh: connect to host x port 22: Connection timed out"), "ok") # command ran
|
|
for err in ("ssh: connect to host 100.1.2.3 port 22: Connection timed out",
|
|
"ssh: connect to host 100.1.2.3 port 22: Connection refused",
|
|
"ssh: connect to host 100.1.2.3 port 22: No route to host",
|
|
"ssh: Could not resolve hostname x: Name or service not known",
|
|
"Connection timed out during banner exchange",
|
|
"kex_exchange_identification: read: Connection reset by peer",
|
|
"Connection closed by 100.1.2.3 port 22",
|
|
"ubuntu@100.1.2.3: Permission denied (publickey).",
|
|
"sign_and_send_pubkey: signing failed for ED25519 \"x\" from agent: agent refused operation"):
|
|
self.assertEqual(classify(255, err), "unreached", err)
|
|
for err in ("@@@@@@@@@@@\n@ WARNING: REMOTE HOST IDENTIFICATION HAS CHANGED! @\nHost key verification failed.",
|
|
"No ED25519 host key is known for agentd-win10 and you have requested strict checking.\n"
|
|
"Host key verification failed."):
|
|
self.assertEqual(classify(255, err), "hostkey")
|
|
for err in ("Connection to 100.1.2.3 closed by remote host.", "Timeout, server 100.1.2.3 not responding.",
|
|
"client_loop: send disconnect: Connection reset by peer", ""):
|
|
self.assertEqual(classify(255, err), "unknown", err) # the command may have run: no fallback
|
|
|
|
|
|
class Prim:
|
|
"""Stand-in transport recording calls; `fail` = exception raised by every call."""
|
|
os = "linux"
|
|
|
|
def __init__(self, fail=None):
|
|
self.fail, self.calls = fail, []
|
|
|
|
def join(self, *p):
|
|
return "/".join(p)
|
|
|
|
def start_run(self, *a):
|
|
self.calls.append(("start_run", a))
|
|
if self.fail:
|
|
raise self.fail
|
|
return "ok"
|
|
|
|
def process_count(self, name):
|
|
self.calls.append(("process_count", name))
|
|
if self.fail:
|
|
raise self.fail
|
|
return 0
|
|
|
|
|
|
class RoutedTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.events = []
|
|
self.log = CallLog(event=self.events.append)
|
|
|
|
def test_unreached_falls_back_once_and_sticks(self):
|
|
p, f = Prim(Unreached("ssh: connect to host x: Connection timed out")), Prim()
|
|
t = RoutedTransport("lin", p, f, self.log)
|
|
self.assertEqual(t.start_run("task", "/r", {}, "run.sh"), "ok")
|
|
self.assertEqual(t.process_count("codex"), 0)
|
|
self.assertEqual(len(p.calls), 1) # direct not retried within the tick
|
|
self.assertEqual([c[0] for c in f.calls], ["start_run", "process_count"])
|
|
self.assertEqual(len(self.events), 1)
|
|
self.assertIn("fallback", self.events[0])
|
|
self.assertEqual(t.join("a", "b"), "a/b") # pure helpers come from the primary
|
|
|
|
def test_host_key_and_after_command_failures_never_fall_back(self):
|
|
for exc in (HostKeyError("host key verification FAILED"), TransportError("ssh timeout after 150s")):
|
|
p, f = Prim(exc), Prim()
|
|
t = RoutedTransport("lin", p, f, self.log)
|
|
with self.assertRaises(type(exc)):
|
|
t.start_run("task", "/r", {}, "run.sh")
|
|
self.assertEqual(f.calls, [])
|
|
|
|
def test_no_fallback_configured_raises(self):
|
|
with self.assertRaises(Unreached):
|
|
RoutedTransport("lin", Prim(Unreached("down")), None, self.log).start_run("t", "/r", {}, "run.sh")
|
|
|
|
def test_busy_check_unreached_is_not_swallowed(self):
|
|
cfg = Config(direct_data())
|
|
for cls, term in ((PosixWarpgateTransport, "lin"), (WindowsWarpgateTransport, "win10")):
|
|
ch = mock.Mock(stdin=True)
|
|
ch.sh.side_effect = ch.ps.side_effect = Unreached("down")
|
|
with self.assertRaises(Unreached):
|
|
cls(cfg, term, channel=ch).process_count("codex")
|
|
ch.sh.side_effect = ch.ps.side_effect = TransportError("exit 1")
|
|
self.assertIsNone(cls(cfg, term, channel=ch).process_count("codex"))
|
|
|
|
|
|
class MissingKeyTest(unittest.TestCase):
|
|
def setUp(self):
|
|
direct_ssh._AGENTS.clear()
|
|
direct_ssh._KEYS.clear()
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.cfg = Config(direct_data(self.tmp.name))
|
|
|
|
def tearDown(self):
|
|
direct_ssh._AGENTS.clear()
|
|
self.tmp.cleanup()
|
|
|
|
def test_missing_key_makes_direct_unavailable_and_uses_the_fallback(self):
|
|
events = []
|
|
log = CallLog(event=events.append)
|
|
env = {k: v for k, v in os.environ.items() if k != "AGENTD_SSH_KEY"}
|
|
with mock.patch.dict(os.environ, env, clear=True), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run") as run:
|
|
t = make_transport(self.cfg, "lin", log=log)
|
|
t.fallback = Prim()
|
|
self.assertEqual(t.process_count("claude"), 0)
|
|
self.assertEqual(t.process_count("claude"), 0)
|
|
run.assert_not_called() # no ssh attempted at all
|
|
self.assertTrue(any("AGENTD_SSH_KEY is not set" in e for e in events))
|
|
self.assertEqual(sum("unavailable in this process" in e for e in events), 1) # once per process
|
|
self.assertEqual(sum("using the warpgate-ssh fallback" in e for e in events), 1) # once per tick
|
|
|
|
def test_empty_key_is_popped_and_reported(self):
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": " "}):
|
|
a = KeyAgent("AGENTD_SSH_KEY").start()
|
|
self.assertNotIn("AGENTD_SSH_KEY", os.environ)
|
|
self.assertIn("not set", a.error)
|
|
|
|
@unittest.skipUnless(all(shutil.which(b) for b in ("ssh-agent", "ssh-add", "ssh-keygen", "ssh")), "needs OpenSSH")
|
|
def test_real_agent_loads_key_from_env_and_cleans_up(self):
|
|
keyfile = Path(self.tmp.name) / "k"
|
|
subprocess.run(["ssh-keygen", "-q", "-t", "ed25519", "-N", "", "-C", "agentd-test", "-f", str(keyfile)], check=True)
|
|
key = keyfile.read_text()
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": key.replace("\n", "\\n")}): # Vault-style escapes
|
|
a = KeyAgent("AGENTD_SSH_KEY").start()
|
|
self.assertNotIn("AGENTD_SSH_KEY", os.environ)
|
|
try:
|
|
self.assertIsNone(a.error)
|
|
self.assertEqual(oct(os.stat(a.dir).st_mode & 0o777), "0o700")
|
|
self.assertEqual(Path(a.pub).read_text().split()[:2], (keyfile.with_suffix(".pub")).read_text().split()[:2])
|
|
for f in Path(a.dir).iterdir(): # nothing on disk holds the private key
|
|
if f.is_file():
|
|
self.assertNotIn("PRIVATE KEY", f.read_text())
|
|
finally:
|
|
a.close()
|
|
self.assertFalse(os.path.exists(a.dir))
|
|
self.assertIsNotNone(a.proc.poll())
|
|
|
|
def test_bad_key_reports_without_leaking(self):
|
|
if not all(shutil.which(b) for b in ("ssh-agent", "ssh-add", "ssh")):
|
|
self.skipTest("needs OpenSSH")
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": SECRET}):
|
|
a = KeyAgent("AGENTD_SSH_KEY").start()
|
|
self.assertIn("ssh-add failed", a.error)
|
|
self.assertNotIn("NOT-A-REAL-KEY", a.error)
|
|
self.assertFalse(os.path.exists(a.dir)) # agent stopped, socket dir gone
|
|
|
|
|
|
class SshCommandTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.cfg = Config(direct_data(self.tmp.name))
|
|
self.events, self.lines = [], []
|
|
self.log = CallLog(event=self.events.append)
|
|
self.log.call = lambda *a, **k: self.lines.append(a)
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def run_with(self, rc, err=b"", out=b"PONG=x\n"):
|
|
seen = {}
|
|
|
|
def fake_run(argv, input=None, stdout=None, stderr=None, env=None, timeout=None):
|
|
seen.update(argv=argv, input=input, env=env, timeout=timeout)
|
|
stdout.write(out)
|
|
stderr.write(err)
|
|
return mock.Mock(returncode=rc)
|
|
return seen, fake_run
|
|
|
|
def test_command_line_options_and_no_secret(self):
|
|
ch = DirectChannel(self.cfg, "lin", self.log)
|
|
seen, fake = self.run_with(0)
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": SECRET, "TOK_LIN": "plane-token", "SSH_AUTH_SOCK": "/x"}), \
|
|
mock.patch("agentd.transports.direct_ssh.key_agent", return_value=FakeAgent), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run", fake):
|
|
rc, out, _ = ch.sh("echo hi", 30, "ping")
|
|
argv = seen["argv"]
|
|
opts = {argv[i + 1] for i, a in enumerate(argv) if a == "-o"}
|
|
ssh_dir = Path(self.tmp.name) / "ssh"
|
|
for o in ("BatchMode=yes", "IdentitiesOnly=yes", "IdentityAgent=" + FakeAgent.sock, "IdentityFile=" + FakeAgent.pub,
|
|
"StrictHostKeyChecking=yes", "UserKnownHostsFile=%s/known_hosts.lin" % ssh_dir,
|
|
"GlobalKnownHostsFile=/dev/null", "HostKeyAlias=agentd-lin", "HostKeyAlgorithms=ssh-ed25519",
|
|
"ConnectTimeout=8", "ServerAliveInterval=15", "ServerAliveCountMax=3", "ForwardAgent=no",
|
|
"ForwardX11=no", "LogLevel=ERROR", "ControlMaster=auto", "ControlPersist=60",
|
|
"ControlPath=%s/%s-%%C" % (ssh_dir, ch.pin), "PasswordAuthentication=no"):
|
|
self.assertIn(o, opts)
|
|
self.assertEqual(argv[:3], ["ssh", "-F", "/dev/null"])
|
|
self.assertEqual(argv[-3:], ["--", "100.127.58.78", SH_COMMAND])
|
|
self.assertEqual(argv[argv.index("-l") + 1], "ubuntu")
|
|
self.assertEqual(seen["input"], b"echo hi") # the script travels on stdin
|
|
self.assertFalse(any("NOT-A-REAL-KEY" in a or "plane-token" in a for a in argv))
|
|
self.assertNotIn("AGENTD_SSH_KEY", seen["env"])
|
|
self.assertNotIn("TOK_LIN", seen["env"])
|
|
self.assertNotIn("SSH_AUTH_SOCK", seen["env"])
|
|
self.assertEqual((Path(ssh_dir) / "known_hosts.lin").read_text(), "agentd-lin %s\n" % LIN_KEY)
|
|
self.assertEqual(oct(os.stat(ssh_dir).st_mode & 0o777), "0o700")
|
|
self.assertEqual((rc, out), (0, "PONG=x\n"))
|
|
self.assertEqual(self.lines[-1][:3], ("lin", "direct", "ping"))
|
|
self.assertEqual(self.lines[-1][4], 0)
|
|
|
|
def test_host_key_failure_is_loud_sticky_and_not_unreached(self):
|
|
ch = DirectChannel(self.cfg, "lin", self.log)
|
|
seen, fake = self.run_with(255, b"No ED25519 host key is known for agentd-lin and you have requested strict "
|
|
b"checking.\r\nHost key verification failed.\r\n", b"")
|
|
with mock.patch("agentd.transports.direct_ssh.key_agent", return_value=FakeAgent), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run", side_effect=fake) as run:
|
|
with self.assertRaises(HostKeyError) as cm:
|
|
ch.sh("true", 30)
|
|
self.assertNotIsInstance(cm.exception, Unreached)
|
|
with self.assertRaises(HostKeyError):
|
|
ch.sh("true", 30)
|
|
self.assertEqual(run.call_count, 1) # no second attempt in this tick
|
|
self.assertTrue(self.events and self.events[0].startswith("SECURITY lin:"))
|
|
|
|
def test_connect_failure_is_unreached_and_timeout_is_not(self):
|
|
ch = DirectChannel(self.cfg, "lin", self.log)
|
|
_, fake = self.run_with(255, b"ssh: connect to host 100.127.58.78 port 22: Connection timed out\n", b"")
|
|
with mock.patch("agentd.transports.direct_ssh.key_agent", return_value=FakeAgent), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run", fake):
|
|
with self.assertRaises(Unreached):
|
|
ch.sh("true", 30)
|
|
with mock.patch("agentd.transports.direct_ssh.key_agent", return_value=FakeAgent), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run",
|
|
side_effect=subprocess.TimeoutExpired("ssh", 30)):
|
|
with self.assertRaises(TransportError) as cm:
|
|
ch.sh("true", 30)
|
|
self.assertNotIsInstance(cm.exception, Unreached) # may have run: resolve_launch probes instead
|
|
|
|
def test_windows_ps_uses_bootstrap_and_stdin(self):
|
|
ch = DirectChannel(self.cfg, "win10", self.log)
|
|
seen, fake = self.run_with(0, out=b"X=ok\r\n")
|
|
with mock.patch("agentd.transports.direct_ssh.key_agent", return_value=FakeAgent), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run", fake):
|
|
self.assertEqual(ch.ps("'X=ok'", 30), "X=ok\n")
|
|
self.assertEqual(seen["argv"][-2:], ["100.123.209.19", PS_COMMAND])
|
|
self.assertEqual(base64.b64decode(seen["input"]).decode("utf-8"), "'X=ok'")
|
|
|
|
|
|
class BootstrapTest(unittest.TestCase):
|
|
def test_encoded_command_round_trip_and_length(self):
|
|
value = PS_COMMAND.split()[-1]
|
|
self.assertEqual(base64.b64decode(value).decode("utf-16-le"), BOOTSTRAP)
|
|
self.assertLess(len(PS_COMMAND), 2000) # cmd.exe limit is 8191
|
|
self.assertIn("[Console]::In.ReadToEnd()", BOOTSTRAP)
|
|
script = "'é 中文 ''q'''\n"
|
|
self.assertTrue(ps_payload(script).isascii())
|
|
self.assertEqual(base64.b64decode(ps_payload(script)).decode("utf-8"), script)
|
|
|
|
|
|
class FakeWinChannel:
|
|
"""Direct-like Windows channel: answers the start / fetch scripts the way PowerShell would."""
|
|
stdin, name, via = True, "ssh fake", "direct"
|
|
|
|
def __init__(self, files=None, haspid=False):
|
|
self.calls, self.files, self.haspid = [], files or {}, haspid
|
|
|
|
def ps(self, script, timeout, op=None):
|
|
self.calls.append((op, script))
|
|
if op == "start":
|
|
if self.haspid:
|
|
return "HASPID\n"
|
|
out = ["WROTE=" + hashlib.sha256(base64.b64decode(b)).hexdigest()
|
|
for b in __import__("re").findall(r"FromBase64String\('([A-Za-z0-9+/=]*)'\)", script)]
|
|
return "\n".join(out + ["LAUNCHED=0"]) + "\n"
|
|
if op == "fetch":
|
|
ops = [l.split()[:4] for l in script.splitlines() if l.split(" ")[0] in ("AdRead", "AdGrep", "AdTail")]
|
|
out = []
|
|
for fn, i, name, *rest in ops:
|
|
text = self.files.get(name.strip("'"))
|
|
if fn == "AdRead":
|
|
out.append("M%s" % i if text is None else "F%s=%s" % (i, base64.b64encode(text.encode()).decode()))
|
|
else:
|
|
if fn == "AdGrep":
|
|
text = next((l for l in (text or "").splitlines() if "thread.started" in l), "")
|
|
else:
|
|
text = "\n".join((text or "").splitlines()[-int(rest[0]):])
|
|
out.append("F%s=%s" % (i, base64.b64encode(text.encode()).decode()))
|
|
return "\n".join(out) + "\n"
|
|
return "FORGOT=ok\n"
|
|
|
|
|
|
class WindowsDirectTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.cfg = config.load(ROOT / "config" / "agentd.yaml")
|
|
|
|
def transport(self, ch):
|
|
return WarpgateSshTransport(self.cfg, "win10", channel=ch)
|
|
|
|
def test_launch_is_one_call_with_exact_utf8_files(self):
|
|
ch = FakeWinChannel()
|
|
t = self.transport(ch)
|
|
ex = make_executor(self.cfg, "win10", "codex", t)
|
|
prompt = "中文 'quoted' \"dq\" $env:X `tick`\r\n" * 2000 # ~70 KB: far beyond any argv limit
|
|
h = ex.start(Run("GW-26-1", t.repo, prompt, "gpt-6-astra", "high"))
|
|
self.assertEqual([c[0] for c in ch.calls], ["start"])
|
|
script = ch.calls[0][1]
|
|
blobs = __import__("re").findall(r"FromBase64String\('([A-Za-z0-9+/=]*)'\)", script)
|
|
self.assertEqual(base64.b64decode(blobs[0]), prompt.encode("utf-8")) # UTF-8, no BOM, bytes exact
|
|
self.assertFalse(base64.b64decode(blobs[1]).startswith(b"\xef\xbb\xbf"))
|
|
self.assertIn("Set-Location -LiteralPath", base64.b64decode(blobs[1]).decode("utf-8"))
|
|
self.assertLess(script.index("'HASPID'; return"), script.index("WriteAllBytes"))
|
|
self.assertLess(script.index("WriteAllBytes"), script.index("schtasks /Run /TN 'agentd-GW-26-1'"))
|
|
self.assertTrue(h["run_dir"].endswith("\\runtime\\agentd\\GW-26-1"))
|
|
|
|
def test_launch_refuses_a_started_directory(self):
|
|
t = self.transport(FakeWinChannel(haspid=True))
|
|
with self.assertRaises(TransportError) as cm:
|
|
t.start_run("agentd-GW-1-1", "C:\\r\\GW-1-1", {"prompt.md": "p", "run.ps1": "r"}, "run.ps1")
|
|
self.assertIn("HASPID", str(cm.exception))
|
|
|
|
def test_collect_is_one_call_plus_forget(self):
|
|
files = {"exit.txt": "\ufeffexit=0 at 2026-10-22T10:00:00", "last.md": "[review] done\n",
|
|
"events.jsonl": '{"type":"x"}\n{"type":"thread.started","thread_id":"th-42"}\n{"type":"turn.completed"}',
|
|
"stderr.txt": "warn 1\nwarn 2"}
|
|
ch = FakeWinChannel(files)
|
|
ex = make_executor(self.cfg, "win10", "codex", self.transport(ch))
|
|
h = {"run_dir": "C:\\r\\GW-1-1", "task": "agentd-GW-1-1", "thread_id": None}
|
|
res = ex.collect(h)
|
|
ex.forget(h)
|
|
self.assertEqual([c[0] for c in ch.calls], ["fetch", "forget"])
|
|
self.assertEqual((res.last_message, res.exit_code, res.thread_id), ("[review] done", 0, "th-42"))
|
|
self.assertEqual(res.tail, 'warn 1\nwarn 2\n{"type":"turn.completed"}')
|
|
|
|
def test_warpgate_fetch_is_one_call_when_it_fits(self):
|
|
cfg = config.load(ROOT / "config" / "agentd.example.yaml")
|
|
t = WindowsWarpgateTransport(cfg, "win10")
|
|
self.assertFalse(t.channel.stdin)
|
|
ops = make_executor(cfg, "win10", "codex", t).COLLECT
|
|
script = t.fetch_script(t.join(t.default_run_root(), "GW-26-1-r2"), ops)
|
|
self.assertTrue(t.channel.fits(script))
|
|
calls = []
|
|
t.ps = lambda s, timeout=150: (calls.append(s), "\n".join("M%d" % i for i in range(len(ops))))[1]
|
|
self.assertEqual(t.fetch("C:\\r", ops), [None, None, None, "", ""])
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertFalse(t.channel.fits("x" * 3000))
|
|
|
|
|
|
class LocalDirect:
|
|
"""POSIX direct channel stand-in: runs the real loader (SH_COMMAND) through a local login-style shell with the
|
|
script on stdin, exactly what sshd does with the remote command."""
|
|
stdin, name, via = True, "ssh local", "direct"
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def sh(self, script, timeout, op=None):
|
|
self.calls.append(op)
|
|
p = subprocess.run(["sh", "-c", SH_COMMAND], input=script.encode("utf-8"), # as sshd: $SHELL -c "<command>"
|
|
capture_output=True, timeout=timeout)
|
|
return p.returncode, p.stdout.decode("utf-8", "replace"), p.stderr.decode("utf-8", "replace")
|
|
|
|
|
|
class PosixDirectTest(unittest.TestCase):
|
|
def setUp(self):
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
root = Path(self.tmp.name)
|
|
(root / "repo").mkdir()
|
|
self.cfg = Config(direct_data(root))
|
|
self.ch = LocalDirect()
|
|
self.t = PosixWarpgateTransport(self.cfg, "lin", channel=self.ch)
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def test_loader_quoting_matches_sh_quote(self):
|
|
self.assertEqual(SH_COMMAND, "bash -c 'exec 3<&0 </dev/null; s=$(cat <&3); exec 3<&-; (eval \"$s\"); r=$?; "
|
|
"[ $r = 255 ] && r=254; exit $r'")
|
|
|
|
def test_loader_never_exits_255(self):
|
|
self.assertEqual(self.ch.sh("echo x >&2; exit 255", 10)[0], 254) # 255 = ssh's own "no session" code
|
|
self.assertEqual(self.ch.sh("exit 3", 10)[0], 3)
|
|
self.assertEqual(self.ch.sh("true", 10)[0], 0)
|
|
|
|
def test_big_prompt_launch_and_collect_are_one_call_each(self):
|
|
ex = make_executor(self.cfg, "lin", "shell", self.t)
|
|
prompt = ("ünïcode 'q' \"d\" $(nope) `x`\n" * 6000) # ~200 KB > CHUNK and > MAX_ARG_STRLEN
|
|
h = ex.start(Run("GW-7-1", self.cfg.terminals["lin"]["repo"], prompt, script="wc -c"))
|
|
self.assertEqual(self.ch.calls, ["start"])
|
|
self.assertEqual(Path(h["run_dir"], "prompt.md").read_text(encoding="utf-8"), prompt)
|
|
for _ in range(100):
|
|
if ex.poll(h) == "finished":
|
|
break
|
|
__import__("time").sleep(0.1)
|
|
n = len(self.ch.calls)
|
|
res = ex.collect(h)
|
|
self.assertEqual(self.ch.calls[n:], ["fetch"])
|
|
self.assertEqual((res.last_message, res.exit_code), (str(len(prompt.encode("utf-8"))), 0))
|
|
with self.assertRaises(TransportError): # pid.txt present: never wiped, never relaunched there
|
|
ex.start(Run("GW-7-1", self.cfg.terminals["lin"]["repo"], "again", script="true"))
|
|
self.assertEqual(Path(h["run_dir"], "prompt.md").read_text(encoding="utf-8"), prompt)
|
|
|
|
def test_fetch_missing_files_and_stdin_is_not_the_scripts(self):
|
|
d = Path(self.tmp.name) / "runs" / "x"
|
|
d.mkdir(parents=True)
|
|
(d / "a.txt").write_text("l1\nl2 needle\nl3\n")
|
|
ops = [("read", "a.txt"), ("read", "none"), ("grep", "a.txt", "needle"), ("grep", "none", "x"),
|
|
("tail", "a.txt", 2), ("tail", "none", 3)]
|
|
self.assertEqual(self.t.fetch(str(d), ops), ["l1\nl2 needle\nl3\n", None, "l2 needle", None, "l2 needle\nl3", ""])
|
|
rc, out, _ = self.ch.sh("cat; echo done", 10) # the script's own stdin is /dev/null
|
|
self.assertEqual((rc, out), (0, "done\n"))
|
|
|
|
def test_parse_fetch_rejects_garbage(self):
|
|
with self.assertRaises(TransportError):
|
|
parse_fetch("nothing here", [("read", "x")])
|
|
|
|
|
|
class CallLogTest(unittest.TestCase):
|
|
def test_line_format_and_once(self):
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
events = []
|
|
log = CallLog(Path(tmp) / "sub" / "remote.log", events.append)
|
|
log.call("win10", "direct", "start", 812.4, 0)
|
|
log.call("win10", "fallback", "busy", 9100, 255, "unreached")
|
|
lines = (Path(tmp) / "sub" / "remote.log").read_text().splitlines()
|
|
self.assertTrue(lines[0].endswith("term=win10 via=direct op=start ms=812 rc=0"))
|
|
self.assertTrue(lines[1].endswith("via=fallback op=busy ms=9100 rc=255 unreached"))
|
|
log.event("a", once="k")
|
|
log.event("a", once="k")
|
|
self.assertEqual(events, ["a"])
|
|
|
|
|
|
class PingCliTest(unittest.TestCase):
|
|
def test_ping_prints_timing_and_forces_channel(self):
|
|
cfg = Config(direct_data())
|
|
seen = {}
|
|
|
|
def fake_make(cfg_, term, log=None, force=None):
|
|
seen.update(term=term, force=force)
|
|
return mock.Mock(ping=mock.Mock(return_value="PONG=Linux 6.8 oregon"))
|
|
buf = io.StringIO()
|
|
with mock.patch("agentd.cli.make_transport", fake_make), redirect_stdout(buf), self.assertRaises(SystemExit) as cm:
|
|
cli.cmd_ping(cfg, cli.parser().parse_args(["ping", "lin", "--transport", "warpgate-ssh", "-n", "2"]))
|
|
self.assertEqual(cm.exception.code, 0)
|
|
self.assertEqual(seen, {"term": "lin", "force": "warpgate-ssh"})
|
|
self.assertEqual(buf.getvalue().count("PONG=Linux 6.8 oregon"), 2)
|
|
|
|
|
|
class ReviewFixesTest(unittest.TestCase):
|
|
"""Follow-ups of the PAD-12 review: classification limits, socket and known_hosts names, key handling."""
|
|
def setUp(self):
|
|
direct_ssh._KEYS.clear()
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
self.cfg = Config(direct_data(self.tmp.name))
|
|
|
|
def tearDown(self):
|
|
direct_ssh._KEYS.clear()
|
|
self.tmp.cleanup()
|
|
|
|
def test_output_or_remote_stderr_never_counts_as_unreached(self):
|
|
self.assertEqual(classify(255, "Connection reset by 100.1.2.3 port 22", "WROTE=ab\n"), "unknown")
|
|
self.assertEqual(classify(255, "git: ssh: connect to host github.com port 22: Connection timed out"), "unknown")
|
|
self.assertEqual(classify(255, "fatal: Could not read from remote repository"), "unknown")
|
|
|
|
def test_no_ed25519_offered_is_a_host_key_failure(self):
|
|
self.assertEqual(classify(255, "Unable to negotiate with 100.1.2.3 port 22: no matching host key type found. "
|
|
"Their offer: ssh-rsa"), "hostkey")
|
|
self.assertEqual(classify(255, "Unable to negotiate with 100.1.2.3 port 22: no matching cipher found. "
|
|
"Their offer: aes128-cbc"), "unreached")
|
|
|
|
def test_reset_after_output_is_not_retried_elsewhere(self):
|
|
ch = DirectChannel(self.cfg, "lin")
|
|
|
|
def fake_run(argv, input=None, stdout=None, stderr=None, env=None, timeout=None):
|
|
stdout.write(b"WROTE=1\n")
|
|
stderr.write(b"Connection reset by 100.127.58.78 port 22\r\n")
|
|
return mock.Mock(returncode=255)
|
|
with mock.patch("agentd.transports.direct_ssh.key_agent", return_value=FakeAgent), \
|
|
mock.patch("agentd.transports.direct_ssh.subprocess.run", fake_run):
|
|
with self.assertRaises(TransportError) as cm:
|
|
ch.sh("true", 30)
|
|
self.assertNotIsInstance(cm.exception, Unreached)
|
|
|
|
def test_control_path_follows_the_pinned_key(self):
|
|
a = DirectChannel(self.cfg, "lin")
|
|
d = direct_data(self.tmp.name)
|
|
d["terminals"]["lin"]["ssh"]["host_key"] = WIN_KEY
|
|
b = DirectChannel(Config(d), "lin")
|
|
path = lambda ch: [x for x in ch.argv(FakeAgent, "true") if x.startswith("ControlPath=")][0] # noqa: E731
|
|
self.assertNotEqual(path(a), path(b))
|
|
self.assertEqual(path(a), "ControlPath=%s/%s-%%C" % (self.cfg.ssh_dir, a.pin))
|
|
self.assertEqual(len(a.pin), 8)
|
|
long_dir = direct_data("/" + "x" * 40)
|
|
self.assertTrue(any("too long" in e for e in config.validate(long_dir))) # 41 + 67 > 104
|
|
|
|
def test_known_hosts_temp_files_are_per_terminal_and_process(self):
|
|
d = direct_ssh.control_dir(self.cfg)
|
|
tmps = []
|
|
real = os.replace
|
|
with mock.patch("agentd.transports.direct_ssh.os.replace",
|
|
side_effect=lambda a, b: (tmps.append(str(a)), real(a, b))):
|
|
DirectChannel(self.cfg, "win10").known_hosts(d)
|
|
DirectChannel(self.cfg, "lin").known_hosts(d)
|
|
self.assertEqual(len(set(tmps)), 2)
|
|
self.assertTrue(all(str(os.getpid()) in t for t in tmps))
|
|
self.assertEqual((d / "known_hosts.win10").read_text(), "agentd-win10 %s\n" % WIN_KEY)
|
|
self.assertEqual((d / "known_hosts.lin").read_text(), "agentd-lin %s\n" % LIN_KEY)
|
|
self.assertEqual(sorted(f.name for f in d.iterdir()), ["known_hosts.lin", "known_hosts.win10"])
|
|
|
|
def test_key_leaves_the_environment_at_cli_start(self):
|
|
buf = io.StringIO()
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": SECRET}), redirect_stdout(buf):
|
|
cli.main(["--config", str(ROOT / "config" / "agentd.yaml"), "check-config"])
|
|
self.assertNotIn("AGENTD_SSH_KEY", os.environ)
|
|
self.assertIn("direct-ssh key AGENTD_SSH_KEY: set", buf.getvalue())
|
|
self.assertNotIn("NOT-A-REAL-KEY", buf.getvalue())
|
|
self.assertTrue(direct_ssh.key_present("AGENTD_SSH_KEY"))
|
|
with mock.patch("agentd.transports.direct_ssh.shutil.which", return_value=None):
|
|
a = KeyAgent("AGENTD_SSH_KEY").start() # the agent takes the key over from the CLI
|
|
self.assertIn("not installed", a.error)
|
|
self.assertFalse(direct_ssh.key_present("AGENTD_SSH_KEY"))
|
|
|
|
def test_local_busy_check_gets_the_scrubbed_env(self):
|
|
cfg = Config(config_data())
|
|
from agentd.transports.local import LocalTransport
|
|
with mock.patch.dict(os.environ, {"AGENTD_SSH_KEY": SECRET, "TOK_OR": "plane-token"}), \
|
|
mock.patch("agentd.transports.local.subprocess.run", return_value=mock.Mock(stdout="2\n")) as run:
|
|
self.assertEqual(LocalTransport(cfg, "oregon").process_count("codex"), 2)
|
|
env = run.call_args.kwargs["env"]
|
|
self.assertNotIn("AGENTD_SSH_KEY", env)
|
|
self.assertNotIn("TOK_OR", env)
|
|
|
|
def test_task_names_are_quoted_in_kill_and_forget(self):
|
|
ch = mock.Mock(stdin=True)
|
|
ch.ps.return_value = "KILLED=ok"
|
|
t = WindowsWarpgateTransport(self.cfg, "win10", channel=ch)
|
|
t.kill("agentd-a'b", "C:\\runs\\x")
|
|
t.forget("agentd-a'b")
|
|
for c in ch.ps.call_args_list:
|
|
self.assertIn("'agentd-a''b'", c.args[0])
|
|
self.assertNotIn("'agentd-a'b'", c.args[0])
|
|
|
|
|
|
class FakeTerm:
|
|
"""Remote transport stand-in for core-level routing tests: pure helpers from the real Windows transport,
|
|
start_run / probe on a dict of run directories shared by the direct and fallback channels. `fail` = exceptions
|
|
for the next calls: Unreached and HostKeyError strike before the command ran, any other after it (the run
|
|
exists). A host-key failure sticks, like DirectChannel's."""
|
|
|
|
def __init__(self, cfg, dirs):
|
|
self.real = WindowsWarpgateTransport(cfg, "win10", channel=mock.Mock(stdin=True))
|
|
self.dirs, self.fail, self.calls, self.hostkey = dirs, [], [], None
|
|
|
|
def __getattr__(self, name):
|
|
return getattr(self.real, name)
|
|
|
|
def _next(self, name):
|
|
self.calls.append(name)
|
|
if self.hostkey:
|
|
raise self.hostkey
|
|
exc = self.fail.pop(0) if self.fail else None
|
|
if isinstance(exc, HostKeyError):
|
|
self.hostkey = exc
|
|
return exc
|
|
|
|
def start_run(self, task, run_dir, files, script_name):
|
|
exc = self._next("start_run")
|
|
if isinstance(exc, (Unreached, HostKeyError)):
|
|
raise exc
|
|
if run_dir in self.dirs:
|
|
raise TransportError("HASPID")
|
|
self.dirs[run_dir] = ["launch.pid"]
|
|
if exc:
|
|
raise exc
|
|
|
|
def probe(self, run_dir, names):
|
|
exc = self._next("probe")
|
|
if exc:
|
|
raise exc
|
|
return list(self.dirs.get(run_dir, ()))
|
|
|
|
|
|
class CoreRoutingTest(unittest.TestCase):
|
|
"""Agentd.start -> codex executor -> RoutedTransport: the double-launch guarantee end to end."""
|
|
|
|
def setUp(self):
|
|
from tests.fakes import Harness
|
|
self.h = Harness(direct_data())
|
|
self.dirs, self.events = {}, []
|
|
self.direct, self.fallback = FakeTerm(self.h.cfg, self.dirs), FakeTerm(self.h.cfg, self.dirs)
|
|
log = CallLog(event=self.events.append)
|
|
self.h.d.executor_factory = lambda term, name: make_executor( # a fresh routed transport per tick
|
|
self.h.cfg, term, name, RoutedTransport(term, self.direct, self.fallback, log))
|
|
|
|
def tearDown(self):
|
|
self.h.close()
|
|
|
|
def test_unreached_launch_moves_to_the_fallback_once(self):
|
|
iid = self.h.plane.add_issue(1)
|
|
self.direct.fail = [Unreached("ssh: connect to host 100.123.209.19 port 22: Connection timed out")]
|
|
self.h.tick()
|
|
rec = self.h.rec(iid)
|
|
self.assertNotIn("launching", rec)
|
|
self.assertEqual(rec["inflight"]["run_id"], "GW-1-1")
|
|
self.assertEqual((self.direct.calls, self.fallback.calls), (["start_run"], ["start_run"]))
|
|
self.assertEqual(len(self.dirs), 1)
|
|
self.assertEqual(self.h.plane.state_of(iid), "In Progress")
|
|
self.assertTrue(any("warpgate-ssh fallback" in e for e in self.events))
|
|
|
|
def test_failure_after_the_command_ran_is_probed_and_adopted_not_relaunched(self):
|
|
iid = self.h.plane.add_issue(1)
|
|
self.direct.fail = [TransportError("ssh 100.123.209.19 timeout after 150s")]
|
|
self.h.tick()
|
|
rec = self.h.rec(iid)
|
|
self.assertIn("launching", rec)
|
|
self.assertNotIn("inflight", rec)
|
|
self.assertEqual(self.fallback.calls, [])
|
|
self.h.tick() # probe finds launch.pid -> adopt
|
|
rec = self.h.rec(iid)
|
|
self.assertNotIn("launching", rec)
|
|
self.assertEqual(rec["inflight"]["handle"]["run_dir"], list(self.dirs)[0])
|
|
self.assertEqual(self.direct.calls, ["start_run", "probe"])
|
|
self.assertEqual((self.fallback.calls, len(self.dirs)), ([], 1))
|
|
|
|
def test_host_key_failure_keeps_the_marker_and_never_falls_back(self):
|
|
iid = self.h.plane.add_issue(1)
|
|
self.direct.fail = [HostKeyError("host key verification FAILED")]
|
|
self.h.tick()
|
|
self.assertIn("launching", self.h.rec(iid))
|
|
self.h.tick() # the probe fails the same way: marker kept, slot held
|
|
rec = self.h.rec(iid)
|
|
self.assertIn("launching", rec)
|
|
self.assertNotIn("inflight", rec)
|
|
self.assertEqual(self.direct.calls, ["start_run", "probe"])
|
|
self.assertEqual((self.fallback.calls, self.dirs), ([], {}))
|
|
self.assertEqual(self.h.plane.state_of(iid), "Todo")
|
|
|
|
|
|
FAKE_CLIENT = """import json, os, sys
|
|
argv = sys.argv[1:]
|
|
with open(os.environ["FAKE_CALLS"], "a") as fh:
|
|
fh.write(json.dumps(argv) + "\\n")
|
|
if os.environ.get("FAKE_MISSING") and "AGENTD_SSH_KEY=AGENTD_SSH_KEY" in argv:
|
|
print("access-hub: credential_unavailable", file=sys.stderr)
|
|
sys.exit(1)
|
|
print("ticked")
|
|
print("agentd stderr line", file=sys.stderr)
|
|
sys.exit(int(os.environ.get("FAKE_RC", "0")))
|
|
"""
|
|
|
|
|
|
@unittest.skipUnless(all(shutil.which(b) for b in ("bash", "flock", "timeout")), "needs bash, flock, timeout")
|
|
class CronScriptTest(unittest.TestCase):
|
|
"""cron/agentd-tick-lxc.sh with a fake agent-access.py: a missing AGENTD_SSH_KEY must not stop the tick."""
|
|
|
|
def setUp(self):
|
|
import re
|
|
self.tmp = tempfile.TemporaryDirectory()
|
|
d = Path(self.tmp.name)
|
|
(d / "client.py").write_text(FAKE_CLIENT)
|
|
text = (ROOT / "cron" / "agentd-tick-lxc.sh").read_text()
|
|
for var, value in (("LOCK", d / "cron.lock"), ("LOG", d / "cron.log"), ("CLIENT", "(python3 -I %s)" % (d / "client.py"))):
|
|
text, n = re.subn(r"(?m)^%s=.*$" % var, "%s=%s" % (var, value), text)
|
|
self.assertEqual(n, 1, var)
|
|
self.script, self.calls, self.log = d / "tick.sh", d / "calls", d / "cron.log"
|
|
self.script.write_text(text)
|
|
|
|
def tearDown(self):
|
|
self.tmp.cleanup()
|
|
|
|
def run_tick(self, **env):
|
|
e = dict(os.environ, FAKE_CALLS=str(self.calls), **env)
|
|
rc = subprocess.run(["bash", str(self.script)], env=e, timeout=60).returncode
|
|
calls = [json.loads(l) for l in self.calls.read_text().splitlines()] if self.calls.exists() else []
|
|
return rc, calls, self.log.read_text() if self.log.exists() else ""
|
|
|
|
def test_normal_tick_passes_the_key(self):
|
|
rc, calls, log = self.run_tick()
|
|
self.assertEqual((rc, len(calls)), (0, 1))
|
|
self.assertIn("--secret-env", calls[0])
|
|
self.assertIn("AGENTD_SSH_KEY=AGENTD_SSH_KEY", calls[0])
|
|
self.assertEqual(calls[0][-3:], ["/opt/plane-agentd/config/agentd.yaml", "tick", "--once"])
|
|
self.assertIn("ticked", log)
|
|
self.assertIn("agentd stderr line", log)
|
|
|
|
def test_missing_vault_key_repeats_the_tick_without_it(self):
|
|
rc, calls, log = self.run_tick(FAKE_MISSING="1")
|
|
self.assertEqual((rc, len(calls)), (0, 2))
|
|
self.assertNotIn("AGENTD_SSH_KEY=AGENTD_SSH_KEY", calls[1])
|
|
self.assertIn("PLANE_OREGON_API_TOKEN=PLANE_OREGON_API_TOKEN", calls[1])
|
|
self.assertIn("AGENTD_SSH_KEY not available from Vault", log)
|
|
self.assertIn("ticked", log)
|
|
|
|
def test_a_failing_tick_is_not_repeated(self):
|
|
rc, calls, _ = self.run_tick(FAKE_RC="1")
|
|
self.assertEqual((rc, len(calls)), (1, 1))
|
|
|
|
def test_lock_held_exits_75(self):
|
|
import fcntl
|
|
with open(Path(self.tmp.name) / "cron.lock", "a") as fh:
|
|
fcntl.flock(fh, fcntl.LOCK_EX | fcntl.LOCK_NB)
|
|
rc, calls, _ = self.run_tick()
|
|
self.assertEqual((rc, calls), (75, []))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|