server-manager/core/ssh_client.py

"""
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)
            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)
        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:
        return (
            self._channel is not None
            and self._channel.get_transport() is not None
            and self._channel.get_transport().is_active()
        )

    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(4096)
                    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:
                pass

    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:
                pass
            self._channel = None
        if self._client:
            try:
                self._client.close()
            except Exception:
                pass
            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)
            sftp.chmod(remote_path, 0o664)
            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}"


def _shell_quote(s: str) -> str:
    return "'" + s.replace("'", "'\\''") + "'"