agent/batfish_client.py

"""
Batfish Integration for Static Network Analysis
Pre-deployment validation of configs without touching live gear
"""

import os
import logging
from typing import Dict, List, Optional, Any
from pathlib import Path
import tempfile
import shutil

logger = logging.getLogger(__name__)


class BatfishAnalysis:
    """Results from Batfish static analysis"""
    def __init__(self):
        self.reachability_passed = False
        self.routing_loops = []
        self.acl_issues = []
        self.undefined_references = []
        self.unused_structures = []
        self.forwarding_errors = []
        self.all_passed = False
    
    def to_dict(self) -> Dict[str, Any]:
        return {
            'reachability_passed': self.reachability_passed,
            'routing_loops': self.routing_loops,
            'acl_issues': self.acl_issues,
            'undefined_references': self.undefined_references,
            'unused_structures': self.unused_structures,
            'forwarding_errors': self.forwarding_errors,
            'all_passed': self.all_passed
        }


class BatfishClient:
    """
    Client for Batfish network analysis
    Performs static analysis on network configurations
    """
    
    def __init__(self, host: str = "localhost", use_batfish: bool = True):
        """
        Initialize Batfish client
        
        Args:
            host: Batfish service hostname
            use_batfish: Enable Batfish (False for mock mode)
        """
        self.host = host
        self.use_batfish = use_batfish
        self.mock_mode = True
        
        if use_batfish:
            try:
                from pybatfish.client.commands import bf_session, bf_set_network, bf_init_snapshot
                from pybatfish.question import bfq, load_questions
                from pybatfish.datamodel import HeaderConstraints
                
                self.bf_session = bf_session
                self.bf_set_network = bf_set_network
                self.bf_init_snapshot = bf_init_snapshot
                self.bfq = bfq
                self.load_questions = load_questions
                self.HeaderConstraints = HeaderConstraints
                
                # Connect to Batfish service
                bf_session.host = host
                load_questions()
                
                self.mock_mode = False
                logger.info(f"Connected to Batfish at {host}")
                
            except ImportError:
                logger.warning("pybatfish not installed - using mock mode")
                logger.info("Install with: pip install pybatfish")
            except Exception as e:
                logger.warning(f"Failed to connect to Batfish: {e}")
                logger.info("Using mock mode")
    
    def analyze_configs(
        self,
        configs: Dict[str, str],
        network_name: str = "overgrowth-analysis"
    ) -> BatfishAnalysis:
        """
        Analyze network configurations
        
        Args:
            configs: Dict mapping device names to config strings
            network_name: Name for this analysis snapshot
            
        Returns:
            BatfishAnalysis with results
        """
        if self.mock_mode:
            return self._mock_analysis(configs)
        
        analysis = BatfishAnalysis()
        
        try:
            # Create temp directory for configs
            snapshot_dir = Path(tempfile.mkdtemp(prefix="batfish_"))
            configs_dir = snapshot_dir / "configs"
            configs_dir.mkdir()
            
            # Write configs to files
            for device_name, config in configs.items():
                config_file = configs_dir / f"{device_name}.cfg"
                config_file.write_text(config)
            
            logger.info(f"Created snapshot with {len(configs)} device configs")
            
            # Initialize Batfish snapshot
            self.bf_set_network(network_name)
            self.bf_init_snapshot(str(snapshot_dir), name="candidate", overwrite=True)
            
            # Run analysis questions
            analysis = self._run_batfish_questions()
            
            # Cleanup temp directory
            shutil.rmtree(snapshot_dir)
            
        except Exception as e:
            logger.error(f"Batfish analysis failed: {e}")
            analysis.all_passed = False
        
        return analysis
    
    def _run_batfish_questions(self) -> BatfishAnalysis:
        """Run Batfish analysis questions"""
        analysis = BatfishAnalysis()
        
        try:
            # 1. Check for undefined references
            logger.info("Checking for undefined references...")
            undef_refs = self.bfq.undefinedReferences().answer().frame()
            if not undef_refs.empty:
                analysis.undefined_references = undef_refs.to_dict('records')
                logger.warning(f"Found {len(undef_refs)} undefined references")
            
            # 2. Check for unused structures
            logger.info("Checking for unused structures...")
            unused = self.bfq.unusedStructures().answer().frame()
            if not unused.empty:
                analysis.unused_structures = unused.to_dict('records')
                logger.info(f"Found {len(unused)} unused structures")
            
            # 3. Check routing loops
            logger.info("Checking for routing loops...")
            loops = self.bfq.detectLoops().answer().frame()
            if not loops.empty:
                analysis.routing_loops = loops.to_dict('records')
                logger.error(f"Found {len(loops)} routing loops!")
            
            # 4. Validate reachability
            logger.info("Validating reachability...")
            reach = self.bfq.reachability().answer().frame()
            analysis.reachability_passed = reach.empty or reach['Action'].str.contains('ACCEPT').any()
            
            # 5. Check for forwarding errors
            logger.info("Checking for forwarding errors...")
            fwd_errors = self.bfq.detectForwardingLoops().answer().frame()
            if not fwd_errors.empty:
                analysis.forwarding_errors = fwd_errors.to_dict('records')
                logger.error(f"Found {len(fwd_errors)} forwarding errors")
            
            # Overall pass/fail
            analysis.all_passed = (
                len(analysis.undefined_references) == 0 and
                len(analysis.routing_loops) == 0 and
                len(analysis.forwarding_errors) == 0 and
                analysis.reachability_passed
            )
            
            if analysis.all_passed:
                logger.info("✓ Batfish analysis PASSED - no critical issues")
            else:
                logger.warning("✗ Batfish analysis found issues")
            
        except Exception as e:
            logger.error(f"Error running Batfish questions: {e}")
            analysis.all_passed = False
        
        return analysis
    
    def _mock_analysis(self, configs: Dict[str, str]) -> BatfishAnalysis:
        """Mock analysis when Batfish unavailable"""
        logger.info("Running mock Batfish analysis...")
        
        analysis = BatfishAnalysis()
        
        # Simple heuristic checks
        for device, config in configs.items():
            # Check for basic issues in config
            if "no ip routing" in config.lower():
                analysis.forwarding_errors.append({
                    'device': device,
                    'issue': 'Routing disabled',
                    'severity': 'WARNING'
                })
            
            # Check for undefined references (simple regex)
            import re
            vlan_refs = re.findall(r'switchport access vlan (\d+)', config, re.IGNORECASE)
            vlan_defs = re.findall(r'vlan (\d+)', config, re.IGNORECASE)
            
            undefined_vlans = set(vlan_refs) - set(vlan_defs)
            for vlan in undefined_vlans:
                analysis.undefined_references.append({
                    'device': device,
                    'type': 'VLAN',
                    'name': vlan,
                    'severity': 'ERROR'
                })
        
        # Mock passes if no critical errors
        analysis.all_passed = len(analysis.undefined_references) == 0
        analysis.reachability_passed = True
        
        logger.info(f"Mock analysis complete: {len(configs)} configs checked")
        return analysis
    
    def validate_acl_behavior(
        self,
        src: str,
        dst: str,
        protocol: str = "TCP",
        dst_port: int = 80
    ) -> bool:
        """
        Test if traffic is permitted by ACLs
        
        Args:
            src: Source IP or network
            dst: Destination IP or network
            protocol: IP protocol (TCP, UDP, ICMP)
            dst_port: Destination port number
            
        Returns:
            True if traffic is permitted
        """
        if self.mock_mode:
            logger.info(f"Mock ACL check: {src} -> {dst}:{dst_port}/{protocol} = PERMIT")
            return True
        
        try:
            # Build header constraints
            headers = self.HeaderConstraints(
                srcIps=src,
                dstIps=dst,
                ipProtocols=[protocol],
                dstPorts=str(dst_port)
            )
            
            # Query reachability with constraints
            result = self.bfq.reachability(headers=headers).answer().frame()
            
            # Check if any flow is accepted
            permitted = not result.empty and result['Action'].str.contains('ACCEPT').any()
            
            logger.info(f"ACL check: {src} -> {dst}:{dst_port}/{protocol} = {'PERMIT' if permitted else 'DENY'}")
            return permitted
            
        except Exception as e:
            logger.error(f"ACL validation failed: {e}")
            return False
    
    def find_routing_issues(self) -> List[Dict[str, Any]]:
        """
        Find routing protocol issues
        
        Returns:
            List of routing issues found
        """
        if self.mock_mode:
            return []
        
        issues = []
        
        try:
            # Check for BGP issues
            logger.info("Checking BGP sessions...")
            bgp_edges = self.bfq.bgpEdges().answer().frame()
            
            for idx, edge in bgp_edges.iterrows():
                if edge.get('Status') != 'ESTABLISHED':
                    issues.append({
                        'type': 'BGP_SESSION_DOWN',
                        'node': edge.get('Node'),
                        'remote': edge.get('Remote_Node'),
                        'severity': 'ERROR'
                    })
            
            # Check for OSPF issues
            logger.info("Checking OSPF neighbors...")
            ospf_edges = self.bfq.ospfEdges().answer().frame()
            
            # Look for missing adjacencies
            # (This is simplified - real check would be more complex)
            
        except Exception as e:
            logger.error(f"Error finding routing issues: {e}")
        
        return issues
    
    def test_failover_scenario(
        self,
        failed_device: str,
        src: str,
        dst: str
    ) -> bool:
        """
        Test if network maintains connectivity when device fails
        
        Args:
            failed_device: Device to simulate failure
            src: Source IP for reachability test
            dst: Destination IP for reachability test
            
        Returns:
            True if network survives failure
        """
        if self.mock_mode:
            logger.info(f"Mock failover test: network survives {failed_device} failure")
            return True
        
        try:
            # Deactivate device
            logger.info(f"Simulating failure of {failed_device}...")
            
            # Test reachability without failed device
            headers = self.HeaderConstraints(srcIps=src, dstIps=dst)
            result = self.bfq.reachability(
                headers=headers,
                forbiddenTransitNodes=failed_device
            ).answer().frame()
            
            survives = not result.empty and result['Action'].str.contains('ACCEPT').any()
            
            if survives:
                logger.info(f"✓ Network survives {failed_device} failure")
            else:
                logger.warning(f"✗ Network fails when {failed_device} is down")
            
            return survives
            
        except Exception as e:
            logger.error(f"Failover test failed: {e}")
            return False
    
    def generate_config_recommendations(self, analysis: BatfishAnalysis) -> List[str]:
        """
        Generate recommendations based on analysis results
        
        Args:
            analysis: Batfish analysis results
            
        Returns:
            List of human-readable recommendations
        """
        recommendations = []
        
        if analysis.undefined_references:
            recommendations.append(
                f"Fix {len(analysis.undefined_references)} undefined references "
                "(VLANs, ACLs, route-maps referenced but not defined)"
            )
        
        if analysis.routing_loops:
            recommendations.append(
                f"Resolve {len(analysis.routing_loops)} routing loops "
                "(will cause packet storms and network meltdown)"
            )
        
        if analysis.forwarding_errors:
            recommendations.append(
                f"Fix {len(analysis.forwarding_errors)} forwarding errors "
                "(traffic will be dropped or blackholed)"
            )
        
        if analysis.unused_structures:
            recommendations.append(
                f"Consider removing {len(analysis.unused_structures)} unused structures "
                "(cleanup for maintainability)"
            )
        
        if not analysis.reachability_passed:
            recommendations.append(
                "Reachability test failed - verify routing and ACLs allow required traffic"
            )
        
        if not recommendations:
            recommendations.append("✓ No issues found - configuration looks good!")
        
        return recommendations