#!/usr/bin/env python3
"""
Complete Network Scanner with pfSense Integration
Combines network scanning with pfSense-specific features
"""

import json
import logging
import argparse
from datetime import datetime
from network_scanner import NetworkScanner, NetworkSegment
from pfsense_scanner import PfSenseScanner

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class IntegratedNetworkScanner:
    """Enhanced network scanner with pfSense integration"""
    
    def __init__(self, config: dict):
        self.config = config
        self.base_scanner = NetworkScanner(config)
        self.pfsense_devices = []
    
    def scan_all(self):
        """Perform complete network scan including pfSense devices"""
        logger.info("Starting integrated network scan...")
        
        # Run base network scan
        self.base_scanner.scan_all()
        
        # Identify and enhance pfSense devices
        self._scan_pfsense_devices()
        
        logger.info("Integrated scan complete")
    
    def _scan_pfsense_devices(self):
        """Find and deeply scan pfSense devices"""
        logger.info("Looking for pfSense devices...")
        
        # Check configured special devices
        special_devices = self.config.get('special_devices', {})
        
        for segment in self.base_scanner.segments:
            for device in segment.devices:
                is_pfsense = False
                
                # Check if device is marked as pfSense in config
                if device.ip in special_devices:
                    if special_devices[device.ip].get('type') == 'firewall' or \
                       special_devices[device.ip].get('os') == 'pfSense':
                        is_pfsense = True
                
                # Check if hostname contains pfsense
                if device.hostname and 'pfsense' in device.hostname.lower():
                    is_pfsense = True
                
                # Check if device looks like a pfSense (has many routes and ports 80/443)
                if device.routes and len(device.routes) > 3 and \
                   80 in device.open_ports and 443 in device.open_ports:
                    is_pfsense = True
                
                if is_pfsense and device.ssh_accessible:
                    logger.info(f"Enhanced scanning pfSense device: {device.ip}")
                    self._enhance_pfsense_device(device)
    
    def _enhance_pfsense_device(self, device):
        """Enhance device info with pfSense-specific data"""
        try:
            scanner = PfSenseScanner(
                device.ip,
                self.config.get('ssh_user', 'root'),
                self.config.get('ssh_key_path')
            )
            
            pfsense_info = scanner.get_full_info()
            
            # Merge pfSense-specific info into device
            device.device_type = 'firewall'
            device.os_type = 'pfSense (FreeBSD)'
            
            # Store additional pfSense data
            if not hasattr(device, 'pfsense_info'):
                device.__dict__['pfsense_info'] = pfsense_info
            
            # Add DHCP leases to known devices
            dhcp_leases = pfsense_info.get('dhcp_leases', [])
            if dhcp_leases:
                logger.info(f"Found {len(dhcp_leases)} DHCP leases on {device.ip}")
            
            # Add VPN info
            vpn_info = pfsense_info.get('vpn', {})
            wireguard = vpn_info.get('wireguard', [])
            if wireguard:
                logger.info(f"Found {len(wireguard)} WireGuard tunnels on {device.ip}")
            
            self.pfsense_devices.append(device)
            
        except Exception as e:
            logger.error(f"Error enhancing pfSense device {device.ip}: {e}")
    
    def export_json(self, filename: str):
        """Export enhanced results to JSON"""
        data = {
            'scan_timestamp': datetime.now().isoformat(),
            'segments': []
        }
        
        for segment in self.base_scanner.segments:
            segment_data = {
                'name': segment.name,
                'cidr': segment.cidr,
                'gateway': segment.gateway,
                'is_vpn': segment.is_vpn,
                'devices': []
            }
            
            for device in segment.devices:
                device_data = {
                    'ip': device.ip,
                    'hostname': device.hostname,
                    'mac': device.mac,
                    'manufacturer': device.manufacturer,
                    'os_type': device.os_type,
                    'os_version': device.os_version,
                    'device_type': device.device_type,
                    'open_ports': device.open_ports,
                    'ssh_accessible': device.ssh_accessible,
                    'services': device.services,
                    'routes': device.routes,
                    'interfaces': device.interfaces
                }
                
                # Add pfSense-specific info if available
                if hasattr(device, 'pfsense_info'):
                    device_data['pfsense_info'] = device.__dict__['pfsense_info']
                
                segment_data['devices'].append(device_data)
            
            data['segments'].append(segment_data)
        
        with open(filename, 'w') as f:
            json.dump(data, f, indent=2)
        
        logger.info(f"Exported results to {filename}")
    
    def print_summary(self):
        """Print enhanced summary"""
        print("\n" + "="*80)
        print("INTEGRATED NETWORK SCAN SUMMARY")
        print("="*80)
        
        total_devices = sum(len(seg.devices) for seg in self.base_scanner.segments)
        print(f"\nTotal Segments: {len(self.base_scanner.segments)}")
        print(f"Total Devices: {total_devices}")
        print(f"pfSense Devices: {len(self.pfsense_devices)}")
        
        for segment in self.base_scanner.segments:
            print(f"\n{'='*80}")
            print(f"📡 Network: {segment.name} ({segment.cidr})")
            if segment.is_vpn:
                print("   🔐 VPN Network")
            print(f"   Devices: {len(segment.devices)}")
            
            for device in segment.devices:
                icon = self._get_device_icon(device.device_type)
                print(f"\n   {icon} {device.ip}")
                
                if device.hostname:
                    print(f"      Hostname: {device.hostname}")
                if device.mac:
                    print(f"      MAC: {device.mac}")
                if device.device_type:
                    print(f"      Type: {device.device_type}")
                if device.os_type:
                    print(f"      OS: {device.os_type} {device.os_version or ''}")
                if device.open_ports:
                    print(f"      Ports: {', '.join(map(str, device.open_ports))}")
                if device.ssh_accessible:
                    print(f"      ✓ SSH Accessible")
                
                # Show pfSense-specific info
                if hasattr(device, 'pfsense_info'):
                    pfsense = device.__dict__['pfsense_info']
                    print(f"      🛡️  pfSense Firewall:")
                    print(f"         Interfaces: {len(pfsense.get('interfaces', []))}")
                    print(f"         Routes: {len(pfsense.get('routes', []))}")
                    print(f"         DHCP Leases: {len(pfsense.get('dhcp_leases', []))}")
                    
                    vpn = pfsense.get('vpn', {})
                    wg_count = len(vpn.get('wireguard', []))
                    ovpn_count = len(vpn.get('openvpn', []))
                    if wg_count:
                        print(f"         WireGuard Tunnels: {wg_count}")
                    if ovpn_count:
                        print(f"         OpenVPN Connections: {ovpn_count}")
                
                if device.routes and len(device.routes) > 0:
                    print(f"      📋 Routing Table ({len(device.routes)} routes):")
                    for route in device.routes[:3]:  # Show first 3
                        dest = route.get('destination', 'unknown')
                        gw = route.get('gateway', 'direct')
                        print(f"         → {dest} via {gw}")
                    if len(device.routes) > 3:
                        print(f"         ... and {len(device.routes) - 3} more")
        
        print("\n" + "="*80)
    
    def _get_device_icon(self, device_type):
        """Get emoji icon for device type"""
        icons = {
            'router': '🔀',
            'firewall': '🛡️',
            'switch': '🔌',
            'server': '🖥️',
            'linux_server': '🐧',
            'windows_client': '💻',
            'client': '💻'
        }
        return icons.get(device_type, '❓')


def main():
    parser = argparse.ArgumentParser(
        description='Integrated Network Scanner with pfSense Support'
    )
    parser.add_argument('-c', '--config', default='config.json',
                       help='Configuration file (default: config.json)')
    parser.add_argument('-o', '--output', default='network_scan.json',
                       help='Output JSON file (default: network_scan.json)')
    parser.add_argument('-v', '--verbose', action='store_true',
                       help='Verbose output')
    parser.add_argument('--generate-svg', action='store_true',
                       help='Automatically generate SVG diagram after scan')
    
    args = parser.parse_args()
    
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    # Load configuration
    try:
        with open(args.config, 'r') as f:
            config = json.load(f)
            logger.info(f"Loaded configuration from {args.config}")
    except FileNotFoundError:
        logger.warning(f"Config file {args.config} not found, using defaults")
        config = {}
    
    # Run integrated scanner
    scanner = IntegratedNetworkScanner(config)
    scanner.scan_all()
    
    # Print summary
    scanner.print_summary()
    
    # Export results
    scanner.export_json(args.output)
    
    print(f"\n✓ Scan complete! Results saved to {args.output}")
    
    # Generate SVG if requested
    if args.generate_svg:
        try:
            from svg_generator import NetworkDiagramGenerator
            
            with open(args.output, 'r') as f:
                scan_data = json.load(f)
            
            svg_file = args.output.replace('.json', '.svg')
            generator = NetworkDiagramGenerator(scan_data)
            generator.generate_svg(svg_file)
            
            print(f"✓ SVG diagram generated: {svg_file}")
        except Exception as e:
            logger.error(f"Error generating SVG: {e}")


if __name__ == '__main__':
    main()