cloudmesh-ai-ssh 0.1.0__py3-none-any.whl

cloudmesh/ai/command/ssh.py

@@ -0,0 +1,307 @@
+import click
+import time
+import sys
+import requests
+from pathlib import Path
+from cloudmesh.ai.common.io import console, path_expand, Table
+from cloudmesh.ai.common.logging_utils import get_contextual_logger
+from cloudmesh.ai.common.telemetry import Telemetry
+from cloudmesh.ai.ssh.tunnel import Tunnel
+from cloudmesh.ai.ssh.exceptions import SSHTunnelError
+from cloudmesh.ai.ssh.ssh_config import SSHConfig
+
+# Initialize Logger and Telemetry
+logger = get_contextual_logger("ssh")
+telemetry = Telemetry("ssh")
+
+def _render_table(rows):
+    """Helper to render a list of dictionaries as a rich table."""
+    if not rows:
+        return None
+    
+    table = Table()
+    
+    # Use keys of the first dictionary as headers
+    headers = list(rows[0].keys())
+    for header in headers:
+        table.add_column(header)
+    
+    # Add rows
+    for row in rows:
+        table.add_row(*[str(row.get(h, "")) for h in headers])
+    
+    return table
+
+# Define the group for the command
+@click.group(name="ssh")
+def ssh_group():
+    """ssh command group."""
+    pass
+
+@click.group(name="list", invoke_without_command=True)
+@click.pass_context
+def list_group(ctx):
+    """List SSH configurations and active tunnels."""
+    if ctx.invoked_subcommand is None:
+        ctx.invoke(hosts_cmd)
+
+@list_group.command(name="hosts")
+def hosts_cmd():
+    """List all hosts defined in the SSH config file."""
+    cfg = SSHConfig()
+    hosts = cfg.list()
+    
+    if not hosts:
+        console.info("No hosts found in SSH config.")
+        return
+
+    rows = []
+    for host in hosts:
+        rows.append({
+            "Host": host,
+            "Hostname": cfg.hostname(host),
+            "User": cfg.username(host),
+            "Details": cfg.get_options(host)
+        })
+    
+    table = _render_table(rows)
+    if table:
+        console.print(table)
+
+@list_group.command(name="tunnel")
+def tunnel_list_cmd():
+    """List all tunnels defined in config and their active status."""
+    cfg = SSHConfig()
+    tunnels = cfg.get_tunnels()
+    
+    if not tunnels:
+        console.info("No tunnels defined in SSH config.")
+        return
+
+    from cloudmesh.ai.ssh.base import SSHBase
+    checker = SSHBase()
+    
+    rows = []
+    for t in tunnels:
+        status = "Unknown"
+        if t["type"] == "Local":
+            try:
+                is_open = checker.is_port_open("localhost", int(t["local_port"]))
+                status = "Active" if is_open else "Inactive"
+            except Exception:
+                status = "Error"
+        
+        rows.append({
+            "Host": t["host"],
+            "Type": t["type"],
+            "Local Port": t["local_port"],
+            "Remote Target": t["remote_target"],
+            "Status": status
+        })
+    
+    table = _render_table(rows)
+    if table:
+        console.print(table)
+
+ssh_group.add_command(list_group)
+
+@ssh_group.command(name="run")
+def run_cmd():
+    """Run the main functionality of ssh."""
+    logger.info("Executing ssh run command")
+    console.ok(f"The ssh extension is running successfully!")
+
+@ssh_group.command(name="tunnel")
+@click.argument("target")
+@click.option("--local-port", type=int, help="Local port to bind to. Defaults to remote port.")
+@click.option("--remote-host", default="localhost", help="Remote host relative to the SSH server. Defaults to localhost.")
+@click.option("--ssh-user", help="SSH username to use.")
+def tunnel_cmd(target, local_port, remote_host, ssh_user):
+    """Create an SSH tunnel.
+    
+    Target should be in the format HOST:PORT (e.g., my-server:8000).
+    """
+    try:
+        if ":" not in target:
+            raise click.BadParameter("Target must be in the format HOST:PORT")
+        
+        ssh_host, remote_port_str = target.split(":", 1)
+        remote_port = int(remote_port_str)
+        
+        # Default local port to remote port if not provided
+        l_port = local_port if local_port else remote_port
+        
+        console.info(f"Setting up tunnel: localhost:{l_port} -> {remote_host}:{remote_port} via {ssh_host}")
+        
+        tunnel = Tunnel(
+            local_port=l_port,
+            remote_host=remote_host,
+            remote_port=remote_port,
+            ssh_host=ssh_host,
+            ssh_user=ssh_user
+        )
+        
+        if tunnel.start():
+            console.ok(f"Tunnel is active. Press Ctrl+C to stop it.")
+            try:
+                while True:
+                    if not tunnel.is_active():
+                        console.warn("Tunnel process terminated unexpectedly.")
+                        break
+                    time.sleep(1)
+            except KeyboardInterrupt:
+                console.info("\nStopping tunnel...")
+                tunnel.stop()
+                console.ok("Tunnel closed.")
+        else:
+            console.error("Failed to start tunnel.")
+            sys.exit(1)
+            
+    except ValueError:
+        raise click.BadParameter("Port must be a valid integer.")
+    except SSHTunnelError as e:
+        console.error(f"SSH Tunnel Error: {e}")
+        sys.exit(1)
+    except Exception as e:
+        console.error(f"Unexpected error: {e}")
+        sys.exit(1)
+
+@ssh_group.command(name="test-path")
+@click.argument("path")
+def test_path_cmd(path):
+    """Example command showing path expansion."""
+    expanded = path_expand(path)
+    console.info(f"Expanded path: {expanded}")
+
+@click.group(name="check", invoke_without_command=True)
+@click.pass_context
+def check_group(ctx):
+    """Check the SSH config file for malformed entries."""
+    if ctx.invoked_subcommand is None:
+        ctx.invoke(check_basic)
+
+@check_group.command(name="basic")
+def check_basic():
+    """Run basic SSH config validation."""
+    cfg = SSHConfig()
+    errors = cfg.check()
+    
+    if not errors:
+        console.ok("SSH config is valid!")
+        return
+
+    console.print(f"[yellow]Found {len(errors)} issue(s) in {cfg.filename}:[/yellow]")
+    
+    rows = []
+    for err in errors:
+        message = err["message"]
+        if ": " in message:
+            category, details = message.split(": ", 1)
+        else:
+            category, details = message, ""
+        
+        rows.append({
+            "Line": err["line"],
+            "Category": category,
+            "Details": details
+        })
+    
+    table = _render_table(rows)
+    if table:
+        console.print(table)
+
+@check_group.command(name="ai")
+@click.option("--api-key", default=None, help="API key for the vLLM server.")
+def check_ai(api_key):
+    """Gather diagnostic data and submit it to the vLLM server for repair."""
+    cfg = SSHConfig()
+    
+    console.info("Gathering SSH configuration diagnostics...")
+    
+    # 1. Get malformed entries
+    errors = cfg.check()
+    
+    # 2. Get raw config content
+    content = cfg.get_content()
+    
+    # 3. Package the data
+    diag_data = {
+        "config_file": str(cfg.filename),
+        "raw_content": content,
+        "errors": errors
+    }
+    
+    # 4. Submit to the vLLM server (OpenAI compatible API)
+    base_url = "http://localhost:17704"
+    
+    # Load default API key if none provided
+    if not api_key:
+        try:
+            key_path = Path("~/gemma/server_master_key.txt").expanduser()
+            if key_path.exists():
+                api_key = key_path.read_text().strip()
+        except Exception as e:
+            logger.debug(f"Could not load default API key from ~/gemma/server_master_key.txt: {e}")
+    
+    headers = {}
+    if api_key:
+        headers["Authorization"] = f"Bearer {api_key}"
+    
+    try:
+        # a. Get the model name from the vLLM server
+        model_response = requests.get(f"{base_url}/v1/models", headers=headers, timeout=10)
+        model_response.raise_for_status()
+        models = model_response.json().get("data", [])
+        if not models:
+            raise Exception("No models found on the vLLM server.")
+        model_name = models[0]["id"]
+        
+        # b. Construct the prompt for the LLM
+        system_prompt = (
+            "You are an expert SSH configuration assistant. Your goal is to analyze the provided "
+            "SSH config diagnostics and provide a corrected version of the config file. "
+            "Explain the issues found and provide the final corrected content clearly."
+        )
+        user_prompt = f"Please analyze these SSH config diagnostics and provide a fix:\n\n{diag_data}"
+        
+        payload = {
+            "model": model_name,
+            "messages": [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": user_prompt}
+            ],
+            "temperature": 0.2
+        }
+        
+        console.info(f"Sending diagnostics to vLLM model {model_name} at {base_url}/v1/chat/completions...")
+        
+        response = requests.post(f"{base_url}/v1/chat/completions", json=payload, headers=headers, timeout=60)
+        response.raise_for_status()
+        
+        result = response.json()
+        answer = result["choices"][0]["message"]["content"]
+        
+        console.ok("AI Analysis complete!")
+        console.print(f"\n[bold green]AI Repair Suggestion:[/bold green]\n\n{answer}")
+            
+    except requests.exceptions.ConnectionError:
+        console.error(f"Could not connect to the vLLM server at {base_url}. Is it running?")
+        _save_fallback(diag_data)
+    except Exception as e:
+        console.error(f"An error occurred during AI analysis: {e}")
+        _save_fallback(diag_data)
+
+def _save_fallback(diag_data):
+    diag_file = Path("ssh_diag.json")
+    import json
+    with open(diag_file, "w") as f:
+        json.dump(diag_data, f, indent=4)
+    console.print(f"[yellow]Diagnostic data saved to {diag_file} as fallback.[/yellow]")
+
+ssh_group.add_command(check_group)
+
+entry_point = ssh_group
+
+def register(cli):
+    """Registers the ssh command group to the main CLI."""
+    cli.add_command(ssh_group, name="ssh")

cloudmesh/ai/ssh/__init__.py

@@ -0,0 +1,4 @@
+from .tunnel import Tunnel
+from .ssh_config import SSHConfig
+
+__all__ = ["Tunnel", "SSHConfig"]

cloudmesh/ai/ssh/authorized_keys.py

@@ -0,0 +1,128 @@
+# Copyright 2026 Gregor von Laszewski
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+import io
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+from cloudmesh.ai.common import logging as ai_log
+from .base import SSHBase
+from .exceptions import SSHError
+
+logger = ai_log.get_logger("ai.ssh.authorized_keys")
+
+
+class AuthorizedKeys(SSHBase):
+    """Class to manage authorized keys."""
+
+    def __init__(self, debug: bool = False):
+        super().__init__(debug=debug)
+        self._order: Dict[int, str] = {}
+        self._keys: Dict[str, str] = {}
+
+    def get_fingerprint_from_public_key(self, pubkey: str) -> str:
+        """Generate the fingerprint of a public key.
+
+        Args:
+            pubkey (str): the value of the public key
+
+        Returns:
+            str: fingerprint
+        """
+        try:
+            # Use ssh-keygen -l -f /dev/stdin to avoid creating temporary files
+            process = self._execute(
+                ["ssh-keygen", "-l", "-f", "/dev/stdin"],
+                input_data=pubkey
+            )
+            output = process.stdout.strip()
+            # Output format: "2048 SHA256:abc... user@host (RSA)"
+            parts = output.split(' ')
+            if len(parts) >= 2:
+                return parts[1]
+            return output
+        except Exception as e:
+            logger.error(f"Failed to get fingerprint for public key: {e}")
+            return ""
+
+    def load(self, path: Union[str, Path]) -> None:
+        """Load the keys from a path.
+
+        Args:
+            path: the filename (path) in which we find the keys
+        """
+        path = self.resolve_path(str(path))
+        if not path.exists():
+            logger.warning(f"Authorized keys file not found: {path}")
+            return
+
+        try:
+            with path.open('r') as fd:
+                for pubkey in map(str.strip, fd):
+                    # skip empty lines and comments
+                    if not pubkey or pubkey.startswith('#'):
+                        continue
+                    self.add(pubkey)
+        except Exception as e:
+            logger.error(f"Error loading authorized keys from {path}: {e}")
+
+    def add(self, pubkey: str):
+        """Add a public key.
+
+        Args:
+            pubkey: the public key string.
+        """
+        fingerprint = self.get_fingerprint_from_public_key(pubkey)
+        if not fingerprint:
+            logger.warning("Could not generate fingerprint for public key; skipping.")
+            return
+
+        if fingerprint not in self._keys:
+            self._order[len(self._keys)] = fingerprint
+            self._keys[fingerprint] = pubkey
+
+    def remove(self, fingerprint: str):
+        """Removes the public key by its fingerprint.
+
+        Args:
+            fingerprint: the fingerprint of the public key to remove.
+        """
+        if fingerprint in self._keys:
+            del self._keys[fingerprint]
+            # Rebuild order to keep it contiguous
+            new_order = {}
+            for i, f in enumerate(self._order.values()):
+                if f != fingerprint:
+                    new_order[i] = f
+            self._order = new_order
+        else:
+            logger.warning(f"Fingerprint {fingerprint} not found in authorized keys.")
+
+    def __str__(self) -> str:
+        with io.StringIO() as sio:
+            for fingerprint in self._order.values():
+                key = self._keys.get(fingerprint)
+                if key:
+                    sio.write(key)
+                    sio.write('\n')
+            return sio.getvalue().strip()
+
+    def __repr__(self) -> str:
+        return f"AuthorizedKeys(keys={len(self._keys)})"
+
+
+if __name__ == '__main__':
+    import sys
+    if len(sys.argv) < 2:
+        print("Usage: python authorized_keys.py <path_to_authorized_keys>")
+        sys.exit(1)
+
+    path = sys.argv[1]
+    auth = AuthorizedKeys()
+    auth.load(path)
+    print(auth)

cloudmesh/ai/ssh/base.py

@@ -0,0 +1,183 @@
+# Copyright 2026 Gregor von Laszewski
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+import subprocess
+import socket
+
+from pathlib import Path
+from typing import List, Optional, Dict, Union
+from dataclasses import dataclass
+from fabric import Connection
+from cloudmesh.ai.common import logging as ai_log
+from .exceptions import SSHError, SSHConnectionError
+
+logger = ai_log.get_logger("ai.ssh.base")
+
+@dataclass
+class CommandResult:
+    """Structured result of a remote command execution."""
+    stdout: str
+    stderr: str
+    exit_code: int
+    command: str
+    host: str
+
+class SSHBase:
+    """Base class for SSH utilities providing shared execution and path logic."""
+
+    def __init__(self, debug: bool = False):
+        self.debug = debug
+        self._connection_pool: Dict[str, Connection] = {}
+
+    def _execute(self, command: List[str], input_data: Optional[str] = None, capture_output: bool = True) -> subprocess.CompletedProcess:
+        """Execute a system command using subprocess.run.
+
+        Args:
+            command: The command and arguments as a list.
+            input_data: Optional string to pass to stdin.
+            capture_output: Whether to capture stdout and stderr.
+
+        Returns:
+            The CompletedProcess object.
+        """
+        if self.debug:
+            logger.debug(f"Executing command: {' '.join(command)}")
+        
+        try:
+            return subprocess.run(
+                command,
+                input=input_data,
+                capture_output=capture_output,
+                text=True,
+                check=True
+            )
+        except subprocess.CalledProcessError as e:
+            logger.error(f"Command failed: {e.stderr}")
+            raise SSHError(f"System command failed: {e.stderr}") from e
+
+    def resolve_path(self, path: Union[str, Path]) -> Path:
+        """Expand and resolve a path."""
+        return Path(path).expanduser().resolve()
+
+    def is_port_open(self, host: str, port: int, timeout: float = 1.0) -> bool:
+        """Check if a port is open on a given host.
+
+        Args:
+            host: the hostname or IP address.
+            port: the port number.
+            timeout: timeout in seconds.
+
+        Returns:
+            bool: True if the port is open, False otherwise.
+        """
+        try:
+            import socket
+            with socket.create_connection((host, port), timeout=timeout):
+                return True
+        except (socket.timeout, ConnectionRefusedError, OSError):
+            return False
+
+    def _get_connection(self, host: str, user: Optional[str] = None) -> Connection:
+        """Get a Fabric connection from the pool or create a new one, with health check.
+
+        Args:
+            host: the hostname or alias.
+            user: optional username.
+
+        Returns:
+            Connection: a Fabric connection object.
+        """
+        pool_key = f"{user}@{host}" if user else host
+        conn = self._connection_pool.get(pool_key)
+        
+        if conn:
+            try:
+                # Heartbeat check: run a lightweight command to verify connection
+                conn.run("true", hide=True, timeout=5)
+            except Exception:
+                if self.debug:
+                    logger.debug(f"Connection for {pool_key} is stale, recreating...")
+                conn = None
+
+        if conn is None:
+            if self.debug:
+                logger.debug(f"Creating new connection for {pool_key}")
+            try:
+                conn = Connection(host=host, user=user)
+                self._connection_pool[pool_key] = conn
+            except Exception as e:
+                raise SSHConnectionError(f"Failed to create connection to {host}: {e}") from e
+            
+        return conn
+
+    def _run_remote(self, host: str, command: str, user: Optional[str] = None, use_sudo: bool = False, use_pty: bool = False) -> CommandResult:
+        """Execute a command on a remote host using Fabric.
+
+        Args:
+            host: the hostname or alias.
+            command: the command to execute.
+            user: optional username.
+            use_sudo: whether to use sudo for execution.
+            use_pty: whether to allocate a pseudo-terminal.
+
+        Returns:
+            CommandResult: structured result of the execution.
+        """
+        if self.debug:
+            logger.debug(f"Executing {'sudo ' if use_sudo else ''}remote command on {host} (pty={use_pty}): {command}")
+
+        try:
+            conn = self._get_connection(host, user)
+            if use_sudo:
+                result = conn.sudo(command, hide=True, pty=use_pty)
+            else:
+                result = conn.run(command, hide=True, pty=use_pty)
+            
+            return CommandResult(
+                stdout=result.stdout.strip(),
+                stderr=result.stderr.strip(),
+                exit_code=result.exited,
+                command=command,
+                host=host
+            )
+        except Exception as e:
+            logger.error(f"Fabric execution failed on {host}: {e}")
+            raise SSHConnectionError(f"Fabric execution failed on {host}: {e}") from e
+
+    def put(self, local_path: Union[str, Path], remote_path: Union[str, Path], host: str, user: Optional[str] = None) -> None:
+        """Upload a local file to a remote host."""
+        if self.debug:
+            logger.debug(f"Uploading {local_path} to {host}:{remote_path}")
+        try:
+            conn = self._get_connection(host, user)
+            conn.put(str(local_path), str(remote_path))
+        except Exception as e:
+            logger.error(f"Fabric put failed on {host}: {e}")
+            raise SSHConnectionError(f"Fabric put failed on {host}: {e}") from e
+
+    def get(self, remote_path: Union[str, Path], local_path: Union[str, Path], host: str, user: Optional[str] = None) -> None:
+        """Download a remote file to a local path."""
+        if self.debug:
+            logger.debug(f"Downloading {remote_path} from {host} to {local_path}")
+        try:
+            conn = self._get_connection(host, user)
+            conn.get(str(remote_path), str(local_path))
+        except Exception as e:
+            logger.error(f"Fabric get failed on {host}: {e}")
+            raise SSHConnectionError(f"Fabric get failed on {host}: {e}") from e
+
+    def close_connections(self) -> None:
+        """Close all active connections in the pool."""
+        if self.debug:
+            logger.debug("Closing all SSH connections in pool.")
+        for conn in self._connection_pool.values():
+            try:
+                conn.close()
+            except Exception as e:
+                logger.error(f"Error closing connection: {e}")
+        self._connection_pool.clear()