agent/device_driver.py

"""
Multi-vendor network device driver library.

Provides unified interface for device connectivity and configuration
across Cisco, Arista, Juniper, and other vendors using Netmiko and NAPALM.
"""

import logging
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, field
from enum import Enum
from datetime import datetime
import time

logger = logging.getLogger(__name__)

# Try to import Netmiko and NAPALM, fall back to mock mode if not available
try:
    from netmiko import ConnectHandler
    from netmiko.exceptions import NetmikoTimeoutException, NetmikoAuthenticationException
    NETMIKO_AVAILABLE = True
except ImportError:
    logger.warning("Netmiko not available - using mock mode")
    NETMIKO_AVAILABLE = False
    ConnectHandler = None
    NetmikoTimeoutException = Exception
    NetmikoAuthenticationException = Exception

try:
    import napalm
    from napalm.base.exceptions import ConnectionException, CommandErrorException
    NAPALM_AVAILABLE = True
except ImportError:
    logger.warning("NAPALM not available - using mock mode")
    NAPALM_AVAILABLE = False
    napalm = None
    ConnectionException = Exception
    CommandErrorException = Exception


class DeviceType(Enum):
    """Supported device types"""
    CISCO_IOS = "cisco_ios"
    CISCO_NXOS = "cisco_nxos"
    CISCO_XE = "cisco_xe"
    ARISTA_EOS = "arista_eos"
    JUNIPER_JUNOS = "juniper_junos"
    GENERIC_SSH = "linux"


class ConnectionStatus(Enum):
    """Device connection status"""
    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    FAILED = "failed"


@dataclass
class DeviceCredentials:
    """Device authentication credentials"""
    hostname: str
    username: str
    password: str
    device_type: DeviceType
    port: int = 22
    secret: Optional[str] = None  # Enable password
    timeout: int = 30
    session_log: Optional[str] = None


@dataclass
class DeviceConnection:
    """Active device connection"""
    credentials: DeviceCredentials
    status: ConnectionStatus = ConnectionStatus.DISCONNECTED
    connection: Any = None
    last_error: Optional[str] = None
    connected_at: Optional[datetime] = None


@dataclass
class CommandResult:
    """Result from device command execution"""
    command: str
    output: str
    success: bool
    error: Optional[str] = None
    duration_seconds: float = 0.0
    timestamp: datetime = field(default_factory=datetime.now)


@dataclass
class ConfigDeploymentResult:
    """Result from configuration deployment"""
    device_id: str
    success: bool
    changes_applied: bool
    config_before: Optional[str] = None
    config_after: Optional[str] = None
    commands_sent: List[str] = field(default_factory=list)
    output: Optional[str] = None
    error: Optional[str] = None
    duration_seconds: float = 0.0
    rollback_available: bool = False


class DeviceDriver:
    """
    Unified device driver for multi-vendor network devices.
    
    Supports Cisco IOS/NXOS/XE, Arista EOS, Juniper JunOS via Netmiko and NAPALM.
    Falls back to mock mode when libraries unavailable (for testing).
    """
    
    def __init__(self, use_napalm: bool = True):
        """
        Initialize device driver.
        
        Args:
            use_napalm: Prefer NAPALM over Netmiko when available
        """
        self.use_napalm = use_napalm and NAPALM_AVAILABLE
        self.use_netmiko = NETMIKO_AVAILABLE
        self.mock_mode = not (self.use_napalm or self.use_netmiko)
        
        self.connections: Dict[str, DeviceConnection] = {}
        
        if self.mock_mode:
            logger.warning("DeviceDriver in MOCK MODE - no real device connections")
        elif self.use_napalm:
            logger.info("DeviceDriver using NAPALM (preferred)")
        else:
            logger.info("DeviceDriver using Netmiko")
    
    def connect(self, credentials: DeviceCredentials) -> DeviceConnection:
        """
        Establish connection to device.
        
        Args:
            credentials: Device credentials and connection info
            
        Returns:
            DeviceConnection object
        """
        device_id = credentials.hostname
        
        if device_id in self.connections:
            conn = self.connections[device_id]
            if conn.status == ConnectionStatus.CONNECTED:
                logger.info(f"Reusing existing connection to {device_id}")
                return conn
        
        # Create new connection object
        conn = DeviceConnection(credentials=credentials, status=ConnectionStatus.CONNECTING)
        
        if self.mock_mode:
            # Mock connection - always succeeds
            time.sleep(0.1)  # Simulate connection delay
            conn.status = ConnectionStatus.CONNECTED
            conn.connected_at = datetime.now()
            conn.connection = {"mock": True, "device_id": device_id}
            logger.info(f"MOCK: Connected to {device_id}")
            
        elif self.use_napalm:
            # Use NAPALM
            try:
                driver_map = {
                    DeviceType.CISCO_IOS: 'ios',
                    DeviceType.CISCO_NXOS: 'nxos',
                    DeviceType.CISCO_XE: 'ios',
                    DeviceType.ARISTA_EOS: 'eos',
                    DeviceType.JUNIPER_JUNOS: 'junos',
                }
                
                driver_name = driver_map.get(credentials.device_type)
                if not driver_name:
                    raise ValueError(f"Unsupported device type for NAPALM: {credentials.device_type}")
                
                driver = napalm.get_network_driver(driver_name)
                
                device = driver(
                    hostname=credentials.hostname,
                    username=credentials.username,
                    password=credentials.password,
                    timeout=credentials.timeout,
                    optional_args={'port': credentials.port}
                )
                
                device.open()
                
                conn.connection = device
                conn.status = ConnectionStatus.CONNECTED
                conn.connected_at = datetime.now()
                logger.info(f"NAPALM: Connected to {device_id}")
                
            except Exception as e:
                conn.status = ConnectionStatus.FAILED
                conn.last_error = str(e)
                logger.error(f"NAPALM connection failed to {device_id}: {e}")
        
        else:
            # Use Netmiko
            try:
                device_params = {
                    'device_type': credentials.device_type.value,
                    'host': credentials.hostname,
                    'username': credentials.username,
                    'password': credentials.password,
                    'port': credentials.port,
                    'timeout': credentials.timeout,
                }
                
                if credentials.secret:
                    device_params['secret'] = credentials.secret
                
                if credentials.session_log:
                    device_params['session_log'] = credentials.session_log
                
                device = ConnectHandler(**device_params)
                
                conn.connection = device
                conn.status = ConnectionStatus.CONNECTED
                conn.connected_at = datetime.now()
                logger.info(f"Netmiko: Connected to {device_id}")
                
            except Exception as e:
                conn.status = ConnectionStatus.FAILED
                conn.last_error = str(e)
                logger.error(f"Netmiko connection failed to {device_id}: {e}")
        
        self.connections[device_id] = conn
        return conn
    
    def disconnect(self, device_id: str):
        """
        Close connection to device.
        
        Args:
            device_id: Device hostname/identifier
        """
        if device_id not in self.connections:
            return
        
        conn = self.connections[device_id]
        
        if conn.status != ConnectionStatus.CONNECTED:
            return
        
        try:
            if self.mock_mode:
                logger.info(f"MOCK: Disconnected from {device_id}")
            
            elif self.use_napalm:
                conn.connection.close()
                logger.info(f"NAPALM: Disconnected from {device_id}")
            
            elif self.use_netmiko:
                conn.connection.disconnect()
                logger.info(f"Netmiko: Disconnected from {device_id}")
            
            conn.status = ConnectionStatus.DISCONNECTED
            
        except Exception as e:
            logger.error(f"Error disconnecting from {device_id}: {e}")
    
    def send_command(self, device_id: str, command: str) -> CommandResult:
        """
        Send single command to device.
        
        Args:
            device_id: Device identifier
            command: Command to execute
            
        Returns:
            CommandResult with output and status
        """
        if device_id not in self.connections:
            return CommandResult(
                command=command,
                output="",
                success=False,
                error="Device not connected"
            )
        
        conn = self.connections[device_id]
        
        if conn.status != ConnectionStatus.CONNECTED:
            return CommandResult(
                command=command,
                output="",
                success=False,
                error=f"Device in {conn.status} state"
            )
        
        start_time = time.time()
        
        try:
            if self.mock_mode:
                # Mock response
                output = f"MOCK OUTPUT for: {command}\n"
                output += "Device successfully executed command (simulated)\n"
                
            elif self.use_napalm:
                # NAPALM CLI command
                output_dict = conn.connection.cli([command])
                output = output_dict.get(command, "")
            
            else:
                # Netmiko
                output = conn.connection.send_command(command)
            
            duration = time.time() - start_time
            
            return CommandResult(
                command=command,
                output=output,
                success=True,
                duration_seconds=duration
            )
            
        except Exception as e:
            duration = time.time() - start_time
            logger.error(f"Command failed on {device_id}: {e}")
            
            return CommandResult(
                command=command,
                output="",
                success=False,
                error=str(e),
                duration_seconds=duration
            )
    
    def get_config(self, device_id: str, config_type: str = "running") -> Optional[str]:
        """
        Retrieve device configuration.
        
        Args:
            device_id: Device identifier
            config_type: Type of config (running, startup, candidate)
            
        Returns:
            Configuration text or None on error
        """
        if device_id not in self.connections:
            logger.error(f"Device {device_id} not connected")
            return None
        
        conn = self.connections[device_id]
        
        try:
            if self.mock_mode:
                # Return mock config
                return f"! Mock {config_type} configuration for {device_id}\nhostname {device_id}\n!\nend\n"
            
            elif self.use_napalm:
                configs = conn.connection.get_config(retrieve=config_type)
                return configs.get(config_type, "")
            
            else:
                # Netmiko - device-specific commands
                if config_type == "running":
                    result = self.send_command(device_id, "show running-config")
                elif config_type == "startup":
                    result = self.send_command(device_id, "show startup-config")
                else:
                    result = self.send_command(device_id, f"show {config_type}-config")
                
                return result.output if result.success else None
                
        except Exception as e:
            logger.error(f"Failed to get {config_type} config from {device_id}: {e}")
            return None
    
    def deploy_config(self, device_id: str, config: str, 
                     dry_run: bool = False,
                     replace: bool = False) -> ConfigDeploymentResult:
        """
        Deploy configuration to device.
        
        Args:
            device_id: Device identifier
            config: Configuration to deploy (as string)
            dry_run: If True, compare only (don't apply)
            replace: If True, replace entire config; if False, merge
            
        Returns:
            ConfigDeploymentResult with deployment status
        """
        if device_id not in self.connections:
            return ConfigDeploymentResult(
                device_id=device_id,
                success=False,
                changes_applied=False,
                error="Device not connected"
            )
        
        conn = self.connections[device_id]
        start_time = time.time()
        
        # Get pre-deployment config for rollback
        config_before = self.get_config(device_id, "running")
        
        try:
            if self.mock_mode:
                # Mock deployment
                time.sleep(0.2)  # Simulate deployment time
                
                # Parse config into commands
                commands = [line.strip() for line in config.split('\n') 
                           if line.strip() and not line.strip().startswith('!')]
                
                output = f"MOCK: Configuration deployed to {device_id}\n"
                output += f"Commands sent: {len(commands)}\n"
                output += f"Mode: {'replace' if replace else 'merge'}\n"
                
                result = ConfigDeploymentResult(
                    device_id=device_id,
                    success=True,
                    changes_applied=not dry_run,
                    config_before=config_before,
                    config_after=config if not dry_run else config_before,
                    commands_sent=commands,
                    output=output,
                    duration_seconds=time.time() - start_time,
                    rollback_available=True
                )
                
            elif self.use_napalm:
                # NAPALM config deployment
                if replace:
                    conn.connection.load_replace_candidate(config=config)
                else:
                    conn.connection.load_merge_candidate(config=config)
                
                # Get diff
                diff = conn.connection.compare_config()
                
                if dry_run:
                    # Discard candidate
                    conn.connection.discard_config()
                    
                    result = ConfigDeploymentResult(
                        device_id=device_id,
                        success=True,
                        changes_applied=False,
                        config_before=config_before,
                        config_after=config_before,
                        output=f"Dry run - diff:\n{diff}",
                        duration_seconds=time.time() - start_time,
                        rollback_available=False
                    )
                else:
                    # Commit config
                    conn.connection.commit_config()
                    
                    config_after = self.get_config(device_id, "running")
                    
                    result = ConfigDeploymentResult(
                        device_id=device_id,
                        success=True,
                        changes_applied=True,
                        config_before=config_before,
                        config_after=config_after,
                        output=f"Config committed - diff:\n{diff}",
                        duration_seconds=time.time() - start_time,
                        rollback_available=True
                    )
            
            else:
                # Netmiko config deployment
                # Parse config into commands
                commands = [line.strip() for line in config.split('\n')
                           if line.strip() and not line.strip().startswith('!')]
                
                if dry_run:
                    output = "Dry run - commands would be:\n" + "\n".join(commands)
                    
                    result = ConfigDeploymentResult(
                        device_id=device_id,
                        success=True,
                        changes_applied=False,
                        config_before=config_before,
                        config_after=config_before,
                        commands_sent=commands,
                        output=output,
                        duration_seconds=time.time() - start_time,
                        rollback_available=False
                    )
                else:
                    # Send config
                    output = conn.connection.send_config_set(commands)
                    
                    # Save config
                    if conn.credentials.device_type in [DeviceType.CISCO_IOS, DeviceType.CISCO_XE]:
                        conn.connection.save_config()
                    
                    config_after = self.get_config(device_id, "running")
                    
                    result = ConfigDeploymentResult(
                        device_id=device_id,
                        success=True,
                        changes_applied=True,
                        config_before=config_before,
                        config_after=config_after,
                        commands_sent=commands,
                        output=output,
                        duration_seconds=time.time() - start_time,
                        rollback_available=True
                    )
            
            logger.info(f"Config deployment to {device_id}: {'dry-run' if dry_run else 'committed'}")
            return result
            
        except Exception as e:
            logger.error(f"Config deployment failed on {device_id}: {e}")
            
            # Try to rollback on error
            if self.use_napalm and not dry_run:
                try:
                    conn.connection.discard_config()
                    logger.info(f"Rolled back failed config on {device_id}")
                except:
                    pass
            
            return ConfigDeploymentResult(
                device_id=device_id,
                success=False,
                changes_applied=False,
                config_before=config_before,
                error=str(e),
                duration_seconds=time.time() - start_time,
                rollback_available=False
            )
    
    def rollback_config(self, device_id: str, config: str) -> ConfigDeploymentResult:
        """
        Rollback to previous configuration.
        
        Args:
            device_id: Device identifier
            config: Previous configuration to restore
            
        Returns:
            ConfigDeploymentResult
        """
        logger.warning(f"Rolling back configuration on {device_id}")
        
        return self.deploy_config(
            device_id=device_id,
            config=config,
            dry_run=False,
            replace=True  # Replace entire config with backup
        )
    
    def verify_connectivity(self, device_id: str) -> bool:
        """
        Verify device is reachable and responsive.
        
        Args:
            device_id: Device identifier
            
        Returns:
            True if device responsive
        """
        if device_id not in self.connections:
            return False
        
        conn = self.connections[device_id]
        
        if conn.status != ConnectionStatus.CONNECTED:
            return False
        
        try:
            # Send simple command to verify
            result = self.send_command(device_id, "show version")
            return result.success
            
        except Exception as e:
            logger.error(f"Connectivity check failed for {device_id}: {e}")
            return False
    
    def get_device_facts(self, device_id: str) -> Optional[Dict[str, Any]]:
        """
        Get device facts (hostname, model, version, uptime, etc.).
        
        Args:
            device_id: Device identifier
            
        Returns:
            Dictionary of facts or None
        """
        if device_id not in self.connections:
            return None
        
        conn = self.connections[device_id]
        
        try:
            if self.mock_mode:
                return {
                    'hostname': device_id,
                    'vendor': 'Mock Vendor',
                    'model': 'Virtual Device',
                    'os_version': '1.0.0',
                    'uptime': 3600,
                    'serial_number': f'MOCK{device_id.upper()[:8]}',
                    'interface_list': ['GigabitEthernet0/1', 'GigabitEthernet0/2']
                }
            
            elif self.use_napalm:
                return conn.connection.get_facts()
            
            else:
                # Parse from show version (device-specific)
                result = self.send_command(device_id, "show version")
                if not result.success:
                    return None
                
                # Basic parsing (would need vendor-specific logic)
                return {
                    'hostname': device_id,
                    'show_version_output': result.output
                }
                
        except Exception as e:
            logger.error(f"Failed to get facts from {device_id}: {e}")
            return None
    
    def disconnect_all(self):
        """Disconnect from all devices"""
        for device_id in list(self.connections.keys()):
            self.disconnect(device_id)
        
        self.connections.clear()
        logger.info("Disconnected from all devices")