457 lines
16 KiB
Python
457 lines
16 KiB
Python
"""SSH remote execution environment with ControlMaster connection persistence."""
|
|
|
|
import logging
|
|
import shlex
|
|
import subprocess
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
from tools.environments.base import BaseEnvironment
|
|
from tools.interrupt import is_interrupted
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SSHEnvironment(BaseEnvironment):
|
|
"""Run commands on a remote machine over SSH.
|
|
|
|
Uses SSH ControlMaster for connection persistence so subsequent
|
|
commands are fast. Security benefit: the agent cannot modify its
|
|
own code since execution happens on a separate machine.
|
|
|
|
Foreground commands are interruptible: the local ssh process is killed
|
|
and a remote kill is attempted over the ControlMaster socket.
|
|
|
|
When ``persistent=True``, a single long-lived bash shell is kept alive
|
|
over SSH and state (cwd, env vars, shell variables) persists across
|
|
``execute()`` calls. Output capture uses file-based IPC on the remote
|
|
host (stdout/stderr/exit-code written to temp files, polled via fast
|
|
ControlMaster one-shot reads).
|
|
"""
|
|
|
|
def __init__(self, host: str, user: str, cwd: str = "~",
|
|
timeout: int = 60, port: int = 22, key_path: str = "",
|
|
persistent: bool = False):
|
|
super().__init__(cwd=cwd, timeout=timeout)
|
|
self.host = host
|
|
self.user = user
|
|
self.port = port
|
|
self.key_path = key_path
|
|
self.persistent = persistent
|
|
|
|
self.control_dir = Path(tempfile.gettempdir()) / "hermes-ssh"
|
|
self.control_dir.mkdir(parents=True, exist_ok=True)
|
|
self.control_socket = self.control_dir / f"{user}@{host}:{port}.sock"
|
|
self._establish_connection()
|
|
|
|
# Persistent shell state
|
|
self._shell_proc: subprocess.Popen | None = None
|
|
self._shell_lock = threading.Lock()
|
|
self._shell_alive = False
|
|
self._session_id: str = ""
|
|
self._remote_stdout: str = ""
|
|
self._remote_stderr: str = ""
|
|
self._remote_status: str = ""
|
|
self._remote_cwd: str = ""
|
|
self._remote_pid: str = ""
|
|
self._remote_shell_pid: int | None = None
|
|
|
|
if self.persistent:
|
|
self._start_persistent_shell()
|
|
|
|
def _build_ssh_command(self, extra_args: list = None) -> list:
|
|
cmd = ["ssh"]
|
|
cmd.extend(["-o", f"ControlPath={self.control_socket}"])
|
|
cmd.extend(["-o", "ControlMaster=auto"])
|
|
cmd.extend(["-o", "ControlPersist=300"])
|
|
cmd.extend(["-o", "BatchMode=yes"])
|
|
cmd.extend(["-o", "StrictHostKeyChecking=accept-new"])
|
|
cmd.extend(["-o", "ConnectTimeout=10"])
|
|
if self.port != 22:
|
|
cmd.extend(["-p", str(self.port)])
|
|
if self.key_path:
|
|
cmd.extend(["-i", self.key_path])
|
|
if extra_args:
|
|
cmd.extend(extra_args)
|
|
cmd.append(f"{self.user}@{self.host}")
|
|
return cmd
|
|
|
|
def _establish_connection(self):
|
|
cmd = self._build_ssh_command()
|
|
cmd.append("echo 'SSH connection established'")
|
|
try:
|
|
result = subprocess.run(cmd, capture_output=True, text=True, timeout=15)
|
|
if result.returncode != 0:
|
|
error_msg = result.stderr.strip() or result.stdout.strip()
|
|
raise RuntimeError(f"SSH connection failed: {error_msg}")
|
|
except subprocess.TimeoutExpired:
|
|
raise RuntimeError(f"SSH connection to {self.user}@{self.host} timed out")
|
|
|
|
# ------------------------------------------------------------------
|
|
# Persistent shell management
|
|
# ------------------------------------------------------------------
|
|
|
|
def _start_persistent_shell(self):
|
|
"""Spawn a long-lived bash shell over SSH."""
|
|
self._session_id = uuid.uuid4().hex[:12]
|
|
prefix = f"/tmp/hermes-ssh-{self._session_id}"
|
|
self._remote_stdout = f"{prefix}-stdout"
|
|
self._remote_stderr = f"{prefix}-stderr"
|
|
self._remote_status = f"{prefix}-status"
|
|
self._remote_cwd = f"{prefix}-cwd"
|
|
self._remote_pid = f"{prefix}-pid"
|
|
|
|
cmd = self._build_ssh_command()
|
|
cmd.append("bash -l")
|
|
|
|
self._shell_proc = subprocess.Popen(
|
|
cmd,
|
|
stdin=subprocess.PIPE,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
text=True,
|
|
)
|
|
self._shell_alive = True
|
|
|
|
# Start daemon thread to drain stdout/stderr and detect shell death
|
|
self._drain_thread = threading.Thread(
|
|
target=self._drain_shell_output, daemon=True
|
|
)
|
|
self._drain_thread.start()
|
|
|
|
# Initialize remote temp files and capture shell PID
|
|
init_script = (
|
|
f"touch {self._remote_stdout} {self._remote_stderr} "
|
|
f"{self._remote_status} {self._remote_cwd} {self._remote_pid}\n"
|
|
f"echo $$ > {self._remote_pid}\n"
|
|
f"pwd > {self._remote_cwd}\n"
|
|
)
|
|
self._send_to_shell(init_script)
|
|
|
|
# Give shell time to initialize and write PID file
|
|
time.sleep(0.3)
|
|
|
|
# Read the remote shell PID
|
|
pid_str = self._read_remote_file(self._remote_pid).strip()
|
|
if pid_str.isdigit():
|
|
self._remote_shell_pid = int(pid_str)
|
|
logger.info("Persistent shell started (session=%s, pid=%d)",
|
|
self._session_id, self._remote_shell_pid)
|
|
else:
|
|
logger.warning("Could not read persistent shell PID (got %r)", pid_str)
|
|
self._remote_shell_pid = None
|
|
|
|
# Update cwd from what the shell reports
|
|
remote_cwd = self._read_remote_file(self._remote_cwd).strip()
|
|
if remote_cwd:
|
|
self.cwd = remote_cwd
|
|
|
|
def _drain_shell_output(self):
|
|
"""Drain the shell's stdout/stderr to prevent pipe deadlock.
|
|
|
|
Also detects when the shell process dies.
|
|
"""
|
|
try:
|
|
for _ in self._shell_proc.stdout:
|
|
pass # Discard — real output goes to temp files
|
|
except Exception:
|
|
pass
|
|
self._shell_alive = False
|
|
|
|
def _send_to_shell(self, text: str):
|
|
"""Write text to the persistent shell's stdin."""
|
|
if not self._shell_alive or self._shell_proc is None:
|
|
return
|
|
try:
|
|
self._shell_proc.stdin.write(text)
|
|
self._shell_proc.stdin.flush()
|
|
except (BrokenPipeError, OSError):
|
|
self._shell_alive = False
|
|
|
|
def _read_remote_file(self, path: str) -> str:
|
|
"""Read a file on the remote host via a one-shot SSH command.
|
|
|
|
Uses ControlMaster so this is very fast (~5ms on LAN).
|
|
"""
|
|
cmd = self._build_ssh_command()
|
|
cmd.append(f"cat {path} 2>/dev/null")
|
|
try:
|
|
result = subprocess.run(
|
|
cmd, capture_output=True, text=True, timeout=10
|
|
)
|
|
return result.stdout
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
return ""
|
|
|
|
def _kill_shell_children(self):
|
|
"""Kill children of the persistent shell (the running command),
|
|
but not the shell itself."""
|
|
if self._remote_shell_pid is None:
|
|
return
|
|
cmd = self._build_ssh_command()
|
|
cmd.append(f"pkill -P {self._remote_shell_pid} 2>/dev/null; true")
|
|
try:
|
|
subprocess.run(cmd, capture_output=True, timeout=5)
|
|
except (subprocess.TimeoutExpired, OSError):
|
|
pass
|
|
|
|
def _execute_persistent(self, command: str, cwd: str, *,
|
|
timeout: int | None = None,
|
|
stdin_data: str | None = None) -> dict:
|
|
"""Execute a command in the persistent shell."""
|
|
# If shell is dead, restart it
|
|
if not self._shell_alive:
|
|
logger.info("Persistent shell died, restarting...")
|
|
self._start_persistent_shell()
|
|
|
|
exec_command, sudo_stdin = self._prepare_command(command)
|
|
effective_timeout = timeout or self.timeout
|
|
|
|
# Fall back to one-shot for commands needing piped stdin
|
|
if stdin_data or sudo_stdin:
|
|
return self._execute_oneshot(
|
|
command, cwd, timeout=timeout, stdin_data=stdin_data
|
|
)
|
|
|
|
with self._shell_lock:
|
|
return self._execute_persistent_locked(
|
|
exec_command, cwd, effective_timeout
|
|
)
|
|
|
|
def _execute_persistent_locked(self, command: str, cwd: str,
|
|
timeout: int) -> dict:
|
|
"""Inner persistent execution — caller must hold _shell_lock."""
|
|
work_dir = cwd or self.cwd
|
|
|
|
# Truncate temp files
|
|
truncate = (
|
|
f": > {self._remote_stdout}\n"
|
|
f": > {self._remote_stderr}\n"
|
|
f": > {self._remote_status}\n"
|
|
)
|
|
self._send_to_shell(truncate)
|
|
|
|
# Escape command for eval — use single quotes with proper escaping
|
|
escaped = command.replace("'", "'\\''")
|
|
|
|
# Send the IPC script
|
|
ipc_script = (
|
|
f"cd {shlex.quote(work_dir)}\n"
|
|
f"eval '{escaped}' < /dev/null > {self._remote_stdout} 2> {self._remote_stderr}\n"
|
|
f"__EC=$?\n"
|
|
f"pwd > {self._remote_cwd}\n"
|
|
f"echo $__EC > {self._remote_status}\n"
|
|
)
|
|
self._send_to_shell(ipc_script)
|
|
|
|
# Poll the status file
|
|
deadline = time.monotonic() + timeout
|
|
poll_interval = 0.05 # 50ms
|
|
|
|
while True:
|
|
if is_interrupted():
|
|
self._kill_shell_children()
|
|
stdout = self._read_remote_file(self._remote_stdout)
|
|
stderr = self._read_remote_file(self._remote_stderr)
|
|
output = self._merge_output(stdout, stderr)
|
|
return {
|
|
"output": output + "\n[Command interrupted]",
|
|
"returncode": 130,
|
|
}
|
|
|
|
if time.monotonic() > deadline:
|
|
self._kill_shell_children()
|
|
stdout = self._read_remote_file(self._remote_stdout)
|
|
stderr = self._read_remote_file(self._remote_stderr)
|
|
output = self._merge_output(stdout, stderr)
|
|
if output:
|
|
return {
|
|
"output": output + f"\n[Command timed out after {timeout}s]",
|
|
"returncode": 124,
|
|
}
|
|
return self._timeout_result(timeout)
|
|
|
|
if not self._shell_alive:
|
|
return {
|
|
"output": "Persistent shell died during execution",
|
|
"returncode": 1,
|
|
}
|
|
|
|
# Check if status file has content (command is done)
|
|
status_content = self._read_remote_file(self._remote_status).strip()
|
|
if status_content:
|
|
break
|
|
|
|
time.sleep(poll_interval)
|
|
|
|
# Read results
|
|
stdout = self._read_remote_file(self._remote_stdout)
|
|
stderr = self._read_remote_file(self._remote_stderr)
|
|
exit_code_str = status_content
|
|
new_cwd = self._read_remote_file(self._remote_cwd).strip()
|
|
|
|
# Parse exit code
|
|
try:
|
|
exit_code = int(exit_code_str)
|
|
except ValueError:
|
|
exit_code = 1
|
|
|
|
# Update cwd
|
|
if new_cwd:
|
|
self.cwd = new_cwd
|
|
|
|
output = self._merge_output(stdout, stderr)
|
|
return {"output": output, "returncode": exit_code}
|
|
|
|
@staticmethod
|
|
def _merge_output(stdout: str, stderr: str) -> str:
|
|
"""Combine stdout and stderr into a single output string."""
|
|
parts = []
|
|
if stdout.strip():
|
|
parts.append(stdout.rstrip("\n"))
|
|
if stderr.strip():
|
|
parts.append(stderr.rstrip("\n"))
|
|
return "\n".join(parts)
|
|
|
|
# ------------------------------------------------------------------
|
|
# One-shot execution (original behavior)
|
|
# ------------------------------------------------------------------
|
|
|
|
def _execute_oneshot(self, command: str, cwd: str = "", *,
|
|
timeout: int | None = None,
|
|
stdin_data: str | None = None) -> dict:
|
|
"""Execute a command via a fresh one-shot SSH invocation."""
|
|
work_dir = cwd or self.cwd
|
|
exec_command, sudo_stdin = self._prepare_command(command)
|
|
wrapped = f'cd {work_dir} && {exec_command}'
|
|
effective_timeout = timeout or self.timeout
|
|
|
|
# Merge sudo password (if any) with caller-supplied stdin_data.
|
|
if sudo_stdin is not None and stdin_data is not None:
|
|
effective_stdin = sudo_stdin + stdin_data
|
|
elif sudo_stdin is not None:
|
|
effective_stdin = sudo_stdin
|
|
else:
|
|
effective_stdin = stdin_data
|
|
|
|
cmd = self._build_ssh_command()
|
|
cmd.extend(["bash", "-c", wrapped])
|
|
|
|
try:
|
|
kwargs = self._build_run_kwargs(timeout, effective_stdin)
|
|
# Remove timeout from kwargs -- we handle it in the poll loop
|
|
kwargs.pop("timeout", None)
|
|
|
|
_output_chunks = []
|
|
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
stdin=subprocess.PIPE if effective_stdin else subprocess.DEVNULL,
|
|
text=True,
|
|
)
|
|
|
|
if effective_stdin:
|
|
try:
|
|
proc.stdin.write(effective_stdin)
|
|
proc.stdin.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def _drain():
|
|
try:
|
|
for line in proc.stdout:
|
|
_output_chunks.append(line)
|
|
except Exception:
|
|
pass
|
|
|
|
reader = threading.Thread(target=_drain, daemon=True)
|
|
reader.start()
|
|
deadline = time.monotonic() + effective_timeout
|
|
|
|
while proc.poll() is None:
|
|
if is_interrupted():
|
|
proc.terminate()
|
|
try:
|
|
proc.wait(timeout=1)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
reader.join(timeout=2)
|
|
return {
|
|
"output": "".join(_output_chunks) + "\n[Command interrupted]",
|
|
"returncode": 130,
|
|
}
|
|
if time.monotonic() > deadline:
|
|
proc.kill()
|
|
reader.join(timeout=2)
|
|
return self._timeout_result(effective_timeout)
|
|
time.sleep(0.2)
|
|
|
|
reader.join(timeout=5)
|
|
return {"output": "".join(_output_chunks), "returncode": proc.returncode}
|
|
|
|
except Exception as e:
|
|
return {"output": f"SSH execution error: {str(e)}", "returncode": 1}
|
|
|
|
# ------------------------------------------------------------------
|
|
# Public interface
|
|
# ------------------------------------------------------------------
|
|
|
|
def execute(self, command: str, cwd: str = "", *,
|
|
timeout: int | None = None,
|
|
stdin_data: str | None = None) -> dict:
|
|
if self.persistent:
|
|
return self._execute_persistent(
|
|
command, cwd, timeout=timeout, stdin_data=stdin_data
|
|
)
|
|
return self._execute_oneshot(
|
|
command, cwd, timeout=timeout, stdin_data=stdin_data
|
|
)
|
|
|
|
def cleanup(self):
|
|
# Persistent shell teardown
|
|
if self.persistent and self._shell_proc is not None:
|
|
# Remove remote temp files
|
|
if self._session_id:
|
|
try:
|
|
cmd = self._build_ssh_command()
|
|
cmd.append(
|
|
f"rm -f /tmp/hermes-ssh-{self._session_id}-*"
|
|
)
|
|
subprocess.run(cmd, capture_output=True, timeout=5)
|
|
except (OSError, subprocess.SubprocessError):
|
|
pass
|
|
|
|
# Close the shell
|
|
try:
|
|
self._shell_proc.stdin.close()
|
|
except Exception:
|
|
pass
|
|
try:
|
|
self._shell_proc.terminate()
|
|
self._shell_proc.wait(timeout=3)
|
|
except Exception:
|
|
try:
|
|
self._shell_proc.kill()
|
|
except Exception:
|
|
pass
|
|
self._shell_alive = False
|
|
self._shell_proc = None
|
|
|
|
# ControlMaster cleanup
|
|
if self.control_socket.exists():
|
|
try:
|
|
cmd = ["ssh", "-o", f"ControlPath={self.control_socket}",
|
|
"-O", "exit", f"{self.user}@{self.host}"]
|
|
subprocess.run(cmd, capture_output=True, timeout=5)
|
|
except (OSError, subprocess.SubprocessError):
|
|
pass
|
|
try:
|
|
self.control_socket.unlink()
|
|
except OSError:
|
|
pass
|