pic/api/wireguard_manager.py

#!/usr/bin/env python3
"""
WireGuard Manager for Personal Internet Cell
Handles WireGuard VPN configuration and peer management
"""

import os
import json
import subprocess
import logging
from datetime import datetime
from typing import Dict, List, Optional, Any
from base_service_manager import BaseServiceManager

logger = logging.getLogger(__name__)

class WireGuardManager(BaseServiceManager):
    """Manages WireGuard VPN configuration and peers"""
    
    def __init__(self, data_dir: str = '/app/data', config_dir: str = '/app/config'):
        super().__init__('wireguard', data_dir, config_dir)
        self.wg_config_dir = os.path.join(config_dir, 'wireguard')
        self.peers_dir = os.path.join(data_dir, 'wireguard', 'peers')
        
        # Ensure directories exist
        os.makedirs(self.wg_config_dir, exist_ok=True)
        os.makedirs(self.peers_dir, exist_ok=True)

    def get_status(self) -> Dict[str, Any]:
        """Get WireGuard service status"""
        try:
            # Check if we're running in Docker environment
            import os
            is_docker = os.path.exists('/.dockerenv') or os.environ.get('DOCKER_CONTAINER') == 'true'
            
            if is_docker:
                # Check if WireGuard container is actually running
                container_running = self._check_wireguard_container_status()
                status = {
                    'running': container_running,
                    'status': 'online' if container_running else 'offline',
                    'interface': 'wg0' if container_running else 'unknown',
                    'peers_count': len(self._get_configured_peers()) if container_running else 0,
                    'total_traffic': self._get_traffic_stats() if container_running else {'bytes_sent': 0, 'bytes_received': 0},
                    'timestamp': datetime.utcnow().isoformat()
                }
            else:
                # Check actual service status in production
                status = {
                    'running': self._check_wireguard_status(),
                    'status': 'online' if self._check_wireguard_status() else 'offline',
                    'interface': 'wg0',
                    'peers_count': len(self._get_configured_peers()),
                    'total_traffic': self._get_traffic_stats(),
                    'timestamp': datetime.utcnow().isoformat()
                }
            
            return status
        except Exception as e:
            return self.handle_error(e, "get_status")

    def test_connectivity(self) -> Dict[str, Any]:
        """Test WireGuard connectivity"""
        try:
            # Test if WireGuard interface exists and is up
            interface_up = self._check_interface_status()
            
            # Test if peers can connect
            peers_connectivity = self._test_peers_connectivity()
            
            results = {
                'interface_up': interface_up,
                'peers_connectivity': peers_connectivity,
                'success': interface_up and all(peers_connectivity.values()),
                'timestamp': datetime.utcnow().isoformat()
            }
            
            return results
        except Exception as e:
            return self.handle_error(e, "test_connectivity")

    def _check_wireguard_status(self) -> bool:
        """Check if WireGuard service is running"""
        try:
            # Check if wg0 interface exists
            result = subprocess.run(['ip', 'link', 'show', 'wg0'], 
                                  capture_output=True, text=True, timeout=5)
            return result.returncode == 0
        except Exception:
            return False

    def _check_wireguard_container_status(self) -> bool:
        """Check if WireGuard Docker container is running"""
        try:
            import docker
            client = docker.from_env()
            containers = client.containers.list(filters={'name': 'cell-wireguard'})
            return len(containers) > 0
        except Exception:
            return False

    def _check_interface_status(self) -> bool:
        """Check if WireGuard interface is up"""
        try:
            result = subprocess.run(['ip', 'link', 'show', 'wg0'], 
                                  capture_output=True, text=True, timeout=5)
            if result.returncode == 0:
                return 'UP' in result.stdout
            return False
        except Exception:
            return False

    def _get_configured_peers(self) -> List[Dict[str, Any]]:
        """Get list of configured peers"""
        peers = []
        try:
            # Read peer configurations from peers directory
            for filename in os.listdir(self.peers_dir):
                if filename.endswith('.conf'):
                    peer_name = filename[:-5]  # Remove .conf extension
                    peer_file = os.path.join(self.peers_dir, filename)
                    
                    with open(peer_file, 'r') as f:
                        content = f.read()
                    
                    # Parse peer configuration
                    peer_config = self._parse_peer_config(content)
                    peer_config['name'] = peer_name
                    peers.append(peer_config)
        except Exception as e:
            logger.error(f"Error reading peer configurations: {e}")
        
        return peers

    def _parse_peer_config(self, content: str) -> Dict[str, Any]:
        """Parse WireGuard peer configuration"""
        config = {}
        lines = content.strip().split('\n')
        
        for line in lines:
            line = line.strip()
            if line.startswith('[Peer]'):
                continue
            elif '=' in line:
                key, value = line.split('=', 1)
                config[key.strip()] = value.strip()
        
        return config

    def _get_traffic_stats(self) -> Dict[str, int]:
        """Get WireGuard traffic statistics"""
        try:
            result = subprocess.run(['wg', 'show', 'wg0', 'transfer'], 
                                  capture_output=True, text=True, timeout=5)
            
            if result.returncode == 0:
                lines = result.stdout.strip().split('\n')
                total_rx = 0
                total_tx = 0
                
                for line in lines:
                    if line.strip():
                        parts = line.split()
                        if len(parts) >= 3:
                            try:
                                rx = int(parts[1])
                                tx = int(parts[2])
                                total_rx += rx
                                total_tx += tx
                            except ValueError:
                                continue
                
                return {
                    'bytes_received': total_rx,
                    'bytes_sent': total_tx
                }
        except Exception as e:
            logger.error(f"Error getting traffic stats: {e}")
        
        return {'bytes_received': 0, 'bytes_sent': 0}

    def _test_peers_connectivity(self) -> Dict[str, bool]:
        """Test connectivity to all peers"""
        connectivity = {}
        peers = self._get_configured_peers()
        
        for peer in peers:
            peer_name = peer.get('name', 'unknown')
            allowed_ips = peer.get('AllowedIPs', '')
            
            if allowed_ips:
                # Extract first IP from AllowedIPs
                ip = allowed_ips.split(',')[0].split('/')[0]
                
                try:
                    # Ping the peer IP
                    result = subprocess.run(['ping', '-c', '1', '-W', '2', ip], 
                                          capture_output=True, text=True, timeout=5)
                    connectivity[peer_name] = result.returncode == 0
                except Exception:
                    connectivity[peer_name] = False
            else:
                connectivity[peer_name] = False
        
        return connectivity

    def get_wireguard_status(self) -> Dict[str, Any]:
        """Get detailed WireGuard status"""
        try:
            status = self.get_status()
            
            # Get peer details
            peers = self._get_configured_peers()
            peer_details = []
            
            for peer in peers:
                peer_detail = {
                    'name': peer.get('name', 'unknown'),
                    'public_key': peer.get('PublicKey', ''),
                    'allowed_ips': peer.get('AllowedIPs', ''),
                    'endpoint': peer.get('Endpoint', ''),
                    'last_handshake': peer.get('LastHandshake', ''),
                    'transfer_rx': peer.get('TransferRx', 0),
                    'transfer_tx': peer.get('TransferTx', 0)
                }
                peer_details.append(peer_detail)
            
            status['peers'] = peer_details
            return status
        except Exception as e:
            return self.handle_error(e, "get_wireguard_status")

    def get_wireguard_peers(self) -> List[Dict[str, Any]]:
        """Get all WireGuard peers"""
        try:
            peers = self._get_configured_peers()
            return peers
        except Exception as e:
            logger.error(f"Error getting WireGuard peers: {e}")
            return []

    def add_wireguard_peer(self, name: str, public_key: str, allowed_ips: str, 
                          endpoint: str = '', persistent_keepalive: int = 25) -> bool:
        """Add a new WireGuard peer"""
        try:
            # Create peer configuration
            peer_config = f"""[Peer]
PublicKey = {public_key}
AllowedIPs = {allowed_ips}
"""
            
            if endpoint:
                peer_config += f"Endpoint = {endpoint}\n"
            
            if persistent_keepalive:
                peer_config += f"PersistentKeepalive = {persistent_keepalive}\n"
            
            # Save peer configuration
            peer_file = os.path.join(self.peers_dir, f'{name}.conf')
            with open(peer_file, 'w') as f:
                f.write(peer_config)
            
            # Reload WireGuard configuration
            self._reload_wireguard_config()
            
            logger.info(f"Added WireGuard peer: {name}")
            return True
        except Exception as e:
            logger.error(f"Failed to add WireGuard peer {name}: {e}")
            return False

    def remove_wireguard_peer(self, name: str) -> bool:
        """Remove a WireGuard peer"""
        try:
            peer_file = os.path.join(self.peers_dir, f'{name}.conf')
            if os.path.exists(peer_file):
                os.remove(peer_file)
                
                # Reload WireGuard configuration
                self._reload_wireguard_config()
                
                logger.info(f"Removed WireGuard peer: {name}")
                return True
            else:
                logger.warning(f"Peer file not found: {peer_file}")
                return False
        except Exception as e:
            logger.error(f"Failed to remove WireGuard peer {name}: {e}")
            return False

    def generate_peer_keys(self, peer_name: str) -> Dict[str, str]:
        """Generate WireGuard keys for a peer"""
        try:
            # Generate private key
            private_key_result = subprocess.run(['wg', 'genkey'], 
                                              capture_output=True, text=True, timeout=10)
            if private_key_result.returncode != 0:
                raise Exception("Failed to generate private key")
            
            private_key = private_key_result.stdout.strip()
            
            # Generate public key from private key
            public_key_result = subprocess.run(['wg', 'pubkey'], 
                                             input=private_key, 
                                             capture_output=True, text=True, timeout=10)
            if public_key_result.returncode != 0:
                raise Exception("Failed to generate public key")
            
            public_key = public_key_result.stdout.strip()
            
            # Save keys to file
            keys_file = os.path.join(self.peers_dir, f'{peer_name}_keys.json')
            keys_data = {
                'private_key': private_key,
                'public_key': public_key,
                'peer_name': peer_name,
                'generated_at': datetime.utcnow().isoformat()
            }
            
            with open(keys_file, 'w') as f:
                json.dump(keys_data, f, indent=2)
            
            logger.info(f"Generated keys for peer: {peer_name}")
            return {
                'private_key': private_key,
                'public_key': public_key,
                'peer_name': peer_name
            }
        except Exception as e:
            logger.error(f"Failed to generate keys for peer {peer_name}: {e}")
            raise

    def _reload_wireguard_config(self):
        """Reload WireGuard configuration by updating the main config file"""
        try:
            # Read the main server configuration
            server_config_path = os.path.join(self.wg_config_dir, 'wg_confs', 'wg0.conf')
            if not os.path.exists(server_config_path):
                logger.error("Server configuration file not found")
                return False
            
            with open(server_config_path, 'r') as f:
                server_content = f.read()
            
            # Find the end of the [Interface] section
            lines = server_content.split('\n')
            interface_end = 0
            for i, line in enumerate(lines):
                if line.strip().startswith('[Peer]'):
                    interface_end = i
                    break
            else:
                interface_end = len(lines)
            
            # Keep only the [Interface] section
            interface_lines = lines[:interface_end]
            
            # Add all peer configurations
            peer_lines = []
            for filename in os.listdir(self.peers_dir):
                if filename.endswith('.conf') and not filename.endswith('_keys.json'):
                    peer_file = os.path.join(self.peers_dir, filename)
                    with open(peer_file, 'r') as f:
                        peer_content = f.read().strip()
                        if peer_content:
                            peer_lines.append('')  # Empty line before peer
                            peer_lines.extend(peer_content.split('\n'))
            
            # Combine interface and peer configurations
            new_content = '\n'.join(interface_lines + peer_lines)
            
            # Write the updated configuration
            with open(server_config_path, 'w') as f:
                f.write(new_content)
            
            # Restart WireGuard container to apply changes
            import subprocess
            result = subprocess.run(['docker', 'restart', 'cell-wireguard'], 
                                  capture_output=True, text=True, timeout=30)
            if result.returncode == 0:
                logger.info("WireGuard configuration reloaded and container restarted")
                return True
            else:
                logger.error(f"Failed to restart WireGuard container: {result.stderr}")
                return False
                
        except Exception as e:
            logger.error(f"Failed to reload WireGuard configuration: {e}")
            return False

    def get_metrics(self) -> Dict[str, Any]:
        """Get WireGuard metrics"""
        try:
            traffic_stats = self._get_traffic_stats()
            peers = self._get_configured_peers()
            
            return {
                'service': 'wireguard',
                'timestamp': datetime.utcnow().isoformat(),
                'status': 'online' if self._check_wireguard_status() else 'offline',
                'peers_count': len(peers),
                'traffic_stats': traffic_stats,
                'interface_status': self._check_interface_status()
            }
        except Exception as e:
            return self.handle_error(e, "get_metrics")

    def restart_service(self) -> bool:
        """Restart WireGuard service"""
        try:
            # Stop WireGuard interface
            subprocess.run(['wg-quick', 'down', 'wg0'], 
                         capture_output=True, text=True, timeout=10)
            
            # Start WireGuard interface
            subprocess.run(['wg-quick', 'up', 'wg0'], 
                         capture_output=True, text=True, timeout=10)
            
            logger.info("WireGuard service restarted")
            return True
        except Exception as e:
            logger.error(f"Failed to restart WireGuard service: {e}")
            return False

    def get_peer_config(self, peer_name: str) -> Optional[str]:
        """Get WireGuard client configuration for a specific peer"""
        try:
            # Get peer information
            peers = self.get_wireguard_peers()
            peer_info = None
            
            for peer in peers:
                if peer.get('name') == peer_name:
                    peer_info = peer
                    break
            
            if not peer_info:
                logger.warning(f"Peer {peer_name} not found")
                return None
            
            # Get server configuration
            server_config = self._get_server_config()
            
            # Generate client configuration
            client_config = self._generate_client_config(peer_info, server_config)
            
            return client_config
            
        except Exception as e:
            logger.error(f"Error getting peer config for {peer_name}: {e}")
            return None

    def _get_server_config(self) -> Dict[str, str]:
        """Get server configuration details"""
        try:
            # Try to read server config file
            server_config_path = os.path.join(self.wg_config_dir, 'wg_confs', 'wg0.conf')
            if os.path.exists(server_config_path):
                with open(server_config_path, 'r') as f:
                    content = f.read()
                
                # Parse server configuration
                lines = content.strip().split('\n')
                server_public_key = None
                server_endpoint = None
                server_private_key = None
                
                # Look for server private key and endpoint
                for line in lines:
                    line = line.strip()
                    if line.startswith('PrivateKey'):
                        server_private_key = line.split('=', 1)[1].strip()
                    elif line.startswith('ListenPort'):
                        port = line.split('=', 1)[1].strip()
                        # Get server IP from environment or detect it
                        server_ip = os.environ.get('WIREGUARD_SERVER_IP')
                        if not server_ip:
                            # Try to get the actual external IP
                            try:
                                import socket
                                import requests
                                # First try to get external IP from a service
                                try:
                                    response = requests.get('https://api.ipify.org', timeout=5)
                                    if response.status_code == 200:
                                        server_ip = response.text.strip()
                                    else:
                                        raise Exception("Failed to get external IP")
                                except Exception:
                                    # Fallback: try to get local IP that's not Docker internal
                                    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                                        s.connect(("8.8.8.8", 80))
                                        local_ip = s.getsockname()[0]
                                        # If it's a Docker internal IP, use localhost for development
                                        if local_ip.startswith('172.') or local_ip.startswith('192.168.'):
                                            server_ip = "localhost"
                                        else:
                                            server_ip = local_ip
                            except Exception:
                                # Ultimate fallback to localhost for development
                                server_ip = "localhost"
                        server_endpoint = f"{server_ip}:{port}"
                
                # Generate public key from private key if we have it
                if server_private_key:
                    try:
                        # Use wg pubkey command to generate public key from private key
                        import subprocess
                        result = subprocess.run(['wg', 'pubkey'], 
                                              input=server_private_key,
                                              capture_output=True, text=True, timeout=5)
                        if result.returncode == 0:
                            server_public_key = result.stdout.strip()
                        else:
                            # Fallback: try to read from existing public key file
                            pubkey_path = os.path.join(self.wg_config_dir, 'publickey')
                            if os.path.exists(pubkey_path):
                                with open(pubkey_path, 'r') as f:
                                    server_public_key = f.read().strip()
                            else:
                                server_public_key = "SERVER_PUBLIC_KEY_PLACEHOLDER"
                    except Exception as e:
                        logger.warning(f"Could not generate public key: {e}")
                        server_public_key = "SERVER_PUBLIC_KEY_PLACEHOLDER"
                else:
                    server_public_key = "SERVER_PUBLIC_KEY_PLACEHOLDER"
                
                # Set default endpoint if not found
                if not server_endpoint:
                    # Try to get the actual server IP
                    server_ip = os.environ.get('WIREGUARD_SERVER_IP')
                    if not server_ip:
                        # Try to get the actual external IP
                        try:
                            import socket
                            import requests
                            # First try to get external IP from a service
                            try:
                                response = requests.get('https://api.ipify.org', timeout=5)
                                if response.status_code == 200:
                                    server_ip = response.text.strip()
                                else:
                                    raise Exception("Failed to get external IP")
                            except Exception:
                                # Fallback: try to get local IP that's not Docker internal
                                with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                                    s.connect(("8.8.8.8", 80))
                                    local_ip = s.getsockname()[0]
                                    # If it's a Docker internal IP, use localhost for development
                                    if local_ip.startswith('172.') or local_ip.startswith('192.168.'):
                                        server_ip = "localhost"
                                    else:
                                        server_ip = local_ip
                        except Exception:
                            # Ultimate fallback to localhost for development
                            server_ip = "localhost"
                    server_endpoint = f"{server_ip}:51820"
                
                return {
                    'public_key': server_public_key,
                    'endpoint': server_endpoint,
                    'allowed_ips': '0.0.0.0/0'
                }
        except Exception as e:
            logger.error(f"Error reading server config: {e}")
        
        # Return default values
        return {
            'public_key': 'SERVER_PUBLIC_KEY_PLACEHOLDER',
            'endpoint': 'YOUR_SERVER_IP:51820',
            'allowed_ips': '0.0.0.0/0'
        }

    def _generate_client_config(self, peer_info: Dict[str, Any], server_config: Dict[str, str]) -> str:
        """Generate WireGuard client configuration"""
        try:
            # Get peer private key from peer data
            peer_private_key = peer_info.get('private_key', 'YOUR_PRIVATE_KEY_HERE')
            
            # Check if IP already has a subnet mask, if not add /32
            peer_ip = peer_info.get('ip', '10.0.0.2')
            peer_address = peer_ip if '/' in peer_ip else f"{peer_ip}/32"
            
            config = f"""[Interface]
PrivateKey = {peer_private_key}
Address = {peer_address}
DNS = 8.8.8.8, 1.1.1.1

[Peer]
PublicKey = {server_config['public_key']}
Endpoint = {server_config['endpoint']}
AllowedIPs = {server_config['allowed_ips']}
PersistentKeepalive = {peer_info.get('persistent_keepalive', 25)}"""
            
            return config
            
        except Exception as e:
            logger.error(f"Error generating client config: {e}")
            return None

    def get_server_config(self) -> Dict[str, str]:
        """Get server configuration details"""
        try:
            # Try to read server config file
            server_config_path = os.path.join(self.wg_config_dir, 'wg_confs', 'wg0.conf')
            logger.info(f"Looking for server config at: {server_config_path}")
            logger.info(f"wg_config_dir is: {self.wg_config_dir}")
            logger.info(f"File exists: {os.path.exists(server_config_path)}")
            if os.path.exists(server_config_path):
                with open(server_config_path, 'r') as f:
                    content = f.read()
                
                # Parse server configuration
                lines = content.strip().split('\n')
                server_public_key = None
                server_endpoint = None
                server_private_key = None
                
                # Look for server private key and endpoint
                for line in lines:
                    line = line.strip()
                    if line.startswith('PrivateKey'):
                        server_private_key = line.split('=', 1)[1].strip()
                        logger.info(f"Found server private key: {server_private_key[:10]}...")
                    elif line.startswith('ListenPort'):
                        port = line.split('=', 1)[1].strip()
                        logger.info(f"Found listen port: {port}")
                        # Get server IP from environment or detect it
                        server_ip = os.environ.get('WIREGUARD_SERVER_IP')
                        if not server_ip:
                            # Try to get the actual external IP
                            try:
                                import socket
                                import requests
                                # First try to get external IP from a service
                                try:
                                    response = requests.get('https://api.ipify.org', timeout=5)
                                    if response.status_code == 200:
                                        server_ip = response.text.strip()
                                        logger.info(f"Got external IP from service: {server_ip}")
                                    else:
                                        raise Exception("Failed to get external IP")
                                except Exception:
                                    # Fallback: try to get local IP that's not Docker internal
                                    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                                        s.connect(("8.8.8.8", 80))
                                        local_ip = s.getsockname()[0]
                                        # If it's a Docker internal IP, use localhost for development
                                        if local_ip.startswith('172.') or local_ip.startswith('192.168.'):
                                            server_ip = "localhost"
                                            logger.info(f"Using localhost for development (Docker internal IP: {local_ip})")
                                        else:
                                            server_ip = local_ip
                                            logger.info(f"Using local IP: {server_ip}")
                            except Exception:
                                # Ultimate fallback to localhost for development
                                server_ip = "localhost"
                                logger.info("Using localhost as ultimate fallback")
                        server_endpoint = f"{server_ip}:{port}"
                        logger.info(f"Set server endpoint: {server_endpoint}")
                
                # Generate public key from private key if we have it
                if server_private_key:
                    try:
                        logger.info("Generating public key from private key...")
                        # Use wg pubkey command to generate public key from private key
                        import subprocess
                        result = subprocess.run(['wg', 'pubkey'], 
                                              input=server_private_key,
                                              capture_output=True, text=True, timeout=5)
                        if result.returncode == 0:
                            server_public_key = result.stdout.strip()
                            logger.info(f"Generated server public key: {server_public_key[:10]}...")
                        else:
                            # Fallback: try to read from existing public key file
                            pubkey_path = os.path.join(self.wg_config_dir, 'publickey')
                            if os.path.exists(pubkey_path):
                                with open(pubkey_path, 'r') as f:
                                    server_public_key = f.read().strip()
                            else:
                                server_public_key = "SERVER_PUBLIC_KEY_PLACEHOLDER"
                    except Exception as e:
                        logger.warning(f"Could not generate public key: {e}")
                        server_public_key = "SERVER_PUBLIC_KEY_PLACEHOLDER"
                else:
                    server_public_key = "SERVER_PUBLIC_KEY_PLACEHOLDER"
                
                # Set default endpoint if not found
                if not server_endpoint:
                    # Try to get the actual server IP
                    server_ip = os.environ.get('WIREGUARD_SERVER_IP')
                    if not server_ip:
                        # Try to get the host IP from Docker network
                        try:
                            import socket
                            # Connect to a remote address to determine local IP
                            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
                                s.connect(("8.8.8.8", 80))
                                server_ip = s.getsockname()[0]
                        except Exception:
                            # Fallback to localhost
                            server_ip = "localhost"
                    server_endpoint = f"{server_ip}:51820"
                
                return {
                    'public_key': server_public_key,
                    'endpoint': server_endpoint
                }
        except Exception as e:
            logger.error(f"Error reading server config: {e}")
        
        # Return default values
        return {
            'public_key': 'SERVER_PUBLIC_KEY_PLACEHOLDER',
            'endpoint': 'YOUR_SERVER_IP:51820'
        }

    def get_peer_status(self, public_key: str) -> Dict[str, Any]:
        """Get status for a specific peer"""
        try:
            # Get WireGuard interface status
            result = subprocess.run(['wg', 'show'], capture_output=True, text=True, check=True)
            wg_output = result.stdout
            
            # Parse the output to find the specific peer
            lines = wg_output.strip().split('\n')
            peer_info = {}
            in_peer = False
            
            for line in lines:
                if line.startswith('peer:') and public_key in line:
                    in_peer = True
                    peer_info['public_key'] = public_key
                elif line.startswith('peer:') and public_key not in line:
                    in_peer = False
                elif in_peer and line.startswith('  allowed ips:'):
                    peer_info['allowed_ips'] = line.split(':', 1)[1].strip()
                elif in_peer and line.startswith('  latest handshake:'):
                    handshake_str = line.split(':', 1)[1].strip()
                    if handshake_str and handshake_str != '(none)':
                        peer_info['latest_handshake'] = handshake_str
                        peer_info['online'] = True
                    else:
                        peer_info['online'] = False
                elif in_peer and line.startswith('  transfer:'):
                    transfer_str = line.split(':', 1)[1].strip()
                    if transfer_str and transfer_str != '(none)':
                        # Parse transfer data (e.g., "1.2 KiB received, 3.4 KiB sent")
                        parts = transfer_str.split(',')
                        if len(parts) >= 2:
                            rx_part = parts[0].strip()
                            tx_part = parts[1].strip()
                            
                            # Extract numbers from strings like "1.2 KiB received"
                            import re
                            rx_match = re.search(r'([\d.]+)\s+(\w+)', rx_part)
                            tx_match = re.search(r'([\d.]+)\s+(\w+)', tx_part)
                            
                            if rx_match and tx_match:
                                rx_value = float(rx_match.group(1))
                                rx_unit = rx_match.group(2)
                                tx_value = float(tx_match.group(1))
                                tx_unit = tx_match.group(2)
                                
                                # Convert to bytes
                                def convert_to_bytes(value, unit):
                                    multipliers = {'B': 1, 'KiB': 1024, 'MiB': 1024**2, 'GiB': 1024**3}
                                    return int(value * multipliers.get(unit, 1))
                                
                                peer_info['transfer_rx'] = convert_to_bytes(rx_value, rx_unit)
                                peer_info['transfer_tx'] = convert_to_bytes(tx_value, tx_unit)
            
            # Set default values if not found
            if 'online' not in peer_info:
                peer_info['online'] = False
            if 'transfer_rx' not in peer_info:
                peer_info['transfer_rx'] = 0
            if 'transfer_tx' not in peer_info:
                peer_info['transfer_tx'] = 0
            if 'latest_handshake' not in peer_info:
                peer_info['latest_handshake'] = None
            
            return peer_info
        except Exception as e:
            logger.error(f"Failed to get peer status for {public_key}: {e}")
            return {'online': False, 'transfer_rx': 0, 'transfer_tx': 0, 'latest_handshake': None}

    def setup_network_configuration(self) -> bool:
        """Setup network configuration for internet access"""
        try:
            logger.info("Setting up network configuration for internet access...")
            
            # Enable IP forwarding
            self._enable_ip_forwarding()
            
            # Configure NAT and routing
            self._configure_nat_routing()
            
            logger.info("Network configuration completed successfully")
            return True
        except Exception as e:
            logger.error(f"Failed to setup network configuration: {e}")
            return False

    def _enable_ip_forwarding(self):
        """Enable IP forwarding"""
        try:
            # Enable IP forwarding in the container
            subprocess.run(['sh', '-c', 'echo 1 > /proc/sys/net/ipv4/ip_forward'], check=True)
            logger.info("IP forwarding enabled")
        except Exception as e:
            logger.error(f"Failed to enable IP forwarding: {e}")
            raise

    def _configure_nat_routing(self):
        """Configure NAT and routing for internet access"""
        try:
            # Get the main network interface
            result = subprocess.run(['ip', 'route', 'show', 'default'], capture_output=True, text=True, check=True)
            main_interface = result.stdout.split()[4]  # Extract interface name
            
            # Configure iptables rules
            rules = [
                # Allow forwarding for WireGuard interface
                f"iptables -A FORWARD -i wg0 -j ACCEPT",
                f"iptables -A FORWARD -o wg0 -j ACCEPT",
                
                # NAT rule for internet access
                f"iptables -t nat -A POSTROUTING -s 10.0.0.0/24 -o {main_interface} -j MASQUERADE",
                
                # Allow established and related connections
                "iptables -A FORWARD -m conntrack --ctstate ESTABLISHED,RELATED -j ACCEPT"
            ]
            
            for rule in rules:
                try:
                    subprocess.run(['sh', '-c', rule], check=True)
                except subprocess.CalledProcessError as e:
                    logger.warning(f"Rule may already exist: {rule} - {e}")
            
            logger.info(f"NAT and routing configured for interface {main_interface}")
        except Exception as e:
            logger.error(f"Failed to configure NAT routing: {e}")
            raise

    def get_network_status(self) -> Dict[str, Any]:
        """Get network configuration status"""
        try:
            status = {
                'ip_forwarding': self._check_ip_forwarding(),
                'nat_rules': self._check_nat_rules(),
                'forwarding_rules': self._check_forwarding_rules(),
                'interface_status': self._check_interface_status(),
                'timestamp': datetime.utcnow().isoformat()
            }
            return status
        except Exception as e:
            logger.error(f"Failed to get network status: {e}")
            return {'error': str(e)}

    def _check_ip_forwarding(self) -> bool:
        """Check if IP forwarding is enabled"""
        try:
            # Check in WireGuard container
            result = subprocess.run(['docker', 'exec', 'cell-wireguard', 'cat', '/proc/sys/net/ipv4/ip_forward'], capture_output=True, text=True, check=True)
            return result.stdout.strip() == '1'
        except:
            return False

    def _check_nat_rules(self) -> bool:
        """Check if NAT rules are configured"""
        try:
            # Check in WireGuard container
            result = subprocess.run(['docker', 'exec', 'cell-wireguard', 'iptables', '-t', 'nat', '-L', 'POSTROUTING', '-n'], capture_output=True, text=True, check=True)
            return 'MASQUERADE' in result.stdout
        except:
            return False

    def _check_forwarding_rules(self) -> bool:
        """Check if forwarding rules are configured"""
        try:
            # Check in WireGuard container
            result = subprocess.run(['docker', 'exec', 'cell-wireguard', 'iptables', '-L', 'FORWARD', '-n'], capture_output=True, text=True, check=True)
            # Check for ACCEPT rules (which indicate forwarding is allowed)
            return 'ACCEPT' in result.stdout and len(result.stdout.strip().split('\n')) > 2
        except:
            return False

    def _check_interface_status(self) -> bool:
        """Check if WireGuard interface is up"""
        try:
            # Check in WireGuard container
            result = subprocess.run(['docker', 'exec', 'cell-wireguard', 'ip', 'link', 'show', 'wg0'], capture_output=True, text=True, check=True)
            return 'UP' in result.stdout
        except:
            return False