jj.wang
4f45d368ad
Before the launch call the planned run (run id, terminal, handle, bookkeeping) is saved as a launching marker. If the launch call fails or agentd dies before recording the run, the next tick probes the run directory (exit.txt / launch.pid / pid.txt): evidence -> adopt the run; nothing -> relaunch with the same run id in <run_id>-r2 (-r3 ...); probe unreachable -> keep the marker, which holds the terminal slot. Transports refuse to clean a run directory holding a pid file. A terminal whose relay fails is skipped for the rest of the tick (heartbeat.down) while other terminals carry on.
382 lines
19 KiB
Python
382 lines
19 KiB
Python
"""Terminals behind Warpgate, reached through the server-access relay client:
|
|
|
|
agent-access.py --config <identity> run --target <target> --timeout N -- <argv...>
|
|
|
|
The client shlex-joins argv and the target account's login shell runs it (cmd.exe on Windows); stdin is
|
|
/dev/null, so everything travels in argv. `relay()` is the single place that builds that call.
|
|
|
|
- os linux / darwin -> PosixWarpgateTransport: every operation is one `bash -c '<script>'` relay call. Files go up
|
|
as base64 chunks (`printf %s '<b64>' | base64 -d`), the runner starts detached with `setsid nohup bash run.sh`
|
|
(its own session: pid == pgid == sid), `poll` reads exit.txt, `kill` signals the process group and session.
|
|
- os windows -> WindowsWarpgateTransport: still drives the helper scripts win10-ps.sh (PowerShell
|
|
-EncodedCommand) and win10-put.sh (chunked upload, SHA-256 verified), which hard-code `--target win10`;
|
|
runs are one-shot scheduled tasks (schtasks) so they outlive the SSH session.
|
|
TODO: move Windows onto relay() (`powershell -EncodedCommand` as argv) and drop the helper scripts."""
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
import posixpath
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
|
|
from . import TransportError
|
|
from .local import RUNNER as POSIX_RUNNER, sh_quote
|
|
|
|
DEFAULT_CLIENT = "~/.local/share/server-access-client/agent-access.py"
|
|
|
|
|
|
def relay_argv(cfg, target, argv, timeout):
|
|
"""argv for one `agent-access.py run` call (the vault client and identity come from the `vault` config)."""
|
|
v = cfg.data.get("vault") or {}
|
|
cmd = [sys.executable or "python3", "-I", os.path.expanduser(v.get("client") or DEFAULT_CLIENT)]
|
|
if v.get("config"):
|
|
cmd += ["--config", v["config"]]
|
|
return cmd + ["run", "--target", target, "--timeout", str(int(timeout)), "--"] + list(argv)
|
|
|
|
|
|
def call(cfg, argv, timeout, extra_env=None):
|
|
"""Run a local relay process (client or helper script) with Plane secrets scrubbed from its environment.
|
|
-> (returncode, stdout, stderr); raises TransportError on timeout / launch failure."""
|
|
env = dict(cfg.scrubbed_env(), **(extra_env or {}))
|
|
try:
|
|
out = subprocess.run(argv, capture_output=True, timeout=timeout + 60, env=env, stdin=subprocess.DEVNULL)
|
|
except subprocess.TimeoutExpired:
|
|
raise TransportError("relay timeout after %ss" % timeout) from None
|
|
except OSError as exc:
|
|
raise TransportError("relay launch failed: %s" % exc) from None
|
|
dec = lambda b: (b or b"").decode("utf-8", "replace").replace("\r", "") # noqa: E731
|
|
return out.returncode, dec(out.stdout), dec(out.stderr)
|
|
|
|
|
|
def relay(cfg, target, argv, timeout=150):
|
|
"""Run argv on a Warpgate target. -> (returncode, stdout, stderr)."""
|
|
return call(cfg, relay_argv(cfg, target, argv, timeout), timeout)
|
|
|
|
|
|
RUNNER = r"""$ErrorActionPreference = 'Continue'
|
|
[Console]::OutputEncoding = [Text.Encoding]::UTF8
|
|
$OutputEncoding = [Text.Encoding]::UTF8
|
|
try {{ Set-Location -LiteralPath {cwd} -ErrorAction Stop }} catch {{
|
|
Set-Content -Path '{d}\exit.txt' -Value ('exit=97 at ' + (Get-Date).ToString('s') + ' (cd failed)'); exit 97 }}
|
|
Set-Content -Path '{d}\pid.txt' -Value $PID
|
|
Set-Content -Path '{d}\started.txt' -Value (Get-Date).ToString('s')
|
|
$prompt = Get-Content -Raw -Encoding UTF8 '{d}\prompt.md'
|
|
$prompt | {command} 2> '{d}\stderr.txt' | ForEach-Object {{ Add-Content -Path '{d}\{stdout}' -Value $_ -Encoding UTF8 }}
|
|
Set-Content -Path '{d}\exit.txt' -Value ("exit=" + $LASTEXITCODE + " at " + (Get-Date).ToString('s'))
|
|
"""
|
|
|
|
|
|
def _b64_payload(out, marker):
|
|
m = re.search(marker + r"=([A-Za-z0-9+/=]*)", out)
|
|
if not m:
|
|
return None
|
|
return base64.b64decode(m.group(1)).decode("utf-8", "replace") if m.group(1) else ""
|
|
|
|
|
|
class WindowsWarpgateTransport:
|
|
os = "windows"
|
|
|
|
def __init__(self, cfg, terminal):
|
|
self.cfg, self.terminal = cfg, terminal
|
|
t = cfg.terminals[terminal]
|
|
self.target = t.get("target", terminal)
|
|
helpers = cfg.helpers_dir(terminal)
|
|
self.ps_script, self.put_script = str(helpers / "win10-ps.sh"), str(helpers / "win10-put.sh")
|
|
self.repo = t["repo"]
|
|
|
|
def default_run_root(self):
|
|
return self.join(self.repo, "runtime", "agentd")
|
|
|
|
def join(self, *parts):
|
|
return "\\".join(p.rstrip("\\") for p in parts)
|
|
|
|
def quote(self, s):
|
|
return "'" + str(s).replace("'", "''") + "'"
|
|
|
|
# ---------------------------------------------------------------- relay primitives
|
|
def ps(self, script, timeout=150):
|
|
rc, out, err = call(self.cfg, [self.ps_script, script], timeout, {"WIN10_TIMEOUT": str(timeout)})
|
|
if rc != 0 and not out.strip():
|
|
raise TransportError("relay exit %s: %s" % (rc, err.strip()[-300:]))
|
|
return out
|
|
|
|
def put(self, local, remote):
|
|
with open(local, "rb") as fh:
|
|
sha = hashlib.sha256(fh.read()).hexdigest()
|
|
rc, out, err = call(self.cfg, [self.put_script, local, remote], 540)
|
|
if ("remote sha256 " + sha) not in out:
|
|
raise TransportError("upload of %s failed: %s" % (remote, (out + err).strip()[-300:]))
|
|
|
|
# ---------------------------------------------------------------- file ops
|
|
def mkdir(self, path, clean=False):
|
|
q = self.quote(path)
|
|
script = "New-Item -ItemType Directory -Force -Path %s | Out-Null\n" % q
|
|
if clean: # never wipe a directory a runner has already started in
|
|
script += ("if ((Test-Path -LiteralPath (Join-Path %s 'pid.txt')) -or (Test-Path -LiteralPath (Join-Path %s 'launch.pid'))) "
|
|
"{ 'HASPID'; return }\n" % (q, q))
|
|
script += "Get-ChildItem -LiteralPath %s -File | Remove-Item -Force\n" % q
|
|
out = self.ps(script + "'MKDIR=ok'")
|
|
if "MKDIR=ok" not in out:
|
|
raise TransportError("mkdir %s failed: %s" % (path, out.strip()[-300:]))
|
|
|
|
def probe(self, run_dir, names):
|
|
out = self.ps("$d = %s\n'PROBE=' + ((@(%s) | Where-Object { Test-Path -LiteralPath (Join-Path $d $_) }) -join ',')"
|
|
% (self.quote(run_dir), ",".join(self.quote(n) for n in names)))
|
|
m = re.search(r"PROBE=([^\n]*)", out)
|
|
if not m:
|
|
raise TransportError("probe %s: unexpected relay output: %s" % (run_dir, out.strip()[-200:]))
|
|
return [n for n in m.group(1).strip().split(",") if n]
|
|
|
|
def write_text(self, path, text):
|
|
fd, tmp = tempfile.mkstemp(prefix="agentd-", suffix=".txt")
|
|
try:
|
|
with os.fdopen(fd, "w", encoding="utf-8", newline="") as fh:
|
|
fh.write(text)
|
|
self.put(tmp, path)
|
|
finally:
|
|
os.unlink(tmp)
|
|
|
|
def read_text(self, path):
|
|
q = self.quote(path)
|
|
out = self.ps("if (Test-Path -LiteralPath %s) { 'B64=' + [Convert]::ToBase64String([IO.File]::ReadAllBytes(%s)) }"
|
|
" else { 'MISSING' }" % (q, q), timeout=200)
|
|
if "MISSING" in out and "B64=" not in out:
|
|
return None
|
|
text = _b64_payload(out, "B64")
|
|
if text is None:
|
|
raise TransportError("read %s: unexpected relay output" % path)
|
|
return text.lstrip("")
|
|
|
|
def grep_first(self, path, needle):
|
|
q = self.quote(path)
|
|
out = self.ps("$m = if (Test-Path -LiteralPath %s) { Select-String -LiteralPath %s -SimpleMatch -Pattern %s | Select-Object -First 1 }\n"
|
|
"'B64=' + $(if ($m) { [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($m.Line)) } else { '' })"
|
|
% (q, q, self.quote(needle)))
|
|
return _b64_payload(out, "B64") or None
|
|
|
|
def tail(self, path, n):
|
|
q = self.quote(path)
|
|
out = self.ps("$t = if (Test-Path -LiteralPath %s) { (Get-Content -LiteralPath %s -Encoding UTF8 -Tail %d) -join \"`n\" } else { '' }\n"
|
|
"'B64=' + [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes([string]$t))" % (q, q, n))
|
|
return _b64_payload(out, "B64") or ""
|
|
|
|
# ---------------------------------------------------------------- processes
|
|
def runner(self, run_dir, cwd, command, stdout_name):
|
|
return "run.ps1", RUNNER.format(cwd=self.quote(cwd), d=run_dir.replace("'", "''"), command=command, stdout=stdout_name)
|
|
|
|
def launch_script(self, task, script_path):
|
|
"""PowerShell that (re)creates and runs the one-shot task. Values sit in '...' literals (`'` doubled). The
|
|
script path is wrapped in \\" so that Windows PowerShell 5.1 (legacy native quoting) hands schtasks
|
|
/TR `... -File "<path>"` intact even when the path contains spaces."""
|
|
return ("schtasks /Delete /F /TN '%(t)s' 2>$null | Out-Null\n"
|
|
"$tr = 'powershell.exe -NoProfile -NonInteractive -ExecutionPolicy Bypass -WindowStyle Hidden -File \\\"%(s)s\\\"'\n"
|
|
"schtasks /Create /F /TN '%(t)s' /SC ONCE /ST 00:00 /TR $tr | Out-Null\n"
|
|
"schtasks /Run /TN '%(t)s' | Out-Null\n"
|
|
"'LAUNCHED=' + $LASTEXITCODE" % dict(t=task.replace("'", "''"), s=script_path.replace("'", "''")))
|
|
|
|
def launch(self, task, script_path, run_dir):
|
|
out = self.ps(self.launch_script(task, script_path))
|
|
if "LAUNCHED=0" not in out:
|
|
raise TransportError("schtasks launch of %s failed: %s" % (task, out.strip()[-300:]))
|
|
|
|
def kill(self, task, run_dir):
|
|
self.ps("$p = Join-Path %s 'pid.txt'\n"
|
|
"if (Test-Path $p) { taskkill /T /F /PID (Get-Content $p | Select-Object -First 1) 2>&1 | Out-Null }\n"
|
|
"schtasks /End /TN '%s' 2>&1 | Out-Null\n"
|
|
"schtasks /Delete /F /TN '%s' 2>&1 | Out-Null\n'KILLED=ok'" % (self.quote(run_dir), task, task))
|
|
|
|
def forget(self, task):
|
|
self.ps("schtasks /Delete /F /TN '%s' 2>&1 | Out-Null; 'FORGOT=ok'" % task)
|
|
|
|
def process_count(self, name):
|
|
try:
|
|
out = self.ps("'COUNT=' + @(Get-Process -Name %s -ErrorAction SilentlyContinue).Count" % self.quote(name))
|
|
except TransportError:
|
|
return None
|
|
m = re.search(r"COUNT=(\d+)", out)
|
|
return int(m.group(1)) if m else None
|
|
|
|
|
|
# ==================================================================== Linux / macOS
|
|
DETACH = {
|
|
# bash puts no background job in its own process group (no job control), so setsid(1) does not fork:
|
|
# $! is the runner itself and becomes pid == pgid == sid.
|
|
"linux": "setsid nohup bash %s",
|
|
# macOS has no setsid(1); perl's POSIX::setsid does the same in-process before exec.
|
|
"darwin": "nohup perl -MPOSIX -e 'POSIX::setsid() or die \"setsid: $!\"; exec @ARGV' bash %s",
|
|
}
|
|
|
|
|
|
class PosixWarpgateTransport:
|
|
"""Linux/macOS terminal through the relay. Each primitive is one `bash -c` round trip (~10 s through
|
|
Vault + Warpgate), so `start_run` stages the run directory and launches it in a single call."""
|
|
CHUNK = 60000 # base64 chars per call: one argv element must stay below 128 KiB (MAX_ARG_STRLEN) on both ends
|
|
|
|
def __init__(self, cfg, terminal):
|
|
self.cfg, self.terminal = cfg, terminal
|
|
t = cfg.terminals[terminal]
|
|
self.os = t["os"]
|
|
self.target = t.get("target", terminal)
|
|
self.repo = t["repo"]
|
|
|
|
def default_run_root(self):
|
|
return posixpath.join(self.repo, "runtime", "agentd")
|
|
|
|
def join(self, *parts):
|
|
return posixpath.join(*parts)
|
|
|
|
def quote(self, s):
|
|
return sh_quote(s)
|
|
|
|
# ---------------------------------------------------------------- relay primitive
|
|
def sh(self, script, timeout=150, marker=None):
|
|
"""Run a bash script on the target. With `marker`, its absence in stdout is an error."""
|
|
rc, out, err = relay(self.cfg, self.target, ["bash", "-c", script], timeout)
|
|
if (marker and marker not in out) or (not marker and rc != 0 and not out.strip()):
|
|
raise TransportError("relay %s exit %s: %s" % (self.target, rc, (out + err).strip()[-300:]))
|
|
return out
|
|
|
|
# ---------------------------------------------------------------- file ops
|
|
def mkdir_script(self, path, clean=False):
|
|
q = self.quote(path)
|
|
s = "mkdir -p %s || exit 1\n" % q
|
|
if clean: # never wipe a directory a runner has already started in
|
|
s += "if [ -e %s ] || [ -e %s ]; then echo HASPID; exit 1; fi\n" % (
|
|
self.quote(posixpath.join(path, "pid.txt")), self.quote(posixpath.join(path, "launch.pid")))
|
|
s += "find %s -maxdepth 1 -type f -exec rm -f {} + || exit 1\n" % q
|
|
return s
|
|
|
|
def mkdir(self, path, clean=False):
|
|
self.sh(self.mkdir_script(path, clean) + "echo MKDIR=ok\n", marker="MKDIR=ok")
|
|
|
|
@staticmethod
|
|
def b64(text):
|
|
return base64.b64encode(text.encode("utf-8")).decode("ascii")
|
|
|
|
def put_script(self, path, b64):
|
|
"""Write a base64 payload (alphabet has no quote characters) to path atomically, echo its byte size."""
|
|
q, part = self.quote(path), self.quote(path + ".part")
|
|
return ("printf %%s '%s' | base64 -d > %s && mv -f %s %s || exit 1\n"
|
|
"echo \"WROTE=$(wc -c < %s | tr -d ' ')\"\n" % (b64, part, part, q, q))
|
|
|
|
def write_text(self, path, text):
|
|
data = self.b64(text)
|
|
size = len(text.encode("utf-8"))
|
|
if len(data) <= self.CHUNK:
|
|
out = self.sh(self.put_script(path, data), marker="WROTE=")
|
|
else: # append chunks to <path>.b64, decode on the last call
|
|
q, qb = self.quote(path), self.quote(path + ".b64")
|
|
chunks = [data[i:i + self.CHUNK] for i in range(0, len(data), self.CHUNK)]
|
|
for i, c in enumerate(chunks):
|
|
s = "printf %%s '%s' %s %s || exit 1\n" % (c, ">" if i == 0 else ">>", qb)
|
|
if i < len(chunks) - 1:
|
|
self.sh(s + "echo CHUNK=ok\n", marker="CHUNK=ok")
|
|
else:
|
|
out = self.sh(s + "base64 -d < %s > %s && rm -f %s || exit 1\necho \"WROTE=$(wc -c < %s | tr -d ' ')\"\n"
|
|
% (qb, q, qb, q), marker="WROTE=")
|
|
if re.findall(r"WROTE=(\d+)", out) != [str(size)]:
|
|
raise TransportError("upload of %s: size mismatch (%s)" % (path, out.strip()[-200:]))
|
|
|
|
def read_text(self, path):
|
|
q = self.quote(path)
|
|
out = self.sh("if [ -f %s ]; then printf B64=; base64 < %s | tr -d '\\n'; echo; else echo MISSING; fi" % (q, q),
|
|
timeout=200)
|
|
if "B64=" not in out:
|
|
if "MISSING" in out:
|
|
return None
|
|
raise TransportError("read %s: unexpected relay output" % path)
|
|
return _b64_payload(out, "B64")
|
|
|
|
def probe(self, run_dir, names):
|
|
d = self.quote(run_dir)
|
|
out = self.sh("printf PROBE=; for f in %s; do [ -e %s/\"$f\" ] && printf '%%s,' \"$f\"; done; echo"
|
|
% (" ".join(self.quote(n) for n in names), d), timeout=60, marker="PROBE=")
|
|
return [n for n in re.search(r"PROBE=([^\n]*)", out).group(1).split(",") if n]
|
|
|
|
def _b64_cmd(self, pipeline):
|
|
out = self.sh("printf B64=; { %s; } 2>/dev/null | base64 | tr -d '\\n'; echo" % pipeline)
|
|
text = _b64_payload(out, "B64")
|
|
if text is None:
|
|
raise TransportError("relay %s: unexpected output" % self.target)
|
|
return text
|
|
|
|
def grep_first(self, path, needle):
|
|
line = self._b64_cmd("grep -F -m1 -e %s %s" % (self.quote(needle), self.quote(path)))
|
|
return line.rstrip("\n") or None
|
|
|
|
def tail(self, path, n):
|
|
return self._b64_cmd("tail -n %d %s" % (int(n), self.quote(path))).rstrip("\n")
|
|
|
|
# ---------------------------------------------------------------- processes
|
|
def runner(self, run_dir, cwd, command, stdout_name):
|
|
return "run.sh", POSIX_RUNNER.format(cwd=self.quote(cwd), d=self.quote(run_dir), command=command,
|
|
stdout=stdout_name)
|
|
|
|
def launch_script(self, run_dir, script_name):
|
|
"""Detach the runner from the SSH session; stdio closed so the channel ends at once. $! -> launch.pid
|
|
(kill falls back to it until run.sh has written pid.txt)."""
|
|
d = self.quote(run_dir)
|
|
return ("cd %s || exit 1\n"
|
|
"%s </dev/null >/dev/null 2>&1 &\n"
|
|
"echo $! > launch.pid\n"
|
|
"echo LAUNCHED=$!\n" % (d, DETACH[self.os] % self.quote(script_name)))
|
|
|
|
def launch(self, task, script_path, run_dir):
|
|
self.sh(self.launch_script(run_dir, posixpath.basename(script_path)), marker="LAUNCHED=")
|
|
|
|
def start_run(self, task, run_dir, files, script_name):
|
|
"""mkdir + clean, write `files` {name: text}, launch `script_name` -- one relay call unless a file is
|
|
larger than CHUNK (those go up first through write_text)."""
|
|
big = {n: t for n, t in files.items() if len(self.b64(t)) > self.CHUNK}
|
|
script = self.mkdir_script(run_dir, clean=True)
|
|
if big:
|
|
self.mkdir(run_dir, clean=True)
|
|
for name, text in big.items():
|
|
self.write_text(self.join(run_dir, name), text)
|
|
script = ""
|
|
sizes = {}
|
|
for name, text in files.items():
|
|
if name not in big:
|
|
script += self.put_script(self.join(run_dir, name), self.b64(text))
|
|
sizes[name] = len(text.encode("utf-8"))
|
|
out = self.sh(script + self.launch_script(run_dir, script_name), marker="LAUNCHED=")
|
|
wrote = re.findall(r"WROTE=(\d+)", out)
|
|
if wrote != [str(v) for v in sizes.values()]:
|
|
raise TransportError("staging %s: size mismatch %s != %s" % (run_dir, wrote, list(sizes.values())))
|
|
|
|
def kill_script(self, run_dir):
|
|
"""TERM then KILL the runner's process group and session. Skipped once exit.txt exists (the runner is
|
|
done and its pid may be reused)."""
|
|
d = self.quote(run_dir)
|
|
return ("d=%s\n"
|
|
"if [ -f \"$d/exit.txt\" ]; then echo KILLED=finished; exit 0; fi\n"
|
|
"p=$(head -n1 \"$d/pid.txt\" 2>/dev/null); [ -n \"$p\" ] || p=$(head -n1 \"$d/launch.pid\" 2>/dev/null)\n"
|
|
"case \"$p\" in ''|*[!0-9]*) echo KILLED=nopid; exit 0;; esac\n"
|
|
"kill -TERM -- \"-$p\" 2>/dev/null; pkill -TERM -s \"$p\" 2>/dev/null\n"
|
|
"for i in 1 2 3 4 5; do pgrep -s \"$p\" >/dev/null 2>&1 || break; sleep 1; done\n"
|
|
"kill -KILL -- \"-$p\" 2>/dev/null; pkill -KILL -s \"$p\" 2>/dev/null\n"
|
|
"sleep 1; echo \"KILLED=ok left=$(pgrep -s \"$p\" 2>/dev/null | wc -l | tr -d ' ')\"\n" % d)
|
|
|
|
def kill(self, task, run_dir):
|
|
self.sh(self.kill_script(run_dir), marker="KILLED=")
|
|
|
|
def forget(self, task):
|
|
pass
|
|
|
|
def process_count(self, name):
|
|
try:
|
|
out = self.sh("echo COUNT=$(pgrep -x %s | wc -l | tr -d ' ')" % self.quote(name), timeout=60)
|
|
except TransportError:
|
|
return None
|
|
m = re.search(r"COUNT=(\d+)", out)
|
|
return int(m.group(1)) if m else None
|
|
|
|
|
|
def WarpgateSshTransport(cfg, terminal):
|
|
"""Factory: the transport class for a warpgate-ssh terminal follows its `os`."""
|
|
if cfg.terminals[terminal]["os"] == "windows":
|
|
return WindowsWarpgateTransport(cfg, terminal)
|
|
return PosixWarpgateTransport(cfg, terminal)
|