mtr-cli 0.1.0__py3-none-any.whl

mtr/ssh.py

@@ -0,0 +1,199 @@
+import os
+import shutil
+import subprocess
+from typing import Generator, Optional
+
+import paramiko
+
+from mtr.logger import get_logger
+
+
+class SSHError(Exception):
+    pass
+
+
+# Constants for batch mode
+BUFFER_SIZE = 32768  # 32KB for better performance
+
+
+def _check_ssh_availability():
+    """Check if ssh command is available."""
+    if shutil.which("ssh") is None:
+        raise SSHError(
+            "SSH command not found. Please install OpenSSH client.\n"
+            "  macOS: brew install openssh\n"
+            "  Ubuntu/Debian: sudo apt-get install openssh-client\n"
+            "  CentOS/RHEL: sudo yum install openssh-clients"
+        )
+
+
+def _check_sshpass_availability():
+    """Check if sshpass command is available."""
+    if shutil.which("sshpass") is None:
+        raise SSHError(
+            "sshpass command not found. Please install sshpass for password authentication.\n"
+            "  macOS: brew install hudochenkov/sshpass/sshpass\n"
+            "  Ubuntu/Debian: sudo apt-get install sshpass\n"
+            "  CentOS/RHEL: sudo yum install sshpass"
+        )
+
+
+class SSHClientWrapper:
+    def __init__(
+        self,
+        host: str,
+        user: str,
+        port: int = 22,
+        key_filename: Optional[str] = None,
+        password: Optional[str] = None,
+    ):
+        self.host = host
+        self.user = user
+        self.port = port
+        self.key_filename = key_filename
+        self.password = password
+        self.client: Optional[paramiko.SSHClient] = None
+
+    def connect(self):
+        logger = get_logger()
+        logger.info(f"Connecting to {self.host}:{self.port} as {self.user}", module="mtr.ssh")
+        self.client = paramiko.SSHClient()
+        self.client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+        connect_kwargs = {
+            "hostname": self.host,
+            "username": self.user,
+            "port": self.port,
+            "timeout": 10,
+        }
+
+        if self.key_filename:
+            expanded_key = os.path.expanduser(self.key_filename)
+            connect_kwargs["key_filename"] = expanded_key
+            logger.info("Using key-based authentication", module="mtr.ssh")
+        elif self.password:
+            connect_kwargs["password"] = self.password
+            logger.info("Using password authentication", module="mtr.ssh")
+        else:
+            logger.info("No authentication method specified, using default", module="mtr.ssh")
+
+        try:
+            self.client.connect(**connect_kwargs)
+            logger.info(f"SSH connection established to {self.host}", module="mtr.ssh")
+        except paramiko.SSHException as e:
+            logger.error(f"Failed to connect to {self.host}: {e}", module="mtr.ssh")
+            raise SSHError(f"Failed to connect to {self.host}: {e}")
+
+    def _build_command(self, command: str, workdir: Optional[str] = None, pre_cmd: Optional[str] = None) -> str:
+        parts = []
+        if workdir:
+            parts.append(f"cd {workdir}")
+        if pre_cmd:
+            parts.append(pre_cmd)
+        parts.append(command)
+        return " && ".join(parts)
+
+    def exec_command_stream(
+        self,
+        command: str,
+        workdir: Optional[str] = None,
+        pre_cmd: Optional[str] = None,
+        pty: bool = True,
+    ) -> Generator[str, None, int]:
+        """
+        Executes command and yields output lines.
+        Returns the exit code.
+        Suitable for batch mode or when interactivity is not required.
+        """
+        logger = get_logger()
+        logger.info(f"Executing command (stream mode): {command}", module="mtr.ssh")
+        logger.debug(f"Workdir: {workdir}, Pre-cmd: {pre_cmd}, PTY: {pty}", module="mtr.ssh")
+
+        if not self.client:
+            raise SSHError("Client not connected")
+
+        full_command = self._build_command(command, workdir, pre_cmd)
+        logger.debug(f"Full command: {full_command}", module="mtr.ssh")
+
+        try:
+            stdin, stdout, stderr = self.client.exec_command(full_command, get_pty=pty)
+            stdin.close()
+
+            logger.debug("Command executed, starting output stream", module="mtr.ssh")
+
+            line_count = 0
+            for line in stdout:
+                line_count += 1
+                if line_count % 100 == 0:
+                    logger.debug(f"Streamed {line_count} lines so far", module="mtr.ssh")
+                yield line
+
+            logger.debug(f"Output stream ended, total lines: {line_count}", module="mtr.ssh")
+
+            exit_code = stdout.channel.recv_exit_status()
+            logger.info(f"Command exited with code: {exit_code}", module="mtr.ssh")
+            return exit_code
+
+        except paramiko.SSHException as e:
+            logger.error(f"Command execution failed: {e}", module="mtr.ssh")
+            raise SSHError(f"Command execution failed: {e}")
+
+    def run_interactive_shell(self, command: str, workdir: Optional[str] = None, pre_cmd: Optional[str] = None) -> int:
+        """
+        Runs an interactive shell using system ssh -t command.
+        This provides full TTY support with proper signal handling.
+        Returns exit code.
+        """
+        logger = get_logger()
+        logger.info(f"Starting interactive shell via ssh -t: {command}", module="mtr.ssh")
+        logger.debug(f"Workdir: {workdir}, Pre-cmd: {pre_cmd}", module="mtr.ssh")
+
+        # Check SSH availability
+        _check_ssh_availability()
+
+        # Check sshpass availability if password is used
+        if self.password:
+            _check_sshpass_availability()
+
+        full_command = self._build_command(command, workdir, pre_cmd)
+
+        # Build ssh command
+        ssh_cmd = ["ssh", "-t"]
+
+        # Port
+        if self.port != 22:
+            ssh_cmd.extend(["-p", str(self.port)])
+
+        # Key authentication
+        if self.key_filename:
+            ssh_cmd.extend(["-i", os.path.expanduser(self.key_filename)])
+
+        # Target host and command
+        target = f"{self.user}@{self.host}"
+        ssh_cmd.extend([target, full_command])
+
+        # Wrap with sshpass if password is provided
+        if self.password:
+            ssh_cmd = ["sshpass", "-p", self.password] + ssh_cmd
+
+        logger.debug(f"Executing: {' '.join(ssh_cmd)}", module="mtr.ssh")
+
+        # Run command with direct stdin/stdout/stderr forwarding
+        try:
+            result = subprocess.run(ssh_cmd)
+            logger.info(f"Interactive shell exited with code: {result.returncode}", module="mtr.ssh")
+            return result.returncode
+        except FileNotFoundError as e:
+            # This shouldn't happen if _check_ssh_availability passed, but just in case
+            logger.error(f"Command not found: {e}", module="mtr.ssh")
+            raise SSHError(f"SSH command execution failed: {e}")
+        except Exception as e:
+            logger.error(f"Interactive shell failed: {e}", module="mtr.ssh")
+            raise SSHError(f"Interactive shell failed: {e}")
+
+    def close(self):
+        if self.client:
+            logger = get_logger()
+            logger.debug("Closing SSH connection", module="mtr.ssh")
+            self.client.close()
+            logger.debug("SSH connection closed", module="mtr.ssh")

mtr/sync.py

@@ -0,0 +1,338 @@
+import fnmatch
+import os
+import shlex
+import shutil
+import subprocess
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+import paramiko
+
+
+class SyncError(Exception):
+    pass
+
+
+class BaseSyncer(ABC):
+    def __init__(self, local_dir: str, remote_dir: str, exclude: List[str]):
+        self.local_dir = local_dir
+        self.remote_dir = remote_dir
+        self.exclude = exclude
+
+    @abstractmethod
+    def sync(self):
+        pass
+
+
+class SftpSyncer(BaseSyncer):
+    def __init__(
+        self,
+        local_dir: str,
+        remote_dir: str,
+        host: str,
+        user: str,
+        key_filename: Optional[str] = None,
+        password: Optional[str] = None,
+        port: int = 22,
+        exclude: List[str] = None,
+    ):
+        super().__init__(local_dir, remote_dir, exclude or [])
+        self.host = host
+        self.user = user
+        self.key_filename = key_filename
+        self.password = password
+        self.port = port
+        self.transport = None
+        self.sftp = None
+
+    def _should_ignore(self, filename: str) -> bool:
+        for pattern in self.exclude:
+            # Handle directory exclusion (basic)
+            if pattern.endswith("/") and filename.startswith(pattern.rstrip("/")):
+                return True
+            if fnmatch.fnmatch(filename, pattern):
+                return True
+        return False
+
+    def _connect(self):
+        try:
+            self.transport = paramiko.Transport((self.host, self.port))
+            connect_kwargs = {"username": self.user}
+
+            if self.key_filename:
+                key_path = os.path.expanduser(self.key_filename)
+                # Try different key types? For now assuming RSA or standard loading
+                # Or just use connect method of SSHClient? No, Sftp is lower level usually,
+                # but we can use SSHClient to get sftp
+
+                # Simpler approach: Use SSHClientWrapper logic or just Paramiko SSHClient
+                client = paramiko.SSHClient()
+                client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+
+                k_kwargs = {
+                    "hostname": self.host,
+                    "username": self.user,
+                    "port": self.port,
+                    "key_filename": key_path,
+                }
+                if self.password:
+                    k_kwargs["password"] = self.password
+
+                client.connect(**k_kwargs)
+                self.sftp = client.open_sftp()
+                return
+
+            if self.password:
+                connect_kwargs["password"] = self.password
+                self.transport.connect(**connect_kwargs)
+                self.sftp = paramiko.SFTPClient.from_transport(self.transport)
+            else:
+                raise SyncError("No auth method provided (need key or password)")
+
+        except Exception as e:
+            raise SyncError(f"SFTP Connection failed: {e}")
+
+    def _ensure_remote_dir(self, remote_path: str):
+        """Recursively create remote directory."""
+        # This is a bit expensive, optimization: assume parent exists or try/except
+        # Simple implementation:
+        dirs = remote_path.split("/")
+        current = ""
+        for d in dirs:
+            if not d:
+                continue
+            current += f"/{d}"
+            try:
+                self.sftp.stat(current)
+            except FileNotFoundError:
+                try:
+                    self.sftp.mkdir(current)
+                except OSError:
+                    pass  # Already exists maybe
+
+    def sync(self):
+        if not self.sftp:
+            self._connect()
+
+        # Ensure base remote dir exists
+        try:
+            self.sftp.stat(self.remote_dir)
+        except FileNotFoundError:
+            self._ensure_remote_dir(self.remote_dir)
+
+        # Walk local tree
+        for root, dirs, files in os.walk(self.local_dir):
+            # Filtering dirs in place to prevent recursion
+            dirs[:] = [d for d in dirs if not self._should_ignore(d)]
+
+            rel_path = os.path.relpath(root, self.local_dir)
+            if rel_path == ".":
+                remote_root = self.remote_dir
+            else:
+                remote_root = os.path.join(self.remote_dir, rel_path)
+
+                # Check/Create remote dir
+                try:
+                    self.sftp.stat(remote_root)
+                except FileNotFoundError:
+                    self.sftp.mkdir(remote_root)
+
+            for file in files:
+                if self._should_ignore(file):
+                    continue
+
+                local_file = os.path.join(root, file)
+                remote_file = os.path.join(remote_root, file)
+
+                # Check sync necessity (Size & Mtime)
+                should_upload = True
+                try:
+                    remote_stat = self.sftp.stat(remote_file)
+                    local_stat = os.stat(local_file)
+
+                    if remote_stat.st_size == local_stat.st_size and int(remote_stat.st_mtime) >= int(local_stat.st_mtime):
+                        should_upload = False
+                except FileNotFoundError:
+                    pass  # Does not exist, must upload
+
+                if should_upload:
+                    # print(f"Uploading {local_file} -> {remote_file}")
+                    self.sftp.put(local_file, remote_file)
+                    # Preserve permissions
+                    mode = os.stat(local_file).st_mode
+                    self.sftp.chmod(remote_file, mode)
+
+        if self.sftp:
+            self.sftp.close()
+        if self.transport:
+            self.transport.close()
+
+    def download(self, remote_path: str, local_path: str):
+        """Download file or directory from remote to local."""
+        if not self.sftp:
+            self._connect()
+
+        try:
+            # Ensure local directory exists
+            local_dir = os.path.dirname(local_path)
+            if local_dir and not os.path.exists(local_dir):
+                os.makedirs(local_dir, exist_ok=True)
+
+            try:
+                import stat
+
+                remote_stat = self.sftp.stat(remote_path)
+                is_dir = stat.S_ISDIR(remote_stat.st_mode)
+
+                if is_dir:
+                    self._download_dir(remote_path, local_path)
+                else:
+                    self._download_file(remote_path, local_path)
+            except FileNotFoundError:
+                raise SyncError(f"Remote path not found: {remote_path}")
+            except Exception as e:
+                raise SyncError(f"Download failed: {e}")
+        finally:
+            if self.sftp:
+                self.sftp.close()
+            if self.transport:
+                self.transport.close()
+
+    def _should_download_file(self, remote_file: str, local_file: str) -> bool:
+        """Check if file should be downloaded based on size and mtime."""
+        try:
+            remote_stat = self.sftp.stat(remote_file)
+            local_stat = os.stat(local_file)
+
+            # Size different, need download
+            if remote_stat.st_size != local_stat.st_size:
+                return True
+
+            # Remote file is newer than local file
+            return int(remote_stat.st_mtime) > int(local_stat.st_mtime)
+        except FileNotFoundError:
+            return True  # Local file doesn't exist, must download
+
+    def _download_file(self, remote_file: str, local_file: str):
+        """Download a single file with incremental check."""
+        if not self._should_download_file(remote_file, local_file):
+            return  # No need to download
+
+        self.sftp.get(remote_file, local_file)
+        # Preserve permissions
+        remote_stat = self.sftp.stat(remote_file)
+        os.chmod(local_file, remote_stat.st_mode)
+
+    def _download_dir(self, remote_dir: str, local_dir: str):
+        """Recursively download a directory."""
+        if not os.path.exists(local_dir):
+            os.makedirs(local_dir, exist_ok=True)
+
+        for entry in self.sftp.listdir_attr(remote_dir):
+            remote_path = f"{remote_dir}/{entry.filename}"
+            local_path = os.path.join(local_dir, entry.filename)
+
+            if self._should_ignore(entry.filename):
+                continue
+
+            import stat
+
+            if stat.S_ISDIR(entry.st_mode):
+                self._download_dir(remote_path, local_path)
+            else:
+                self._download_file(remote_path, local_path)
+
+
+class RsyncSyncer(BaseSyncer):
+    def __init__(
+        self,
+        local_dir: str,
+        remote_dir: str,
+        host: str,
+        user: str,
+        key_filename: Optional[str] = None,
+        password: Optional[str] = None,
+        port: int = 22,
+        exclude: List[str] = None,
+    ):
+        super().__init__(local_dir, remote_dir, exclude or [])
+        self.host = host
+        self.user = user
+        self.key_filename = key_filename
+        self.password = password
+        self.port = port
+
+    def _build_ssh_options(self) -> str:
+        """Build SSH options string for rsync."""
+        opts = f"ssh -p {self.port}"
+        if self.key_filename:
+            opts += f" -i {self.key_filename}"
+        return opts
+
+    def _build_rsync_base(self) -> List[str]:
+        """Build rsync base command with common options."""
+        cmd = ["rsync", "-azq"]
+
+        # Add excludes
+        for item in self.exclude:
+            cmd.append(f"--exclude={item}")
+
+        # SSH options
+        cmd.extend(["-e", self._build_ssh_options()])
+
+        return cmd
+
+    def _wrap_with_sshpass(self, cmd: List[str]) -> List[str]:
+        """Wrap command with sshpass if password authentication is used."""
+        if self.password and not self.key_filename:
+            return ["sshpass", "-p", self.password] + cmd
+        return cmd
+
+    def _check_sshpass(self):
+        """Check if sshpass is available when password authentication is used."""
+        if self.password and not self.key_filename:
+            if not shutil.which("sshpass"):
+                raise SyncError("Rsync with password requires 'sshpass'. Please install it or use SSH Key.")
+
+    def _build_rsync_command(self) -> List[str]:
+        """Build rsync command for uploading (local -> remote)."""
+        # Ensure local dir ends with / to sync contents, not the dir itself
+        src = self.local_dir if self.local_dir.endswith("/") else f"{self.local_dir}/"
+        dest = f"{self.user}@{self.host}:{shlex.quote(self.remote_dir)}"
+
+        cmd = self._build_rsync_base()
+        cmd.extend([src, dest])
+
+        return self._wrap_with_sshpass(cmd)
+
+    def _build_rsync_download_command(self, remote_path: str, local_path: str) -> List[str]:
+        """Build rsync command for downloading (remote -> local)."""
+        src = f"{self.user}@{self.host}:{shlex.quote(remote_path)}"
+
+        cmd = self._build_rsync_base()
+        cmd.extend([src, local_path])
+
+        return self._wrap_with_sshpass(cmd)
+
+    def sync(self):
+        self._check_sshpass()
+        cmd = self._build_rsync_command()
+        try:
+            subprocess.run(cmd, check=True)
+        except subprocess.CalledProcessError as e:
+            raise SyncError(f"Rsync failed with exit code {e.returncode}")
+
+    def download(self, remote_path: str, local_path: str):
+        """Download file or directory from remote to local."""
+        self._check_sshpass()
+
+        # Ensure local parent directory exists
+        local_dir = os.path.dirname(local_path)
+        if local_dir and not os.path.exists(local_dir):
+            os.makedirs(local_dir, exist_ok=True)
+
+        cmd = self._build_rsync_download_command(remote_path, local_path)
+        try:
+            subprocess.run(cmd, check=True)
+        except subprocess.CalledProcessError as e:
+            raise SyncError(f"Rsync download failed with exit code {e.returncode}")