jj.wang
560 lines
28 KiB
Python
560 lines
28 KiB
Python
"""Remote terminals (Windows and Linux/macOS) and the Warpgate channel.
|
|
|
|
The transport classes build the scripts; a *channel* carries them (design §9):
|
|
- WarpgateChannel (`transport: warpgate-ssh`, also the direct-ssh fallback): 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 (agent-access.py run_remote uses stdin=DEVNULL), so everything travels in argv. `relay()` is the single
|
|
place that builds that call.
|
|
- DirectChannel (`transport: direct-ssh`, agentd/transports/direct_ssh.py): OpenSSH over Tailscale; the script goes
|
|
on stdin, so every operation (including a launch with its prompt) is one call.
|
|
|
|
- os linux / darwin -> PosixWarpgateTransport: every operation is one `bash -c '<script>'` call. On the relay files go
|
|
up as base64 chunks (`printf %s '<b64>' | base64 -d`), on the direct channel in the same single call. The runner
|
|
starts detached with `setsid nohup bash run.sh` (its own session: pid == pgid == sid), `poll` reads exit.txt,
|
|
`fetch` returns everything `collect` needs in one call, `kill` signals the process group and session.
|
|
- os windows -> WindowsWarpgateTransport: PowerShell scripts; runs are one-shot scheduled tasks (schtasks) so they
|
|
outlive the SSH session. On the relay it still drives the helper scripts win10-ps.sh (PowerShell -EncodedCommand,
|
|
argv-size bound) and win10-put.sh (chunked upload, SHA-256 verified), which hard-code `--target win10`; on the
|
|
direct channel mkdir + clean + prompt.md + run.ps1 + schtasks is one call.
|
|
The class names are historical: the channel, not the class, decides the path."""
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
import posixpath
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
|
|
from . import NULL_LOG, TransportError, Unreached, fetch_each, operation, parse_fetch
|
|
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)
|
|
|
|
|
|
PS_ARGV_MAX = 7600 # win10-ps.sh refuses longer -EncodedCommand values (cmd.exe line limit 8191)
|
|
PS_PRELUDE = 120 # characters win10-ps.sh puts in front of every script
|
|
|
|
|
|
class WarpgateChannel:
|
|
"""Calls through agent-access.py run (argv only). Each call is logged with its duration."""
|
|
stdin = False
|
|
|
|
def __init__(self, cfg, terminal, log=None, label="warpgate"):
|
|
self.cfg, self.terminal, self.log, self.via = cfg, terminal, log or NULL_LOG, label
|
|
t = cfg.terminals[terminal]
|
|
self.target = t.get("target", terminal)
|
|
self.name = "relay %s" % self.target
|
|
if t["os"] == "windows":
|
|
helpers = cfg.helpers_dir(terminal)
|
|
self.ps_script, self.put_script = str(helpers / "win10-ps.sh"), str(helpers / "win10-put.sh")
|
|
|
|
def timed(self, op, fn, *args):
|
|
t0 = time.monotonic()
|
|
try:
|
|
res = fn(*args)
|
|
except TransportError:
|
|
self.log.call(self.terminal, self.via, op, (time.monotonic() - t0) * 1000, "error")
|
|
raise
|
|
self.log.call(self.terminal, self.via, op, (time.monotonic() - t0) * 1000, res[0])
|
|
return res
|
|
|
|
def fits(self, script):
|
|
"""Whether win10-ps.sh accepts the script (its -EncodedCommand stays under PS_ARGV_MAX)."""
|
|
units = len(script.encode("utf-16-le")) // 2 + PS_PRELUDE
|
|
return (2 * units + 2) // 3 * 4 <= PS_ARGV_MAX
|
|
|
|
def ps(self, script, timeout, op=None):
|
|
rc, out, err = self.timed(op, 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, op=None):
|
|
with open(local, "rb") as fh:
|
|
sha = hashlib.sha256(fh.read()).hexdigest()
|
|
rc, out, err = self.timed(op, 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:]))
|
|
|
|
def sh(self, script, timeout, op=None):
|
|
return self.timed(op, relay, self.cfg, self.target, ["bash", "-c", script], 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'))
|
|
"""
|
|
|
|
|
|
# Windows fetch helpers: F<i>=<base64> per result, M<i> = file missing (parse_fetch). Names must not collide with
|
|
# PowerShell's built-in aliases (aliases win over functions: `r` is Invoke-History).
|
|
PS_FETCH = r"""function AdRead($i,$n){$p=Join-Path $d $n; if(Test-Path -LiteralPath $p){"F$i="+[Convert]::ToBase64String([IO.File]::ReadAllBytes($p))}else{"M$i"}}
|
|
function AdOut($i,$t){"F$i="+[Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes([string]$t))}
|
|
function AdGrep($i,$n,$x){$p=Join-Path $d $n; $m=$null; if(Test-Path -LiteralPath $p){$m=Select-String -LiteralPath $p -SimpleMatch -Pattern $x|Select-Object -First 1}; if($m){AdOut $i $m.Line}else{"F$i="}}
|
|
function AdTail($i,$n,$k){$p=Join-Path $d $n; $t=''; if(Test-Path -LiteralPath $p){$t=(Get-Content -LiteralPath $p -Encoding UTF8 -Tail $k) -join "`n"}; AdOut $i $t}
|
|
"""
|
|
|
|
|
|
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, channel=None, log=None, label="warpgate"):
|
|
self.cfg, self.terminal = cfg, terminal
|
|
t = cfg.terminals[terminal]
|
|
self.target = t.get("target", terminal)
|
|
self.channel = channel or WarpgateChannel(cfg, terminal, log, label)
|
|
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("'", "''") + "'"
|
|
|
|
# ---------------------------------------------------------------- channel primitives
|
|
def ps(self, script, timeout=150):
|
|
return self.channel.ps(script, timeout, getattr(self, "_op", None))
|
|
|
|
def put(self, local, remote):
|
|
self.channel.put(local, remote, getattr(self, "_op", None))
|
|
|
|
@operation("ping")
|
|
def ping(self):
|
|
out = self.ps("'PONG=' + $PSVersionTable.PSVersion.ToString() + ' ' + $env:COMPUTERNAME", timeout=60)
|
|
m = re.search(r"PONG=[^\n]*", out)
|
|
if not m:
|
|
raise TransportError("ping %s: unexpected output: %s" % (self.terminal, out.strip()[-200:]))
|
|
return m.group(0)
|
|
|
|
# ---------------------------------------------------------------- file ops
|
|
def mkdir_script(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
|
|
return script
|
|
|
|
@operation("mkdir")
|
|
def mkdir(self, path, clean=False):
|
|
out = self.ps(self.mkdir_script(path, clean) + "'MKDIR=ok'")
|
|
if "MKDIR=ok" not in out:
|
|
raise TransportError("mkdir %s failed: %s" % (path, out.strip()[-300:]))
|
|
|
|
@operation("probe")
|
|
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_script(self, path, data):
|
|
"""Write bytes exactly as given (UTF-8 without BOM for our texts), echo their SHA-256."""
|
|
return ("$p = %s\n[IO.File]::WriteAllBytes($p, [Convert]::FromBase64String('%s'))\n"
|
|
"'WROTE=' + (Get-FileHash -LiteralPath $p -Algorithm SHA256).Hash.ToLower()\n"
|
|
% (self.quote(path), base64.b64encode(data).decode("ascii")))
|
|
|
|
@operation("write")
|
|
def write_text(self, path, text):
|
|
if self.channel.stdin:
|
|
data = text.encode("utf-8")
|
|
out = self.ps(self.write_script(path, data))
|
|
if re.findall(r"WROTE=([0-9a-f]{64})", out) != [hashlib.sha256(data).hexdigest()]:
|
|
raise TransportError("upload of %s failed: %s" % (path, out.strip()[-300:]))
|
|
return
|
|
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)
|
|
|
|
@operation("read")
|
|
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("\ufeff")
|
|
|
|
@operation("grep")
|
|
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
|
|
|
|
@operation("tail")
|
|
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 ""
|
|
|
|
def fetch_script(self, run_dir, ops):
|
|
lines = ["$d = %s" % self.quote(run_dir), PS_FETCH.rstrip("\n")]
|
|
for i, o in enumerate(ops):
|
|
fn = {"read": "AdRead", "grep": "AdGrep", "tail": "AdTail"}[o[0]]
|
|
arg = "" if o[0] == "read" else " " + (self.quote(o[2]) if o[0] == "grep" else str(int(o[2])))
|
|
lines.append("%s %d %s%s" % (fn, i, self.quote(o[1]), arg))
|
|
return "\n".join(lines) + "\n"
|
|
|
|
@operation("fetch")
|
|
def fetch(self, run_dir, ops):
|
|
"""Several reads of one run directory in one call (collect)."""
|
|
script = self.fetch_script(run_dir, ops)
|
|
if not self.channel.stdin and not self.channel.fits(script):
|
|
return fetch_each(self, run_dir, ops)
|
|
res = parse_fetch(self.ps(script, timeout=200), ops, run_dir)
|
|
return [r.lstrip("\ufeff") if o[0] == "read" and r else r for o, r in zip(ops, res)]
|
|
|
|
# ---------------------------------------------------------------- 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("'", "''")))
|
|
|
|
@operation("launch")
|
|
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 start_script(self, task, run_dir, files, script_name):
|
|
"""mkdir + clean (HASPID refusal) + every file + schtasks, as one script. -> (script, [sha256 per file])"""
|
|
script, shas = "$ErrorActionPreference = 'Stop'\n" + self.mkdir_script(run_dir, clean=True), []
|
|
for name, text in files.items():
|
|
data = text.encode("utf-8")
|
|
script += self.write_script(self.join(run_dir, name), data)
|
|
shas.append(hashlib.sha256(data).hexdigest())
|
|
script += "$ErrorActionPreference = 'Continue'\n" + self.launch_script(task, self.join(run_dir, script_name))
|
|
return script, shas
|
|
|
|
@operation("start")
|
|
def start_run(self, task, run_dir, files, script_name):
|
|
"""Stage the run directory and start the task: one call on the direct channel; on the relay (argv only)
|
|
mkdir, chunked uploads and the launch go one by one."""
|
|
if not self.channel.stdin:
|
|
self.mkdir(run_dir, clean=True)
|
|
for name, text in files.items():
|
|
self.write_text(self.join(run_dir, name), text)
|
|
return self.launch(task, self.join(run_dir, script_name), run_dir)
|
|
script, shas = self.start_script(task, run_dir, files, script_name)
|
|
out = self.ps(script)
|
|
if re.search(r"^HASPID$", out, re.M):
|
|
raise TransportError("refusing to clean %s: it holds a pid file (HASPID)" % run_dir)
|
|
wrote = re.findall(r"WROTE=([0-9a-f]{64})", out)
|
|
if wrote != shas:
|
|
raise TransportError("staging %s: hash mismatch (%d/%d files): %s" % (run_dir, len(wrote), len(shas),
|
|
out.strip()[-300:]))
|
|
if "LAUNCHED=0" not in out:
|
|
raise TransportError("schtasks launch of %s failed: %s" % (task, out.strip()[-300:]))
|
|
|
|
@operation("kill")
|
|
def kill(self, task, run_dir):
|
|
t = task.replace("'", "''")
|
|
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), t, t))
|
|
|
|
@operation("forget")
|
|
def forget(self, task):
|
|
self.ps("schtasks /Delete /F /TN '%s' 2>&1 | Out-Null; 'FORGOT=ok'" % task.replace("'", "''"))
|
|
|
|
@operation("busy")
|
|
def process_count(self, name):
|
|
try:
|
|
out = self.ps("'COUNT=' + @(Get-Process -Name %s -ErrorAction SilentlyContinue).Count" % self.quote(name))
|
|
except Unreached:
|
|
raise # a direct-ssh terminal retries this on its fallback
|
|
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. Each primitive is one `bash -c` round trip (~10 s through Vault + Warpgate, well under
|
|
a second on a reused direct connection), so `start_run` stages the run directory and launches it in a single
|
|
call and `fetch` reads everything `collect` needs in one."""
|
|
CHUNK = 60000 # relay only: base64 chars per call; one argv element must stay below 128 KiB (MAX_ARG_STRLEN)
|
|
|
|
def __init__(self, cfg, terminal, channel=None, log=None, label="warpgate"):
|
|
self.cfg, self.terminal = cfg, terminal
|
|
t = cfg.terminals[terminal]
|
|
self.os = t["os"]
|
|
self.target = t.get("target", terminal)
|
|
self.channel = channel or WarpgateChannel(cfg, terminal, log, label)
|
|
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)
|
|
|
|
# ---------------------------------------------------------------- channel 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 = self.channel.sh(script, timeout, getattr(self, "_op", None))
|
|
if (marker and marker not in out) or (not marker and rc != 0 and not out.strip()):
|
|
raise TransportError("%s exit %s: %s" % (self.channel.name, rc, (out + err).strip()[-300:]))
|
|
return out
|
|
|
|
@operation("ping")
|
|
def ping(self):
|
|
out = self.sh("echo \"PONG=$(uname -sr) $(uname -n)\"", timeout=60, marker="PONG=")
|
|
return re.search(r"PONG=[^\n]*", out).group(0)
|
|
|
|
# ---------------------------------------------------------------- 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
|
|
|
|
@operation("mkdir")
|
|
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))
|
|
|
|
@operation("write")
|
|
def write_text(self, path, text):
|
|
data = self.b64(text)
|
|
size = len(text.encode("utf-8"))
|
|
if len(data) <= self.CHUNK or self.channel.stdin: # stdin carries the script: no argv limit
|
|
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:]))
|
|
|
|
@operation("read")
|
|
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")
|
|
|
|
@operation("probe")
|
|
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
|
|
|
|
@operation("grep")
|
|
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
|
|
|
|
@operation("tail")
|
|
def tail(self, path, n):
|
|
return self._b64_cmd("tail -n %d %s" % (int(n), self.quote(path))).rstrip("\n")
|
|
|
|
def fetch_script(self, run_dir, ops):
|
|
s = "d=%s\n" % self.quote(run_dir)
|
|
for i, o in enumerate(ops):
|
|
f = "\"$d\"/%s" % self.quote(o[1])
|
|
if o[0] == "read":
|
|
s += ("if [ -f %s ]; then printf F%d=; base64 < %s | tr -d '\\n'; echo; else echo M%d; fi\n"
|
|
% (f, i, f, i))
|
|
else:
|
|
cmd = ("grep -F -m1 -e %s %s" % (self.quote(o[2]), f) if o[0] == "grep"
|
|
else "tail -n %d %s" % (int(o[2]), f))
|
|
s += "printf F%d=; { %s; } 2>/dev/null | base64 | tr -d '\\n'; echo\n" % (i, cmd)
|
|
return s + "echo FETCH=ok\n"
|
|
|
|
@operation("fetch")
|
|
def fetch(self, run_dir, ops):
|
|
"""Several reads of one run directory in one call (collect)."""
|
|
return parse_fetch(self.sh(self.fetch_script(run_dir, ops), timeout=200, marker="FETCH=ok"), ops, run_dir)
|
|
|
|
# ---------------------------------------------------------------- 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)))
|
|
|
|
@operation("launch")
|
|
def launch(self, task, script_path, run_dir):
|
|
self.sh(self.launch_script(run_dir, posixpath.basename(script_path)), marker="LAUNCHED=")
|
|
|
|
@operation("start")
|
|
def start_run(self, task, run_dir, files, script_name):
|
|
"""mkdir + clean, write `files` {name: text}, launch `script_name` -- one call; on the relay a file larger
|
|
than CHUNK goes up first through write_text."""
|
|
big = {} if self.channel.stdin else {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)
|
|
|
|
@operation("kill")
|
|
def kill(self, task, run_dir):
|
|
self.sh(self.kill_script(run_dir), marker="KILLED=")
|
|
|
|
def forget(self, task):
|
|
pass
|
|
|
|
@operation("busy")
|
|
def process_count(self, name):
|
|
try:
|
|
out = self.sh("echo COUNT=$(pgrep -x %s | wc -l | tr -d ' ')" % self.quote(name), timeout=60)
|
|
except Unreached:
|
|
raise # a direct-ssh terminal retries this on its fallback
|
|
except TransportError:
|
|
return None
|
|
m = re.search(r"COUNT=(\d+)", out)
|
|
return int(m.group(1)) if m else None
|
|
|
|
|
|
def WarpgateSshTransport(cfg, terminal, channel=None, log=None, label="warpgate"):
|
|
"""Factory: the transport class for a remote terminal follows its `os`; the channel defaults to Warpgate."""
|
|
cls = WindowsWarpgateTransport if cfg.terminals[terminal]["os"] == "windows" else PosixWarpgateTransport
|
|
return cls(cfg, terminal, channel=channel, log=log, label=label)
|