281 lines
7.7 KiB
Python
281 lines
7.7 KiB
Python
# Adapted from https://github.com/numtide/deploykit
|
|
|
|
import logging
|
|
import os
|
|
import shlex
|
|
import socket
|
|
import subprocess
|
|
import sys
|
|
import types
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from shlex import quote
|
|
from tempfile import TemporaryDirectory
|
|
from typing import Any
|
|
|
|
from clan_lib.errors import ClanError
|
|
|
|
from clan_cli.cmd import CmdOut, RunOpts, run
|
|
from clan_cli.colors import AnsiColor
|
|
from clan_cli.nix import nix_shell
|
|
from clan_cli.ssh.host_key import HostKeyCheck
|
|
|
|
cmdlog = logging.getLogger(__name__)
|
|
|
|
|
|
# Seconds until a message is printed when _run produces no output.
|
|
NO_OUTPUT_TIMEOUT = 20
|
|
|
|
|
|
@dataclass
|
|
class Host:
|
|
host: str
|
|
user: str | None = None
|
|
port: int | None = None
|
|
private_key: Path | None = None
|
|
password: str | None = None
|
|
forward_agent: bool = False
|
|
command_prefix: str | None = None
|
|
host_key_check: HostKeyCheck = HostKeyCheck.ASK
|
|
meta: dict[str, Any] = field(default_factory=dict)
|
|
verbose_ssh: bool = False
|
|
ssh_options: dict[str, str] = field(default_factory=dict)
|
|
tor_socks: bool = False
|
|
|
|
_temp_dir: TemporaryDirectory | None = None
|
|
|
|
def __enter__(self) -> "Host":
|
|
directory = None
|
|
if sys.platform == "darwin" and os.environ.get("TMPDIR", "").startswith(
|
|
"/var/folders/"
|
|
):
|
|
# macOS's tmpdir is too long for unix domain sockets
|
|
directory = "/tmp/"
|
|
|
|
self._temp_dir = TemporaryDirectory(prefix="clan-ssh-", dir=directory)
|
|
return self
|
|
|
|
def __exit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_value: BaseException | None,
|
|
traceback: types.TracebackType | None,
|
|
) -> None:
|
|
try:
|
|
if self._temp_dir:
|
|
self._temp_dir.cleanup()
|
|
except OSError:
|
|
pass
|
|
|
|
def __post_init__(self) -> None:
|
|
if not self.command_prefix:
|
|
self.command_prefix = self.host
|
|
if not self.user:
|
|
self.user = "root"
|
|
|
|
def __str__(self) -> str:
|
|
return self.target
|
|
|
|
@property
|
|
def target(self) -> str:
|
|
return f"{self.user}@{self.host}"
|
|
|
|
@classmethod
|
|
def from_host(cls, host: "Host") -> "Host":
|
|
return cls(
|
|
host=host.host,
|
|
user=host.user,
|
|
port=host.port,
|
|
private_key=host.private_key,
|
|
forward_agent=host.forward_agent,
|
|
command_prefix=host.command_prefix,
|
|
host_key_check=host.host_key_check,
|
|
meta=host.meta.copy(),
|
|
verbose_ssh=host.verbose_ssh,
|
|
ssh_options=host.ssh_options.copy(),
|
|
)
|
|
|
|
def run_local(
|
|
self,
|
|
cmd: list[str],
|
|
opts: RunOpts | None = None,
|
|
extra_env: dict[str, str] | None = None,
|
|
) -> CmdOut:
|
|
"""
|
|
Command to run locally for the host
|
|
"""
|
|
if opts is None:
|
|
opts = RunOpts()
|
|
env = opts.env or os.environ.copy()
|
|
if extra_env:
|
|
env.update(extra_env)
|
|
|
|
displayed_cmd = " ".join(cmd)
|
|
cmdlog.info(
|
|
f"$ {displayed_cmd}",
|
|
extra={
|
|
"command_prefix": self.command_prefix,
|
|
"color": AnsiColor.GREEN.value,
|
|
},
|
|
)
|
|
opts.env = env
|
|
opts.prefix = self.command_prefix
|
|
return run(cmd, opts)
|
|
|
|
def run(
|
|
self,
|
|
cmd: list[str],
|
|
opts: RunOpts | None = None,
|
|
become_root: bool = False,
|
|
extra_env: dict[str, str] | None = None,
|
|
tty: bool = False,
|
|
verbose_ssh: bool = False,
|
|
) -> CmdOut:
|
|
"""
|
|
Command to run on the host via ssh
|
|
"""
|
|
if extra_env is None:
|
|
extra_env = {}
|
|
|
|
if opts is None:
|
|
opts = RunOpts()
|
|
|
|
# If we are not root and we need to become root, prepend sudo
|
|
sudo = ""
|
|
if become_root and self.user != "root":
|
|
sudo = "sudo -- "
|
|
|
|
# Quote all added environment variables
|
|
env_vars = []
|
|
for k, v in extra_env.items():
|
|
env_vars.append(f"{shlex.quote(k)}={shlex.quote(v)}")
|
|
|
|
if opts.prefix is None:
|
|
opts.prefix = self.command_prefix
|
|
# always set needs_user_terminal to True because ssh asks for passwords
|
|
opts.needs_user_terminal = True
|
|
|
|
if opts.cwd is not None:
|
|
msg = "cwd is not supported for remote commands"
|
|
raise ClanError(msg)
|
|
|
|
# Build a pretty command for logging
|
|
displayed_cmd = ""
|
|
export_cmd = ""
|
|
if env_vars:
|
|
export_cmd = f"export {' '.join(env_vars)}; "
|
|
displayed_cmd += export_cmd
|
|
displayed_cmd += " ".join(cmd)
|
|
cmdlog.info(
|
|
f"$ {displayed_cmd}",
|
|
extra={
|
|
"command_prefix": self.command_prefix,
|
|
"color": AnsiColor.GREEN.value,
|
|
},
|
|
)
|
|
|
|
# Build the ssh command
|
|
bash_cmd = export_cmd
|
|
if opts.shell:
|
|
bash_cmd += " ".join(cmd)
|
|
opts.shell = False
|
|
else:
|
|
bash_cmd += 'exec "$@"'
|
|
# FIXME we assume bash to be present here? Should be documented...
|
|
ssh_cmd = [
|
|
*self.ssh_cmd(verbose_ssh=verbose_ssh, tty=tty),
|
|
"--",
|
|
f"{sudo}bash -c {quote(bash_cmd)} -- {' '.join(map(quote, cmd))}",
|
|
]
|
|
|
|
# Run the ssh command
|
|
return run(ssh_cmd, opts)
|
|
|
|
def nix_ssh_env(
|
|
self, env: dict[str, str] | None, local_ssh: bool = True
|
|
) -> dict[str, str]:
|
|
if env is None:
|
|
env = {}
|
|
env["NIX_SSHOPTS"] = " ".join(self.ssh_cmd_opts(local_ssh=local_ssh))
|
|
return env
|
|
|
|
def ssh_cmd_opts(
|
|
self,
|
|
local_ssh: bool = True,
|
|
) -> list[str]:
|
|
ssh_opts = ["-A"] if self.forward_agent else []
|
|
if self.port:
|
|
ssh_opts.extend(["-p", str(self.port)])
|
|
|
|
for k, v in self.ssh_options.items():
|
|
ssh_opts.extend(["-o", f"{k}={shlex.quote(v)}"])
|
|
|
|
ssh_opts.extend(self.host_key_check.to_ssh_opt())
|
|
|
|
if self.private_key:
|
|
ssh_opts.extend(["-i", str(self.private_key)])
|
|
|
|
if local_ssh and self._temp_dir:
|
|
ssh_opts.extend(["-o", "ControlPersist=30m"])
|
|
ssh_opts.extend(
|
|
[
|
|
"-o",
|
|
f"ControlPath={Path(self._temp_dir.name) / 'clan-%h-%p-%r'}",
|
|
]
|
|
)
|
|
ssh_opts.extend(["-o", "ControlMaster=auto"])
|
|
|
|
return ssh_opts
|
|
|
|
def ssh_cmd(
|
|
self,
|
|
verbose_ssh: bool = False,
|
|
tty: bool = False,
|
|
) -> list[str]:
|
|
packages = []
|
|
password_args = []
|
|
if self.password:
|
|
packages.append("sshpass")
|
|
password_args = [
|
|
"sshpass",
|
|
"-p",
|
|
self.password,
|
|
]
|
|
|
|
ssh_opts = self.ssh_cmd_opts()
|
|
if verbose_ssh or self.verbose_ssh:
|
|
ssh_opts.extend(["-v"])
|
|
if tty:
|
|
ssh_opts.extend(["-t"])
|
|
|
|
if self.tor_socks:
|
|
packages.append("netcat")
|
|
ssh_opts.append("-o")
|
|
ssh_opts.append("ProxyCommand=nc -x 127.0.0.1:9050 -X 5 %h %p")
|
|
|
|
cmd = [
|
|
*password_args,
|
|
"ssh",
|
|
self.target,
|
|
*ssh_opts,
|
|
]
|
|
|
|
return nix_shell(packages, cmd)
|
|
|
|
def interactive_ssh(self) -> None:
|
|
subprocess.run(self.ssh_cmd())
|
|
|
|
|
|
def is_ssh_reachable(host: Host) -> bool:
|
|
with socket.socket(
|
|
socket.AF_INET6 if ":" in host.host else socket.AF_INET, socket.SOCK_STREAM
|
|
) as sock:
|
|
sock.settimeout(2)
|
|
try:
|
|
sock.connect((host.host, host.port or 22))
|
|
sock.close()
|
|
except OSError:
|
|
return False
|
|
else:
|
|
return True
|