#!/usr/bin/env python3
"""
Network Scanner - Comprehensive network topology discovery tool
Scans local and VPN networks, gathers device info, routing tables, and generates
structured data for network diagram visualization.
"""

import subprocess
import ipaddress
import json
import re
import socket
import argparse
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass, asdict
from concurrent.futures import ThreadPoolExecutor, as_completed
import logging

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


@dataclass
class Device:
    """Represents a network device"""
    ip: str
    hostname: Optional[str] = None
    mac: Optional[str] = None
    manufacturer: Optional[str] = None
    os_type: Optional[str] = None
    os_version: Optional[str] = None
    device_type: Optional[str] = None  # router, switch, server, client, etc.
    open_ports: List[int] = None
    ssh_accessible: bool = False
    services: List[str] = None
    routes: List[Dict] = None
    interfaces: List[Dict] = None
    
    def __post_init__(self):
        if self.open_ports is None:
            self.open_ports = []
        if self.services is None:
            self.services = []
        if self.routes is None:
            self.routes = []
        if self.interfaces is None:
            self.interfaces = []


@dataclass
class NetworkSegment:
    """Represents a network segment"""
    name: str
    cidr: str
    gateway: Optional[str] = None
    vlan: Optional[int] = None
    is_vpn: bool = False
    devices: List[Device] = None
    
    def __post_init__(self):
        if self.devices is None:
            self.devices = []


class NetworkScanner:
    """Main network scanner class"""
    
    def __init__(self, config: Dict):
        self.config = config
        self.segments: List[NetworkSegment] = []
        self.ssh_user = config.get('ssh_user', 'root')
        self.ssh_key = config.get('ssh_key_path')
        self.timeout = config.get('timeout', 2)
        
    def discover_networks(self) -> List[str]:
        """Discover all network segments from local routing table"""
        logger.info("Discovering network segments...")
        networks = []
        
        try:
            # Get local routing table
            result = subprocess.run(
                ['ip', 'route', 'show'],
                capture_output=True,
                text=True,
                timeout=5
            )
            
            for line in result.stdout.splitlines():
                # Parse routes like: 192.168.1.0/24 dev eth0 proto kernel scope link src 192.168.1.10
                match = re.search(r'(\d+\.\d+\.\d+\.\d+/\d+)', line)
                if match:
                    network = match.group(1)
                    try:
                        ipaddress.ip_network(network, strict=False)
                        if network not in networks:
                            networks.append(network)
                            logger.info(f"Found network: {network}")
                    except ValueError:
                        continue
            
            # Add configured networks
            for net in self.config.get('additional_networks', []):
                if net not in networks:
                    networks.append(net)
                    logger.info(f"Added configured network: {net}")
                    
        except Exception as e:
            logger.error(f"Error discovering networks: {e}")
        
        return networks
    
    def ping_sweep(self, network: str) -> List[str]:
        """Perform ping sweep to find live hosts"""
        logger.info(f"Performing ping sweep on {network}")
        live_hosts = []
        
        try:
            net = ipaddress.ip_network(network, strict=False)
            hosts = list(net.hosts())
            
            # Limit to reasonable size
            if len(hosts) > 254:
                logger.warning(f"Large network {network}, limiting scan")
                hosts = hosts[:254]
            
            with ThreadPoolExecutor(max_workers=50) as executor:
                future_to_ip = {
                    executor.submit(self._ping_host, str(ip)): str(ip)
                    for ip in hosts
                }
                
                for future in as_completed(future_to_ip):
                    ip = future_to_ip[future]
                    try:
                        if future.result():
                            live_hosts.append(ip)
                            logger.info(f"Host alive: {ip}")
                    except Exception as e:
                        logger.debug(f"Error checking {ip}: {e}")
        
        except Exception as e:
            logger.error(f"Error in ping sweep: {e}")
        
        return live_hosts
    
    def _ping_host(self, ip: str) -> bool:
        """Ping a single host"""
        try:
            result = subprocess.run(
                ['ping', '-c', '1', '-W', str(self.timeout), ip],
                capture_output=True,
                timeout=self.timeout + 1
            )
            return result.returncode == 0
        except:
            return False
    
    def get_device_info(self, ip: str) -> Device:
        """Gather comprehensive information about a device"""
        logger.info(f"Gathering info for {ip}")
        device = Device(ip=ip)
        
        # Get hostname
        device.hostname = self._get_hostname(ip)
        
        # Get MAC address
        device.mac = self._get_mac_address(ip)
        
        # Port scan for common services
        device.open_ports = self._scan_common_ports(ip)
        
        # Try SSH access
        if 22 in device.open_ports:
            device.ssh_accessible = self._test_ssh(ip)
            
            if device.ssh_accessible:
                # Gather detailed info via SSH
                self._gather_ssh_info(device)
        
        # Identify device type based on ports and services
        device.device_type = self._identify_device_type(device)
        
        return device
    
    def _get_hostname(self, ip: str) -> Optional[str]:
        """Get hostname via reverse DNS"""
        try:
            hostname, _, _ = socket.gethostbyaddr(ip)
            return hostname
        except:
            return None
    
    def _get_mac_address(self, ip: str) -> Optional[str]:
        """Get MAC address from ARP table"""
        try:
            # First ensure the host is in ARP table
            subprocess.run(['ping', '-c', '1', '-W', '1', ip], 
                         capture_output=True, timeout=2)
            
            result = subprocess.run(
                ['ip', 'neigh', 'show', ip],
                capture_output=True,
                text=True,
                timeout=2
            )
            
            match = re.search(r'([0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2}:[0-9a-f]{2})', 
                            result.stdout, re.IGNORECASE)
            if match:
                return match.group(1).upper()
        except:
            pass
        return None
    
    def _scan_common_ports(self, ip: str) -> List[int]:
        """Scan common ports"""
        common_ports = [22, 80, 443, 8080, 8443, 3389, 445, 139, 21, 23, 25, 53, 3306, 5432]
        open_ports = []
        
        for port in common_ports:
            if self._check_port(ip, port):
                open_ports.append(port)
                logger.debug(f"{ip}:{port} - OPEN")
        
        return open_ports
    
    def _check_port(self, ip: str, port: int) -> bool:
        """Check if a port is open"""
        try:
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(1)
            result = sock.connect_ex((ip, port))
            sock.close()
            return result == 0
        except:
            return False
    
    def _test_ssh(self, ip: str) -> bool:
        """Test SSH connectivity"""
        try:
            cmd = ['ssh', '-o', 'ConnectTimeout=3', 
                   '-o', 'StrictHostKeyChecking=no',
                   '-o', 'BatchMode=yes']
            
            if self.ssh_key:
                cmd.extend(['-i', self.ssh_key])
            
            cmd.extend([f'{self.ssh_user}@{ip}', 'echo', 'test'])
            
            result = subprocess.run(cmd, capture_output=True, timeout=5)
            return result.returncode == 0
        except:
            return False
    
    def _gather_ssh_info(self, device: Device):
        """Gather detailed information via SSH"""
        logger.info(f"Gathering SSH info from {device.ip}")
        
        # Get OS info
        device.os_type, device.os_version = self._get_os_info(device.ip)
        
        # Get network interfaces
        device.interfaces = self._get_interfaces(device.ip)
        
        # Get routing table
        device.routes = self._get_routes(device.ip)
        
        # Get running services
        device.services = self._get_services(device.ip)
    
    def _ssh_exec(self, ip: str, command: str) -> Optional[str]:
        """Execute command via SSH"""
        try:
            cmd = ['ssh', '-o', 'ConnectTimeout=3',
                   '-o', 'StrictHostKeyChecking=no',
                   '-o', 'BatchMode=yes']
            
            if self.ssh_key:
                cmd.extend(['-i', self.ssh_key])
            
            cmd.extend([f'{self.ssh_user}@{ip}', command])
            
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
            if result.returncode == 0:
                return result.stdout
        except Exception as e:
            logger.debug(f"SSH exec error on {ip}: {e}")
        return None
    
    def _get_os_info(self, ip: str) -> Tuple[Optional[str], Optional[str]]:
        """Get OS type and version"""
        output = self._ssh_exec(ip, 'cat /etc/os-release 2>/dev/null || uname -a')
        if output:
            if 'NAME=' in output:
                match = re.search(r'NAME="?([^"\n]+)"?', output)
                name = match.group(1) if match else None
                match = re.search(r'VERSION="?([^"\n]+)"?', output)
                version = match.group(1) if match else None
                return name, version
            else:
                return 'Unix-like', output.strip()
        return None, None
    
    def _get_interfaces(self, ip: str) -> List[Dict]:
        """Get network interfaces"""
        interfaces = []
        output = self._ssh_exec(ip, 'ip -j addr show 2>/dev/null || ip addr show')
        
        if output:
            try:
                # Try JSON format first
                data = json.loads(output)
                for iface in data:
                    interfaces.append({
                        'name': iface.get('ifname'),
                        'state': iface.get('operstate'),
                        'mac': iface.get('address'),
                        'addresses': [addr.get('local') for addr in iface.get('addr_info', [])]
                    })
            except json.JSONDecodeError:
                # Parse text format
                current_iface = None
                for line in output.splitlines():
                    if not line.startswith(' '):
                        match = re.match(r'\d+: (\S+):', line)
                        if match:
                            current_iface = {'name': match.group(1), 'addresses': []}
                            interfaces.append(current_iface)
                    elif 'inet ' in line and current_iface:
                        match = re.search(r'inet (\S+)', line)
                        if match:
                            current_iface['addresses'].append(match.group(1))
        
        return interfaces
    
    def _get_routes(self, ip: str) -> List[Dict]:
        """Get routing table"""
        routes = []
        output = self._ssh_exec(ip, 'ip route show')
        
        if output:
            for line in output.splitlines():
                route = {'raw': line}
                
                # Parse destination
                if line.startswith('default'):
                    route['destination'] = 'default'
                    match = re.search(r'via (\S+)', line)
                    if match:
                        route['gateway'] = match.group(1)
                else:
                    match = re.match(r'(\S+)', line)
                    if match:
                        route['destination'] = match.group(1)
                
                # Parse interface
                match = re.search(r'dev (\S+)', line)
                if match:
                    route['interface'] = match.group(1)
                
                routes.append(route)
        
        return routes
    
    def _get_services(self, ip: str) -> List[str]:
        """Get running services"""
        services = []
        output = self._ssh_exec(ip, 'systemctl list-units --type=service --state=running --no-pager --no-legend 2>/dev/null')
        
        if output:
            for line in output.splitlines():
                match = re.match(r'(\S+\.service)', line)
                if match:
                    services.append(match.group(1))
        
        return services[:20]  # Limit to top 20
    
    def _identify_device_type(self, device: Device) -> str:
        """Identify device type based on available info"""
        if device.routes and len(device.routes) > 5:
            return 'router'
        elif 80 in device.open_ports or 443 in device.open_ports:
            if device.hostname and 'pfsense' in device.hostname.lower():
                return 'firewall'
            return 'server'
        elif 3389 in device.open_ports:
            return 'windows_client'
        elif 22 in device.open_ports:
            return 'linux_server'
        else:
            return 'client'
    
    def scan_network(self, network: str, name: str = None, is_vpn: bool = False) -> NetworkSegment:
        """Scan a complete network segment"""
        logger.info(f"Scanning network segment: {network}")
        
        segment = NetworkSegment(
            name=name or network,
            cidr=network,
            is_vpn=is_vpn
        )
        
        # Find live hosts
        live_hosts = self.ping_sweep(network)
        
        # Gather device info
        with ThreadPoolExecutor(max_workers=10) as executor:
            future_to_ip = {
                executor.submit(self.get_device_info, ip): ip
                for ip in live_hosts
            }
            
            for future in as_completed(future_to_ip):
                try:
                    device = future.result()
                    segment.devices.append(device)
                except Exception as e:
                    logger.error(f"Error getting device info: {e}")
        
        return segment
    
    def scan_all(self):
        """Scan all configured networks"""
        logger.info("Starting comprehensive network scan...")
        
        # Discover networks
        networks = self.discover_networks()
        
        # Scan each network
        for network in networks:
            try:
                segment = self.scan_network(network)
                self.segments.append(segment)
            except Exception as e:
                logger.error(f"Error scanning {network}: {e}")
        
        logger.info(f"Scan complete. Found {len(self.segments)} segments")
    
    def export_json(self, filename: str):
        """Export results to JSON"""
        data = {
            'scan_timestamp': None,  # Will be set by caller
            'segments': [
                {
                    'name': seg.name,
                    'cidr': seg.cidr,
                    'gateway': seg.gateway,
                    'is_vpn': seg.is_vpn,
                    'devices': [asdict(dev) for dev in seg.devices]
                }
                for seg in self.segments
            ]
        }
        
        with open(filename, 'w') as f:
            json.dump(data, f, indent=2)
        
        logger.info(f"Exported results to {filename}")
    
    def print_summary(self):
        """Print a human-readable summary"""
        print("\n" + "="*80)
        print("NETWORK SCAN SUMMARY")
        print("="*80)
        
        for segment in self.segments:
            print(f"\n📡 Network: {segment.name} ({segment.cidr})")
            print(f"   Devices found: {len(segment.devices)}")
            
            for device in segment.devices:
                print(f"\n   🖥️  {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"      Open Ports: {', '.join(map(str, device.open_ports))}")
                if device.ssh_accessible:
                    print(f"      ✓ SSH Accessible")
                    if device.routes:
                        print(f"      Routes: {len(device.routes)}")
                    if device.interfaces:
                        print(f"      Interfaces: {len(device.interfaces)}")


def main():
    parser = argparse.ArgumentParser(description='Network Scanner for Diagram Generation')
    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')
    
    args = parser.parse_args()
    
    if args.verbose:
        logger.setLevel(logging.DEBUG)
    
    # Load configuration
    try:
        with open(args.config, 'r') as f:
            config = json.load(f)
    except FileNotFoundError:
        logger.warning(f"Config file {args.config} not found, using defaults")
        config = {}
    
    # Run scanner
    scanner = NetworkScanner(config)
    scanner.scan_all()
    
    # Print summary
    scanner.print_summary()
    
    # Export results
    from datetime import datetime
    scanner.export_json(args.output)
    
    print(f"\n✓ Scan complete! Results saved to {args.output}")


if __name__ == '__main__':
    main()