cisco-conf-generator/backend/app/ssh/__init__.py

"""SSH and Netmiko integration for device push"""
from typing import Tuple, Optional
import asyncio
from netmiko import ConnectHandler, NetmikoTimeoutException, NetmikoAuthenticationException
import logging

logger = logging.getLogger(__name__)


class SSHConnectionError(Exception):
    """SSH connection failed"""
    pass


class SSHExecutor:
    """
    Handles SSH connections to Cisco devices via Netmiko.
    
    CRITICAL SECURITY:
    - Credentials are passed in-memory only, never stored
    - SSH connection is established per-push
    - All commands are logged for audit
    """
    
    def __init__(self, hostname: str, ip_address: str, username: str, password: str,
                 device_type: str = "cisco_ios", ssh_port: int = 22, timeout: int = 30):
        """
        Initialize SSH executor.
        
        Args:
            hostname: Device hostname
            ip_address: Device IP/FQDN
            username: SSH username
            password: SSH password
            device_type: Netmiko device type (e.g., "cisco_ios", "cisco_xe")
            ssh_port: SSH port
            timeout: Connection timeout in seconds
        """
        self.hostname = hostname
        self.ip_address = ip_address
        self.username = username
        self.password = password
        self.device_type = device_type
        self.ssh_port = ssh_port
        self.timeout = timeout
        self.connection = None
    
    def connect(self) -> bool:
        """
        Establish SSH connection.
        
        Returns:
            True if successful
        
        Raises:
            SSHConnectionError if connection fails
        """
        try:
            device = {
                "device_type": self.device_type,
                "host": self.ip_address,
                "username": self.username,
                "password": self.password,
                "port": self.ssh_port,
                "timeout": self.timeout,
            }
            
            self.connection = ConnectHandler(**device)
            logger.info(f"SSH connected to {self.hostname} ({self.ip_address})")
            return True
        
        except NetmikoAuthenticationException as e:
            logger.error(f"SSH authentication failed for {self.hostname}: {e}")
            raise SSHConnectionError(f"Authentication failed: {e}")
        
        except NetmikoTimeoutException as e:
            logger.error(f"SSH timeout connecting to {self.hostname}: {e}")
            raise SSHConnectionError(f"Connection timeout: {e}")
        
        except Exception as e:
            logger.error(f"SSH connection error for {self.hostname}: {e}")
            raise SSHConnectionError(f"Connection failed: {e}")
    
    def send_commands(self, commands: list) -> str:
        """
        Send commands to device and return output.
        
        Args:
            commands: List of CLI commands
        
        Returns:
            Device output
        """
        if not self.connection:
            raise SSHConnectionError("Not connected")
        
        output = ""
        try:
            for cmd in commands:
                logger.debug(f"Sending: {cmd}")
                output += self.connection.send_command(cmd)
                output += "\n"
        
        except Exception as e:
            logger.error(f"Command execution failed: {e}")
            raise SSHConnectionError(f"Command failed: {e}")
        
        return output
    
    def get_running_config(self) -> str:
        """
        Retrieve running configuration (for backup before push).
        
        Returns:
            Full running-config
        """
        if not self.connection:
            raise SSHConnectionError("Not connected")
        
        try:
            config = self.connection.send_command("show running-config")
            logger.info(f"Retrieved running-config from {self.hostname}")
            return config
        
        except Exception as e:
            logger.error(f"Failed to get running-config: {e}")
            raise SSHConnectionError(f"Get config failed: {e}")
    
    def save_running_config(self) -> bool:
        """
        Save running-config to startup-config.
        
        Returns:
            True if successful
        """
        if not self.connection:
            raise SSHConnectionError("Not connected")
        
        try:
            output = self.connection.send_command("write memory")
            logger.info(f"Saved config on {self.hostname}")
            return "OK" in output or "successfully" in output.lower()
        
        except Exception as e:
            logger.error(f"Failed to save config: {e}")
            raise SSHConnectionError(f"Save failed: {e}")
    
    def disconnect(self):
        """Close SSH connection"""
        if self.connection:
            try:
                self.connection.disconnect()
                logger.info(f"SSH disconnected from {self.hostname}")
            except Exception as e:
                logger.warning(f"Error disconnecting from {self.hostname}: {e}")
        
        self.connection = None
    
    def __enter__(self):
        """Context manager support"""
        self.connect()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager cleanup"""
        self.disconnect()


async def push_to_device_async(
    hostname: str,
    ip_address: str,
    commands: list,
    username: str,
    password: str,
    dry_run: bool = False,
    save_config: bool = True,
) -> Tuple[bool, str, Optional[str]]:
    """
    Push commands to device asynchronously.
    
    Args:
        hostname, ip_address, commands, username, password: Device info
        dry_run: If True, only validate, don't execute
        save_config: If True, backup running-config before push
    
    Returns:
        (success, message, error_msg)
    """
    executor = SSHExecutor(hostname, ip_address, username, password)
    
    try:
        # Run in thread pool (netmiko is blocking)
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, executor.connect)
        
        # Get backup if requested
        backup = None
        if save_config:
            backup = await loop.run_in_executor(None, executor.get_running_config)
        
        # Dry run mode
        if dry_run:
            executor.disconnect()
            return True, "Dry-run completed. Configuration validated.", None
        
        # Send commands
        output = await loop.run_in_executor(None, executor.send_commands, commands)
        
        # Save config
        if save_config:
            await loop.run_in_executor(None, executor.save_running_config)
        
        executor.disconnect()
        return True, "Configuration pushed successfully", None
    
    except SSHConnectionError as e:
        executor.disconnect()
        return False, "", str(e)
    
    except Exception as e:
        executor.disconnect()
        logger.error(f"Unexpected error during push: {e}")
        return False, "", str(e)