573 lines
20 KiB
Python
573 lines
20 KiB
Python
"""
|
|
SSH client wrapper — connect, exec, sftp, key management via paramiko.
|
|
"""
|
|
|
|
import os
|
|
import platform
|
|
import socket
|
|
import threading
|
|
import time
|
|
import paramiko
|
|
from core.logger import log
|
|
|
|
|
|
def _create_bound_socket(bind_ip: str, hostname: str, port: int, timeout: int) -> socket.socket:
|
|
"""Create a TCP socket bound to a specific local IP address."""
|
|
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
sock.settimeout(timeout)
|
|
sock.bind((bind_ip, 0))
|
|
sock.connect((hostname, port))
|
|
return sock
|
|
|
|
|
|
def _connect_client(server: dict, key_path: str, timeout: int = 15) -> paramiko.SSHClient:
|
|
"""Create and authenticate a paramiko SSHClient. Shared by SSHClientWrapper and ShellSession."""
|
|
client = paramiko.SSHClient()
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
|
|
hostname = server["ip"]
|
|
port = server.get("port", 22)
|
|
bind_ip = server.get("bind_interface")
|
|
|
|
kwargs = {
|
|
"hostname": hostname,
|
|
"port": port,
|
|
"username": server.get("user", "root"),
|
|
"timeout": timeout,
|
|
"banner_timeout": timeout,
|
|
}
|
|
|
|
if bind_ip:
|
|
kwargs["sock"] = _create_bound_socket(bind_ip, hostname, port, timeout)
|
|
|
|
# Try key first
|
|
if key_path and os.path.exists(key_path):
|
|
try:
|
|
kwargs["key_filename"] = key_path
|
|
client.connect(**kwargs)
|
|
transport = client.get_transport()
|
|
if transport is not None:
|
|
transport.set_keepalive(30)
|
|
return client
|
|
except paramiko.AuthenticationException:
|
|
log.debug(f"Key auth failed for {server.get('alias', '?')}, trying password")
|
|
del kwargs["key_filename"]
|
|
client = paramiko.SSHClient()
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
if bind_ip:
|
|
kwargs["sock"] = _create_bound_socket(bind_ip, hostname, port, timeout)
|
|
except Exception as e:
|
|
log.debug(f"Key connect failed: {e}")
|
|
del kwargs["key_filename"]
|
|
client = paramiko.SSHClient()
|
|
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
|
|
if bind_ip:
|
|
kwargs["sock"] = _create_bound_socket(bind_ip, hostname, port, timeout)
|
|
|
|
# Fallback to password
|
|
password = server.get("password", "")
|
|
if password:
|
|
kwargs["password"] = password
|
|
kwargs["look_for_keys"] = False
|
|
kwargs["allow_agent"] = False
|
|
client.connect(**kwargs)
|
|
transport = client.get_transport()
|
|
if transport is not None:
|
|
transport.set_keepalive(30)
|
|
return client
|
|
|
|
raise Exception(f"No auth method for {server.get('alias', 'unknown')}")
|
|
|
|
|
|
class ShellSession:
|
|
"""Persistent interactive shell session over SSH."""
|
|
|
|
def __init__(self, server: dict, key_path: str, cols: int = 80, rows: int = 24):
|
|
self.server = server
|
|
self.key_path = key_path
|
|
self.cols = cols
|
|
self.rows = rows
|
|
self._client: paramiko.SSHClient | None = None
|
|
self._channel: paramiko.Channel | None = None
|
|
self._running = False
|
|
self._read_thread: threading.Thread | None = None
|
|
|
|
# Callbacks — set by the owner
|
|
self.on_data = None # on_data(data: bytes)
|
|
self.on_disconnect = None # on_disconnect()
|
|
|
|
@property
|
|
def connected(self) -> bool:
|
|
try:
|
|
return (
|
|
self._channel is not None
|
|
and self._channel.get_transport() is not None
|
|
and self._channel.get_transport().is_active()
|
|
)
|
|
except Exception:
|
|
return False
|
|
|
|
def connect(self):
|
|
self._client = _connect_client(self.server, self.key_path)
|
|
self._channel = self._client.invoke_shell(
|
|
term="xterm-256color",
|
|
width=self.cols,
|
|
height=self.rows,
|
|
)
|
|
self._channel.settimeout(0.1)
|
|
self._running = True
|
|
self._read_thread = threading.Thread(target=self._read_loop, daemon=True)
|
|
self._read_thread.start()
|
|
|
|
def _read_loop(self):
|
|
try:
|
|
while self._running:
|
|
try:
|
|
data = self._channel.recv(65536)
|
|
if not data:
|
|
break
|
|
if self.on_data:
|
|
self.on_data(data)
|
|
except TimeoutError:
|
|
continue
|
|
except OSError:
|
|
break
|
|
except Exception as e:
|
|
log.debug(f"ShellSession read loop error: {e}")
|
|
finally:
|
|
if self._running:
|
|
self._running = False
|
|
if self.on_disconnect:
|
|
self.on_disconnect()
|
|
|
|
def send(self, data: bytes):
|
|
if self._channel and self._running:
|
|
try:
|
|
self._channel.sendall(data)
|
|
except OSError:
|
|
self._running = False
|
|
if self.on_disconnect:
|
|
self.on_disconnect()
|
|
|
|
def resize(self, cols: int, rows: int):
|
|
self.cols = cols
|
|
self.rows = rows
|
|
if self._channel and self._running:
|
|
try:
|
|
self._channel.resize_pty(width=cols, height=rows)
|
|
except OSError:
|
|
pass
|
|
|
|
def disconnect(self):
|
|
self._running = False
|
|
if self._channel:
|
|
try:
|
|
self._channel.close()
|
|
except Exception as e:
|
|
log.debug(f"ShellSession channel close: {e}")
|
|
self._channel = None
|
|
if self._client:
|
|
try:
|
|
self._client.close()
|
|
except Exception as e:
|
|
log.debug(f"ShellSession client close: {e}")
|
|
self._client = None
|
|
|
|
def reconnect(self):
|
|
self.disconnect()
|
|
time.sleep(0.2)
|
|
self.connect()
|
|
|
|
|
|
class SSHClientWrapper:
|
|
def __init__(self, server: dict, key_path: str = ""):
|
|
self.server = server
|
|
self.key_path = key_path or os.path.expanduser("~/.ssh/id_ed25519")
|
|
self._client: paramiko.SSHClient | None = None
|
|
|
|
def connect(self) -> paramiko.SSHClient:
|
|
client = _connect_client(self.server, self.key_path)
|
|
self._client = client
|
|
return client
|
|
|
|
def disconnect(self):
|
|
if self._client:
|
|
try:
|
|
self._client.close()
|
|
except Exception:
|
|
pass
|
|
self._client = None
|
|
|
|
def exec_command(self, command: str, use_sudo: bool = True) -> tuple[str, str, int]:
|
|
"""Execute command. Auto-sudo if user != root and use_sudo=True."""
|
|
client = self.connect()
|
|
stdin = stdout = stderr = None
|
|
try:
|
|
user = self.server.get("user", "root")
|
|
need_sudo = use_sudo and user != "root"
|
|
|
|
if need_sudo:
|
|
full_cmd = f"export TERM=xterm; sudo -S -p '' bash -c {_shell_quote(command)}"
|
|
else:
|
|
full_cmd = f"export TERM=xterm; {command}"
|
|
|
|
stdin, stdout, stderr = client.exec_command(full_cmd, timeout=120, get_pty=True)
|
|
|
|
if need_sudo:
|
|
password = self.server.get("password", "")
|
|
stdin.write(password + "\n")
|
|
stdin.flush()
|
|
|
|
exit_code = stdout.channel.recv_exit_status()
|
|
out = stdout.read().decode("utf-8", errors="replace")
|
|
err = stderr.read().decode("utf-8", errors="replace")
|
|
|
|
# Strip sudo noise
|
|
err_lines = [l for l in err.splitlines()
|
|
if not l.startswith("[sudo]") and "password for" not in l.lower()]
|
|
err = "\n".join(err_lines).strip()
|
|
|
|
return out, err, exit_code
|
|
finally:
|
|
# Close channels explicitly
|
|
for ch in (stdin, stdout, stderr):
|
|
if ch:
|
|
try:
|
|
ch.close()
|
|
except Exception:
|
|
pass
|
|
client.close()
|
|
|
|
def upload(self, local_path: str, remote_path: str, progress_cb=None):
|
|
client = self.connect()
|
|
try:
|
|
sftp = client.open_sftp()
|
|
if progress_cb:
|
|
sftp.put(local_path, remote_path, callback=progress_cb)
|
|
else:
|
|
sftp.put(local_path, remote_path)
|
|
try:
|
|
sftp.chmod(remote_path, 0o664)
|
|
except OSError:
|
|
pass # Windows OpenSSH doesn't support chmod
|
|
sftp.close()
|
|
finally:
|
|
client.close()
|
|
|
|
def download(self, remote_path: str, local_path: str, progress_cb=None):
|
|
client = self.connect()
|
|
try:
|
|
sftp = client.open_sftp()
|
|
if progress_cb:
|
|
sftp.get(remote_path, local_path, callback=progress_cb)
|
|
else:
|
|
sftp.get(remote_path, local_path)
|
|
sftp.close()
|
|
finally:
|
|
client.close()
|
|
|
|
def check_connection(self) -> bool:
|
|
try:
|
|
client = _connect_client(self.server, self.key_path, timeout=5)
|
|
client.close()
|
|
return True
|
|
except Exception:
|
|
return False
|
|
|
|
def install_key(self) -> str:
|
|
pub_key_path = self.key_path + ".pub"
|
|
if not os.path.exists(pub_key_path):
|
|
raise FileNotFoundError(f"No public key at {pub_key_path}")
|
|
|
|
with open(pub_key_path, "r") as f:
|
|
pub_key = f.read().strip()
|
|
|
|
out, _, _ = self.exec_command(
|
|
f'grep -c "{pub_key}" ~/.ssh/authorized_keys 2>/dev/null || echo 0',
|
|
use_sudo=False
|
|
)
|
|
if out.strip() != "0":
|
|
return "Key already installed"
|
|
|
|
command = (
|
|
f'mkdir -p ~/.ssh && chmod 700 ~/.ssh && '
|
|
f'echo "{pub_key}" >> ~/.ssh/authorized_keys && '
|
|
f'chmod 600 ~/.ssh/authorized_keys && '
|
|
f'echo "KEY_OK"'
|
|
)
|
|
out, err, code = self.exec_command(command, use_sudo=False)
|
|
if "KEY_OK" in out:
|
|
return "Key installed successfully"
|
|
raise Exception(f"Key install failed: {err or out}")
|
|
|
|
def generate_key(self) -> str:
|
|
if os.path.exists(self.key_path):
|
|
return f"Key already exists: {self.key_path}"
|
|
|
|
os.makedirs(os.path.dirname(self.key_path), exist_ok=True)
|
|
key = paramiko.Ed25519Key.generate()
|
|
key.write_private_key_file(self.key_path)
|
|
|
|
# Set restrictive permissions on private key (Unix)
|
|
if platform.system() != "Windows":
|
|
os.chmod(self.key_path, 0o600)
|
|
|
|
pub_key = f"ssh-ed25519 {key.get_base64()} server-manager"
|
|
with open(self.key_path + ".pub", "w") as f:
|
|
f.write(pub_key + "\n")
|
|
|
|
log.info(f"SSH key generated: {self.key_path}")
|
|
return f"Key generated: {self.key_path}"
|
|
|
|
|
|
class SFTPSession:
|
|
"""Persistent SFTP session for file browser."""
|
|
|
|
def __init__(self, server: dict, key_path: str):
|
|
self.server = server
|
|
self.key_path = key_path
|
|
self._client: paramiko.SSHClient | None = None
|
|
self._sftp: paramiko.SFTPClient | None = None
|
|
self.sudo_mode: bool = False
|
|
|
|
@property
|
|
def connected(self) -> bool:
|
|
try:
|
|
return (
|
|
self._client is not None
|
|
and self._sftp is not None
|
|
and self._client.get_transport() is not None
|
|
and self._client.get_transport().is_active()
|
|
)
|
|
except Exception:
|
|
return False
|
|
|
|
def connect(self):
|
|
self._client = _connect_client(self.server, self.key_path)
|
|
self._sftp = self._client.open_sftp()
|
|
|
|
def reconnect(self):
|
|
"""Disconnect and re-establish SFTP session."""
|
|
self.disconnect()
|
|
time.sleep(0.2)
|
|
self.connect()
|
|
|
|
def disconnect(self):
|
|
if self._sftp:
|
|
try:
|
|
self._sftp.close()
|
|
except Exception as e:
|
|
log.debug(f"SFTPSession sftp close: {e}")
|
|
self._sftp = None
|
|
if self._client:
|
|
try:
|
|
self._client.close()
|
|
except Exception as e:
|
|
log.debug(f"SFTPSession client close: {e}")
|
|
self._client = None
|
|
|
|
def listdir_attr(self, path: str) -> list:
|
|
return self._sftp.listdir_attr(path)
|
|
|
|
def stat(self, path: str):
|
|
return self._sftp.stat(path)
|
|
|
|
def mkdir(self, path: str):
|
|
self._sftp.mkdir(path)
|
|
|
|
def rmdir(self, path: str):
|
|
self._sftp.rmdir(path)
|
|
|
|
def remove(self, path: str):
|
|
self._sftp.remove(path)
|
|
|
|
def rename(self, old: str, new: str):
|
|
self._sftp.rename(old, new)
|
|
|
|
def upload(self, local_path: str, remote_path: str, progress_cb=None):
|
|
if progress_cb:
|
|
self._sftp.put(local_path, remote_path, callback=progress_cb)
|
|
else:
|
|
self._sftp.put(local_path, remote_path)
|
|
|
|
def download(self, remote_path: str, local_path: str, progress_cb=None):
|
|
if progress_cb:
|
|
self._sftp.get(remote_path, local_path, callback=progress_cb)
|
|
else:
|
|
self._sftp.get(remote_path, local_path)
|
|
|
|
def normalize(self, path: str) -> str:
|
|
return self._sftp.normalize(path)
|
|
|
|
# ── Recursive operations ──
|
|
|
|
def upload_dir(self, local_dir: str, remote_dir: str, progress_cb=None, file_cb=None):
|
|
"""Recursively upload a local directory to remote."""
|
|
all_files = []
|
|
for root, dirs, files in os.walk(local_dir):
|
|
rel = os.path.relpath(root, local_dir)
|
|
remote_sub = remote_dir if rel == "." else remote_dir + "/" + rel.replace("\\", "/")
|
|
try:
|
|
self._sftp.mkdir(remote_sub)
|
|
except IOError:
|
|
pass
|
|
for f in files:
|
|
local_file = os.path.join(root, f)
|
|
remote_file = remote_sub + "/" + f
|
|
all_files.append((local_file, remote_file))
|
|
|
|
for idx, (local_file, remote_file) in enumerate(all_files):
|
|
if file_cb:
|
|
file_cb(idx + 1, len(all_files), os.path.basename(local_file))
|
|
self._sftp.put(local_file, remote_file, callback=progress_cb)
|
|
|
|
def download_dir(self, remote_dir: str, local_dir: str, progress_cb=None, file_cb=None):
|
|
"""Recursively download a remote directory to local."""
|
|
all_files = []
|
|
self._walk_remote(remote_dir, remote_dir, all_files)
|
|
for idx, (remote_file, rel_path) in enumerate(all_files):
|
|
local_file = os.path.join(local_dir, rel_path)
|
|
os.makedirs(os.path.dirname(local_file), exist_ok=True)
|
|
if file_cb:
|
|
file_cb(idx + 1, len(all_files), os.path.basename(remote_file))
|
|
self._sftp.get(remote_file, local_file, callback=progress_cb)
|
|
|
|
def _walk_remote(self, base: str, current: str, result: list):
|
|
"""Recursively walk remote directory, collecting (abs_path, relative_path) tuples."""
|
|
import stat as stat_mod
|
|
for attr in self._sftp.listdir_attr(current):
|
|
full = current + "/" + attr.filename
|
|
rel = full[len(base):].lstrip("/")
|
|
if stat_mod.S_ISDIR(attr.st_mode or 0):
|
|
self._walk_remote(base, full, result)
|
|
else:
|
|
result.append((full, rel.replace("/", os.sep)))
|
|
|
|
def rmdir_recursive(self, path: str):
|
|
"""Recursively delete a remote directory."""
|
|
import stat as stat_mod
|
|
for attr in self._sftp.listdir_attr(path):
|
|
child = path + "/" + attr.filename
|
|
if stat_mod.S_ISDIR(attr.st_mode or 0):
|
|
self.rmdir_recursive(child)
|
|
else:
|
|
self._sftp.remove(child)
|
|
self._sftp.rmdir(path)
|
|
|
|
# ── Sudo operations ──
|
|
|
|
def exec_command(self, cmd: str) -> str:
|
|
"""Execute command via SSH with optional sudo wrapper."""
|
|
if not self._client:
|
|
raise Exception("Not connected")
|
|
password = self.server.get("password", "")
|
|
user = self.server.get("user", "root")
|
|
if self.sudo_mode and user != "root":
|
|
full_cmd = f"sudo -S -p '' bash -c {_shell_quote(cmd)}"
|
|
else:
|
|
full_cmd = cmd
|
|
stdin, stdout, stderr = self._client.exec_command(full_cmd, timeout=30)
|
|
if self.sudo_mode and user != "root" and password:
|
|
# Wait briefly for sudo prompt to appear before sending password
|
|
time.sleep(0.1)
|
|
stdin.write(password + "\n")
|
|
stdin.flush()
|
|
return stdout.read().decode("utf-8", errors="replace")
|
|
|
|
def listdir_attr_sudo(self, path: str) -> list:
|
|
"""List directory using sudo ls -la, returning objects with .filename/.st_size/.st_mtime/.st_mode."""
|
|
output = self.exec_command(f"ls -la --time-style=+%s {_shell_quote(path)}")
|
|
results = []
|
|
for line in output.strip().splitlines():
|
|
if line.startswith("total "):
|
|
continue
|
|
# With --time-style=+%s the columns are:
|
|
# perms links owner group size mtime name [-> target if symlink]
|
|
# Use maxsplit=8 to preserve spaces in filenames, giving us:
|
|
# [0=perms, 1=links, 2=owner, 3=group, 4=size, 5=mtime, 6=name...]
|
|
parts = line.split(None, 8)
|
|
if len(parts) < 7:
|
|
continue
|
|
perms = parts[0]
|
|
if not (perms.startswith("d") or perms.startswith("l") or perms.startswith("-")):
|
|
continue
|
|
try:
|
|
size = int(parts[4])
|
|
except ValueError:
|
|
size = 0
|
|
continue
|
|
try:
|
|
mtime = int(parts[5])
|
|
except ValueError:
|
|
mtime = 0
|
|
continue
|
|
name = parts[6]
|
|
# Handle cases where name contains " -> " (symlinks) or has spaces
|
|
if len(parts) > 7:
|
|
# This means the filename itself contained spaces and was split
|
|
name = parts[6] + " " + parts[7]
|
|
# Strip symlink target (e.g. "name -> target")
|
|
name = name.split(" -> ")[0].strip()
|
|
if name in (".", ".."):
|
|
continue
|
|
mode = _parse_ls_perms(perms)
|
|
entry = _SudoFileAttr(name, size, mtime, mode)
|
|
results.append(entry)
|
|
return results
|
|
|
|
def upload_sudo(self, local_path: str, remote_path: str, progress_cb=None):
|
|
"""Upload via SFTP to /tmp then sudo mv to destination."""
|
|
import random
|
|
tmp_name = f"/tmp/.sm_upload_{random.randint(100000, 999999)}"
|
|
if progress_cb:
|
|
self._sftp.put(local_path, tmp_name, callback=progress_cb)
|
|
else:
|
|
self._sftp.put(local_path, tmp_name)
|
|
self.exec_command(f"mv {_shell_quote(tmp_name)} {_shell_quote(remote_path)}")
|
|
|
|
def download_sudo(self, remote_path: str, local_path: str, progress_cb=None):
|
|
"""Copy via sudo to /tmp then download via SFTP."""
|
|
import random
|
|
tmp_name = f"/tmp/.sm_download_{random.randint(100000, 999999)}"
|
|
self.exec_command(f"cp {_shell_quote(remote_path)} {_shell_quote(tmp_name)} && chmod 644 {_shell_quote(tmp_name)}")
|
|
try:
|
|
if progress_cb:
|
|
self._sftp.get(tmp_name, local_path, callback=progress_cb)
|
|
else:
|
|
self._sftp.get(tmp_name, local_path)
|
|
finally:
|
|
self.exec_command(f"rm -f {_shell_quote(tmp_name)}")
|
|
|
|
|
|
class _SudoFileAttr:
|
|
"""Mimics paramiko SFTPAttributes for sudo ls output."""
|
|
def __init__(self, filename: str, st_size: int, st_mtime: int, st_mode: int):
|
|
self.filename = filename
|
|
self.st_size = st_size
|
|
self.st_mtime = st_mtime
|
|
self.st_mode = st_mode
|
|
|
|
|
|
def _parse_ls_perms(perms: str) -> int:
|
|
"""Parse ls -l permission string like 'drwxr-xr-x' into stat mode int."""
|
|
import stat as stat_mod
|
|
mode = 0
|
|
if perms[0] == "d":
|
|
mode |= stat_mod.S_IFDIR
|
|
elif perms[0] == "l":
|
|
mode |= stat_mod.S_IFLNK
|
|
else:
|
|
mode |= stat_mod.S_IFREG
|
|
mapping = "rwxrwxrwx"
|
|
bits = [stat_mod.S_IRUSR, stat_mod.S_IWUSR, stat_mod.S_IXUSR,
|
|
stat_mod.S_IRGRP, stat_mod.S_IWGRP, stat_mod.S_IXGRP,
|
|
stat_mod.S_IROTH, stat_mod.S_IWOTH, stat_mod.S_IXOTH]
|
|
for i, (ch, bit) in enumerate(zip(perms[1:10], bits)):
|
|
if ch != "-":
|
|
mode |= bit
|
|
return mode
|
|
|
|
|
|
def _shell_quote(s: str) -> str:
|
|
return "'" + s.replace("'", "'\\''") + "'"
|