#!/usr/bin/env python3
"""
Comprehensive Network Diagram Generator
Combines network scanning with pfSense XML parsing for complete network topology
"""

import json
import argparse
import logging
from pathlib import Path
from typing import Dict, List, Any
from network_scanner import NetworkScanner
from pfsense_xml_parser import PfSenseXMLParser
from svg_generator import NetworkDiagramGenerator

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


class ComprehensiveNetworkMapper:
    """Complete network mapping solution combining multiple data sources"""

    def __init__(self, config: Dict):
        self.config = config
        self.network_data = {}
        self.pfsense_data = {}
        self.combined_data = {
            'scan_timestamp': None,
            'data_sources': [],
            'segments': [],
            'pfsense_firewalls': [],
            'wireguard_networks': [],
            'openvpn_networks': [],
            'routing_table': [],
            'dhcp_leases': [],
            'static_mappings': []
        }

    def load_network_scan(self, scan_file: str):
        """Load network scan data"""
        logger.info(f"Loading network scan data from {scan_file}")
        try:
            with open(scan_file, 'r') as f:
                self.network_data = json.load(f)
            self.combined_data['data_sources'].append('network_scan')
            logger.info("Network scan data loaded successfully")
        except Exception as e:
            logger.error(f"Error loading network scan: {e}")

    def load_pfsense_configs(self, xml_files: List[str]):
        """Load and parse pfSense XML configuration files"""
        for xml_file in xml_files:
            if Path(xml_file).exists():
                logger.info(f"Parsing pfSense config: {xml_file}")
                parser = PfSenseXMLParser(xml_file)
                data = parser.parse_all()

                if data:
                    hostname = data.get('hostname', f'pfsense_{len(self.pfsense_data)}')
                    self.pfsense_data[hostname] = data
                    self.combined_data['data_sources'].append(f'pfsense_{hostname}')
                    self.combined_data['pfsense_firewalls'].append(data)
                    logger.info(f"pfSense config {hostname} loaded successfully")
                else:
                    logger.error(f"Failed to parse pfSense config: {xml_file}")
            else:
                logger.warning(f"pfSense config file not found: {xml_file}")

    def merge_network_data(self):
        """Merge all data sources into comprehensive network map"""
        logger.info("Merging network data sources...")

        # Start with network scan data if available
        if self.network_data:
            self.combined_data.update({
                'scan_timestamp': self.network_data.get('scan_timestamp'),
                'segments': self.network_data.get('segments', [])
            })

        # Process pfSense data
        self._process_pfsense_data()

        # Extract additional network information
        self._extract_network_topology()

        logger.info("Network data merging complete")

    def _process_pfsense_data(self):
        """Process and integrate pfSense configuration data"""
        for hostname, pfsense in self.pfsense_data.items():
            logger.info(f"Processing pfSense data for {hostname}")

            # Add interfaces as network segments
            interfaces = pfsense.get('interfaces', {})
            for iface_name, iface_data in interfaces.items():
                if iface_data.get('ipaddr') and iface_data.get('ipaddr') != 'dhcp':
                    try:
                        # Create network segment from interface
                        ip = iface_data['ipaddr']
                        subnet = iface_data.get('subnet', '24')
                        network_cidr = f"{ip.rsplit('.', 1)[0]}.0/{subnet}"

                        segment = {
                            'name': f"{hostname}_{iface_name}",
                            'cidr': network_cidr,
                            'gateway': iface_data.get('gateway', ip),
                            'interface': iface_data.get('interface'),
                            'pfsense_host': hostname,
                            'is_vpn': iface_data.get('type') in ['wireguard', 'openvpn'],
                            'devices': []
                        }

                        # Add gateway device
                        gateway_device = {
                            'ip': ip,
                            'hostname': hostname,
                            'device_type': 'firewall',
                            'os_type': 'pfSense',
                            'interface': iface_name,
                            'pfsense_config': True
                        }
                        segment['devices'].append(gateway_device)

                        self.combined_data['segments'].append(segment)

                        # Track VPN networks
                        if iface_data.get('type') == 'wireguard':
                            self.combined_data['wireguard_networks'].append({
                                'network': network_cidr,
                                'interface': iface_data.get('interface'),
                                'pfsense': hostname,
                                'peers': pfsense.get('wireguard', {}).get('peers', [])
                            })
                        elif iface_data.get('type') == 'openvpn':
                            self.combined_data['openvpn_networks'].append({
                                'network': network_cidr,
                                'interface': iface_data.get('interface'),
                                'pfsense': hostname
                            })

                    except Exception as e:
                        logger.debug(f"Error processing interface {iface_name}: {e}")

            # Add static routes
            static_routes = pfsense.get('static_routes', [])
            for route in static_routes:
                route_info = {
                    'network': route.get('network'),
                    'gateway': route.get('gateway'),
                    'description': route.get('description'),
                    'pfsense': hostname
                }
                self.combined_data['routing_table'].append(route_info)

            # Add DHCP information
            dhcp_config = pfsense.get('dhcp', {})
            for iface_name, dhcp_data in dhcp_config.items():
                if dhcp_data.get('static_mappings'):
                    for mapping in dhcp_data['static_mappings']:
                        mapping_info = {
                            'ip': mapping.get('ipaddr'),
                            'mac': mapping.get('mac'),
                            'hostname': mapping.get('hostname'),
                            'description': mapping.get('description'),
                            'interface': iface_name,
                            'pfsense': hostname
                        }
                        self.combined_data['static_mappings'].append(mapping_info)

    def _extract_network_topology(self):
        """Extract network topology information from all sources"""
        logger.info("Extracting network topology...")

        # Process WireGuard peer information
        for wg_net in self.combined_data['wireguard_networks']:
            for peer in wg_net.get('peers', []):
                for allowed_ip in peer.get('allowed_ips', []):
                    # Create segment for remote networks
                    if allowed_ip.get('mask') and allowed_ip['mask'] != '32':
                        remote_segment = {
                            'name': f"WG_Remote_{allowed_ip['address']}_{allowed_ip['mask']}",
                            'cidr': f"{allowed_ip['address']}/{allowed_ip['mask']}",
                            'gateway': wg_net.get('network', '').split('/')[0].rsplit('.', 1)[0] + '.1',
                            'is_vpn': True,
                            'vpn_type': 'wireguard',
                            'pfsense_host': wg_net.get('pfsense'),
                            'devices': []
                        }
                        self.combined_data['segments'].append(remote_segment)

    def generate_svg_diagram(self, output_file: str):
        """Generate SVG network diagram"""
        logger.info(f"Generating SVG diagram: {output_file}")

        # Use the existing SVG generator
        generator = NetworkDiagramGenerator(self.combined_data)
        generator.generate_svg(output_file)

        logger.info(f"SVG diagram generated: {output_file}")

    def export_comprehensive_json(self, output_file: str):
        """Export comprehensive network data"""
        logger.info(f"Exporting comprehensive data to {output_file}")

        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(self.combined_data, f, indent=2, ensure_ascii=False)

        logger.info(f"Comprehensive network data exported to {output_file}")

    def print_summary(self):
        """Print comprehensive network summary"""
        print("\n" + "="*80)
        print("COMPREHENSIVE NETWORK MAPPING SUMMARY")
        print("="*80)

        print(f"\nData Sources: {', '.join(self.combined_data['data_sources'])}")
        print(f"Network Segments: {len(self.combined_data['segments'])}")
        print(f"pfSense Firewalls: {len(self.combined_data['pfsense_firewalls'])}")
        print(f"WireGuard Networks: {len(self.combined_data['wireguard_networks'])}")
        print(f"OpenVPN Networks: {len(self.combined_data['openvpn_networks'])}")
        print(f"Static Routes: {len(self.combined_data['routing_table'])}")
        print(f"DHCP Static Mappings: {len(self.combined_data['static_mappings'])}")

        # Network segments summary
        print(f"\n{'='*80}")
        print("NETWORK SEGMENTS")
        print(f"{'='*80}")

        for segment in self.combined_data['segments']:
            vpn_indicator = " (VPN)" if segment.get('is_vpn') else ""
            pfsense_indicator = f" [{segment.get('pfsense_host')}]" if segment.get('pfsense_host') else ""
            print(f"\n📡 {segment['name']}{vpn_indicator}{pfsense_indicator}")
            print(f"   CIDR: {segment['cidr']}")
            print(f"   Gateway: {segment.get('gateway', 'N/A')}")
            print(f"   Devices: {len(segment.get('devices', []))}")

            # Show devices
            for device in segment.get('devices', [])[:5]:  # Show first 5
                print(f"   • {device.get('ip', 'N/A')} - {device.get('hostname', 'N/A')} ({device.get('device_type', 'N/A')})")

            if len(segment.get('devices', [])) > 5:
                print(f"   ... and {len(segment.get('devices', [])) - 5} more devices")

        # pfSense summary
        if self.combined_data['pfsense_firewalls']:
            print(f"\n{'='*80}")
            print("PFSENSE FIREWALLS")
            print(f"{'='*80}")

            for pfsense in self.combined_data['pfsense_firewalls']:
                hostname = pfsense.get('hostname', 'Unknown')
                print(f"\n🛡️  {hostname}")
                print(f"   Interfaces: {len(pfsense.get('interfaces', {}))}")
                print(f"   Static Routes: {len(pfsense.get('static_routes', []))}")

                wg_config = pfsense.get('wireguard', {})
                if wg_config.get('enabled'):
                    print(f"   WireGuard: {len(wg_config.get('tunnels', []))} tunnels, {len(wg_config.get('peers', []))} peers")

                dhcp_config = pfsense.get('dhcp', {})
                total_mappings = sum(len(dhcp.get('static_mappings', [])) for dhcp in dhcp_config.values())
                print(f"   DHCP Static Mappings: {total_mappings}")

        # VPN Networks
        if self.combined_data['wireguard_networks'] or self.combined_data['openvpn_networks']:
            print(f"\n{'='*80}")
            print("VPN NETWORKS")
            print(f"{'='*80}")

            for wg_net in self.combined_data['wireguard_networks']:
                print(f"\n🔐 WireGuard: {wg_net['network']} via {wg_net['pfsense']}")
                print(f"   Interface: {wg_net['interface']}")
                print(f"   Peers: {len(wg_net.get('peers', []))}")

            for ovpn_net in self.combined_data['openvpn_networks']:
                print(f"\n🔐 OpenVPN: {ovpn_net['network']} via {ovpn_net['pfsense']}")
                print(f"   Interface: {ovpn_net['interface']}")

        print(f"\n{'='*80}")
        print("NETWORK TOPOLOGY ANALYSIS COMPLETE")
        print(f"{'='*80}")


def main():
    """Command line interface"""
    parser = argparse.ArgumentParser(
        description='Comprehensive Network Diagram Generator with pfSense Integration'
    )
    parser.add_argument('-c', '--config', default='config.json',
                       help='Network scanner configuration file')
    parser.add_argument('-s', '--scan-data', help='Existing network scan JSON file')
    parser.add_argument('-p', '--pfsense-xml', nargs='+',
                       help='pfSense XML configuration files')
    parser.add_argument('-o', '--output', default='comprehensive_network.json',
                       help='Output comprehensive JSON file')
    parser.add_argument('--svg', help='Generate SVG diagram')
    parser.add_argument('-v', '--verbose', action='store_true',
                       help='Verbose output')

    args = parser.parse_args()

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Load configuration
    config = {}
    if Path(args.config).exists():
        try:
            with open(args.config, 'r') as f:
                config = json.load(f)
        except Exception as e:
            logger.warning(f"Could not load config {args.config}: {e}")

    # Initialize comprehensive mapper
    mapper = ComprehensiveNetworkMapper(config)

    # Load network scan data if provided
    if args.scan_data and Path(args.scan_data).exists():
        mapper.load_network_scan(args.scan_data)
    else:
        logger.info("No network scan data provided - will use pfSense data only")

    # Load pfSense configurations
    pfsense_files = args.pfsense_xml or []
    # Auto-detect pfSense XML files if none specified
    if not pfsense_files:
        xml_files = list(Path('.').glob('config-*.xml'))
        if xml_files:
            pfsense_files = [str(f) for f in xml_files]
            logger.info(f"Auto-detected pfSense XML files: {pfsense_files}")

    if pfsense_files:
        mapper.load_pfsense_configs(pfsense_files)
    else:
        logger.warning("No pfSense XML files found")

    # Merge all data
    mapper.merge_network_data()

    # Print summary
    mapper.print_summary()

    # Export comprehensive data
    mapper.export_comprehensive_json(args.output)
    print(f"\n✅ Comprehensive network data saved to: {args.output}")

    # Generate SVG if requested
    if args.svg:
        svg_file = args.svg
        mapper.generate_svg_diagram(svg_file)
        print(f"✅ SVG network diagram generated: {svg_file}")


if __name__ == '__main__':
    main()