core/openvpn_manager.py

"""
OpenVPN Manager Module

Manages OpenVPN server integration with the Virtual ISP Stack
"""

import os
import json
import subprocess
import threading
import time
import logging
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict
import ipaddress

logger = logging.getLogger(__name__)

@dataclass
class VPNClient:
    """Represents a connected VPN client"""
    client_id: str
    common_name: str
    ip_address: str
    connected_at: float
    bytes_received: int = 0
    bytes_sent: int = 0
    status: str = "connected"
    routed_through_vpn: bool = False

@dataclass
class VPNServerStatus:
    """Represents VPN server status"""
    is_running: bool
    connected_clients: int
    total_bytes_received: int
    total_bytes_sent: int
    uptime: float
    server_ip: str
    server_port: int

class OpenVPNManager:
    """Manages OpenVPN server and client connections with traffic routing"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.server_config_path = "/etc/openvpn/server/server.conf"
        self.status_log_path = "/tmp/openvpn/openvpn-status.log"
        self.clients: Dict[str, VPNClient] = {}
        self.server_process = None
        self.is_running = False
        self.start_time = None
        
        # VPN network configuration
        self.vpn_network = ipaddress.IPv4Network("10.8.0.0/24")
        self.vpn_server_ip = "10.8.0.1"
        self.vpn_port = 1194
        
        # Integration with ISP stack
        self.dhcp_server = None
        self.nat_engine = None
        self.firewall = None
        self.router = None
        self.traffic_router = None  # New traffic router component
        
        # Status monitoring thread
        self.monitor_thread = None
        self.monitor_running = False
        
        # Client configuration storage
        self.config_storage_path = "/tmp/vpn_client_configs"
        os.makedirs(self.config_storage_path, exist_ok=True)
        
    def set_isp_components(self, dhcp_server=None, nat_engine=None, firewall=None, router=None, traffic_router=None):
        """Set references to ISP stack components for integration"""
        self.dhcp_server = dhcp_server
        self.nat_engine = nat_engine
        self.firewall = firewall
        self.router = router
        self.traffic_router = traffic_router
        
        # Configure traffic router with other components
        if self.traffic_router:
            self.traffic_router.set_components(
                nat_engine=nat_engine,
                firewall=firewall,
                dhcp_server=dhcp_server
            )
        
    def start_server(self) -> bool:
        """Start the OpenVPN server with traffic routing"""
        try:
            if self.is_running:
                logger.warning("OpenVPN server is already running")
                return True
            
            # Ensure configuration exists
            if not os.path.exists(self.server_config_path):
                logger.error(f"OpenVPN server configuration not found: {self.server_config_path}")
                return False
            
            # Start traffic router first
            if self.traffic_router and not self.traffic_router.is_running:
                if not self.traffic_router.start():
                    logger.error("Failed to start traffic router")
                    return False
            

            # Start OpenVPN server
            self.server_process = subprocess.Popen(['sudo', 'openvpn', '--config', self.server_config_path], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
            self.is_running = True
            self.start_time = time.time()
            logger.info("OpenVPN server started successfully")
            
            # Start monitoring thread
            self.start_monitoring()
            
            # Configure firewall rules for VPN
            self._configure_vpn_firewall()
            
            # Configure NAT for VPN traffic
            self._configure_vpn_nat()
            
            return True
                
        except Exception as e:
            logger.error(f"Error starting OpenVPN server: {e}")
            return False
    
    def stop_server(self) -> bool:
        """Stop the OpenVPN server and traffic routing"""
        try:
            if not self.is_running:
                logger.warning("OpenVPN server is not running")
                return True
            
            # Stop monitoring
            self.stop_monitoring()
            
            # Remove all client routes before stopping
            if self.traffic_router:
                for client_id in list(self.clients.keys()):
                    self.traffic_router.remove_client_route(client_id)
            
            # Stop OpenVPN server
            if self.server_process:
                self.server_process.terminate()
                self.server_process.wait(timeout=5)
                if self.server_process.poll() is None:
                    self.server_process.kill()
                self.server_process = None
            self.is_running = False
            self.start_time = None
            self.clients.clear()
            logger.info("OpenVPN server stopped successfully")
            return True
            
        except Exception as e:
            logger.error(f"Error stopping OpenVPN server: {e}")
            return False
    
    def start_monitoring(self):
        """Start the client monitoring thread"""
        if self.monitor_thread and self.monitor_thread.is_alive():
            return
        
        self.monitor_running = True
        self.monitor_thread = threading.Thread(target=self._monitor_clients, daemon=True)
        self.monitor_thread.start()
        logger.info("Started OpenVPN client monitoring")
    
    def stop_monitoring(self):
        """Stop the client monitoring thread"""
        self.monitor_running = False
        if self.monitor_thread:
            self.monitor_thread.join(timeout=5)
        logger.info("Stopped OpenVPN client monitoring")
    
    def _monitor_clients(self):
        """Monitor connected VPN clients"""
        while self.monitor_running:
            try:
                self._update_client_status()
                time.sleep(10)  # Update every 10 seconds
            except Exception as e:
                logger.error(f"Error monitoring VPN clients: {e}")
                time.sleep(30)  # Wait longer on error
    
    def _update_client_status(self):
        """Update client status from OpenVPN status log and manage traffic routing"""
        try:
            with open(self.status_log_path, 'r') as f:
                lines = f.readlines()

            new_clients = {}
            client_section = False
            for line in lines:
                if line.startswith('ROUTING TABLE'):
                    client_section = False
                if client_section and not line.startswith('GLOBAL STATS'):
                    parts = line.strip().split(',')
                    if len(parts) >= 5:
                        common_name = parts[0]
                        real_ip_port = parts[1]
                        virtual_ip = parts[2]
                        bytes_received = int(parts[3])
                        bytes_sent = int(parts[4])
                        connected_since = float(parts[5]) # Assuming this is a timestamp

                        # Extract IP address from real_ip_port (e.g., 1.2.3.4:12345)
                        ip_address = real_ip_port.split(':')[0]

                        client = VPNClient(
                            client_id=common_name,
                            common_name=common_name,
                            ip_address=virtual_ip,
                            connected_at=connected_since,
                            bytes_received=bytes_received,
                            bytes_sent=bytes_sent,
                            status="connected",
                            routed_through_vpn=True
                        )
                        new_clients[common_name] = client
                if line.startswith('COMMON NAME'):
                    client_section = True
            self.clients = new_clients
            
        except Exception as e:
            logger.error(f"Error updating client status: {e}")
    
    def _sync_with_dhcp(self):
        """Sync VPN clients with DHCP server"""
        try:
            for client in self.clients.values():
                if client.ip_address != "unknown":
                    # Register VPN client IP with DHCP server
                    # This allows the ISP stack to track VPN clients
                    if hasattr(self.dhcp_server, 'register_static_lease'):
                        self.dhcp_server.register_static_lease(
                            client.common_name,
                            client.ip_address,
                            "VPN Client"
                        )
        except Exception as e:
            logger.error(f"Error syncing with DHCP: {e}")
    
    def _configure_vpn_firewall(self):
        """Configure firewall rules for VPN traffic"""
        try:
            if not self.firewall:
                return
            
            # Add firewall rules for VPN
            vpn_rules = [
                {
                    "rule_id": "allow_openvpn",
                    "priority": 10,
                    "action": "ACCEPT",
                    "direction": "BOTH",
                    "dest_port": str(self.vpn_port),
                    "protocol": "UDP",
                    "description": "Allow OpenVPN traffic",
                    "enabled": True
                },
                {
                    "rule_id": "allow_vpn_network",
                    "priority": 11,
                    "action": "ACCEPT",
                    "direction": "BOTH",
                    "source_network": str(self.vpn_network),
                    "description": "Allow VPN client network traffic",
                    "enabled": True
                }
            ]
            
            for rule in vpn_rules:
                if hasattr(self.firewall, 'add_rule'):
                    self.firewall.add_rule(rule)
            
            logger.info("Configured firewall rules for VPN")
            
        except Exception as e:
            logger.error(f"Error configuring VPN firewall: {e}")
    
    def _configure_vpn_nat(self):
        """Configure NAT for VPN traffic"""
        try:
            # NAT configuration will be handled by the external environment (e.g., HuggingFace Spaces setup)
            # or by the underlying network infrastructure. We are removing direct iptables calls.
            logger.info("Skipping direct iptables NAT configuration as per instructions.")
            
        except Exception as e:
            logger.error(f"Error configuring VPN NAT: {e}")
    
    def get_server_status(self) -> VPNServerStatus:
        """Get current server status"""
        total_bytes_received = sum(client.bytes_received for client in self.clients.values())
        total_bytes_sent = sum(client.bytes_sent for client in self.clients.values())
        uptime = time.time() - self.start_time if self.start_time else 0
        
        return VPNServerStatus(
            is_running=self.is_running,
            connected_clients=len(self.clients),
            total_bytes_received=total_bytes_received,
            total_bytes_sent=total_bytes_sent,
            uptime=uptime,
            server_ip=self.vpn_server_ip,
            server_port=self.vpn_port
        )    
    def get_connected_clients(self) -> List[Dict[str, Any]]:
        """Get list of connected clients"""
        return [asdict(client) for client in self.clients.values()]
    
    def disconnect_client(self, client_id: str) -> bool:
        """Disconnect a specific client"""
        try:
            if client_id not in self.clients:
                return False
            
            # Send kill signal to specific client
            # This requires OpenVPN management interface, simplified for now
            logger.info(f"Disconnecting client: {client_id}")
            
            # Remove from clients dict
            del self.clients[client_id]
            return True
            
        except Exception as e:
            logger.error(f"Error disconnecting client {client_id}: {e}")
            return False
    
    def generate_client_config(self, client_name: str, server_ip: str) -> str:
        """Generate client configuration file with embedded certificates"""
        try:
            # Read real CA certificate
            ca_cert_path = "/etc/openvpn/server/ca.crt"
            with open(ca_cert_path, 'r') as f:
                ca_cert = f.read()

            client_cert_path = f"/home/ubuntu/easy-rsa/pki/issued/{client_name}.crt"
            with open(client_cert_path, 'r') as f:
                client_cert = f.read()

            client_key_path = f"/home/ubuntu/easy-rsa/pki/private/{client_name}.key"
            with open(client_key_path, 'r') as f:
                client_key = f.read()

            # Generate complete client configuration
            client_config = f"""# OpenVPN Client Configuration for {client_name}
# Generated by Virtual ISP Stack
# Server: {server_ip}:{self.vpn_port}

client
dev tun
proto udp
remote {server_ip} {self.vpn_port}
resolv-retry infinite
nobind
persist-key
persist-tun
cipher AES-256-CBC
auth SHA256
verb 3
key-direction 1
redirect-gateway def1 bypass-dhcp
dhcp-option DNS 8.8.8.8
dhcp-option DNS 8.8.4.4
remote-cert-tls server

# Embedded CA Certificate
<ca>
{ca_cert}
</ca>

# Embedded Client Certificate
<cert>
{client_cert}
</cert>

# Embedded Client Private Key
<key>
{client_key}
</key>

# TLS Authentication Key (optional, for extra security)
# <tls-auth>
# -----BEGIN OpenVPN Static key V1-----
# [TLS-AUTH-KEY-CONTENT-WOULD-GO-HERE]
# -----END OpenVPN Static key V1-----
# </tls-auth>
"""
            
            logger.info(f"Generated client configuration for {client_name}")
            return client_config
            
        except Exception as e:
            logger.error(f"Error generating client config: {e}")
            return ""
    
    def save_client_config(self, client_name: str, config_content: str) -> bool:
        """Save client configuration to storage"""
        try:
            config_file_path = os.path.join(self.config_storage_path, f"{client_name}.ovpn")
            with open(config_file_path, 'w') as f:
                f.write(config_content)
            
            logger.info(f"Saved client configuration for {client_name}")
            return True
            
        except Exception as e:
            logger.error(f"Error saving client config for {client_name}: {e}")
            return False
    
    def load_client_config(self, client_name: str) -> str:
        """Load client configuration from storage"""
        try:
            config_file_path = os.path.join(self.config_storage_path, f"{client_name}.ovpn")
            if not os.path.exists(config_file_path):
                return ""
            
            with open(config_file_path, 'r') as f:
                config_content = f.read()
            
            logger.info(f"Loaded client configuration for {client_name}")
            return config_content
            
        except Exception as e:
            logger.error(f"Error loading client config for {client_name}: {e}")
            return ""
    
    def list_client_configs(self) -> List[str]:
        """List all stored client configurations"""
        try:
            config_files = []
            if os.path.exists(self.config_storage_path):
                for filename in os.listdir(self.config_storage_path):
                    if filename.endswith('.ovpn'):
                        client_name = filename[:-5]  # Remove .ovpn extension
                        config_files.append(client_name)
            
            return config_files
            
        except Exception as e:
            logger.error(f"Error listing client configs: {e}")
            return []
    
    def delete_client_config(self, client_name: str) -> bool:
        """Delete client configuration from storage"""
        try:
            config_file_path = os.path.join(self.config_storage_path, f"{client_name}.ovpn")
            if os.path.exists(config_file_path):
                os.remove(config_file_path)
                logger.info(f"Deleted client configuration for {client_name}")
                return True
            else:
                logger.warning(f"Client configuration for {client_name} not found")
                return False
                
        except Exception as e:
            logger.error(f"Error deleting client config for {client_name}: {e}")
            return False
    
    def generate_and_save_client_config(self, client_name: str, server_ip: str) -> str:
        """Generate client configuration and save it to storage"""
        try:
            config_content = self.generate_client_config(client_name, server_ip)
            if config_content:
                if self.save_client_config(client_name, config_content):
                    return config_content
            return ""
            
        except Exception as e:
            logger.error(f"Error generating and saving client config for {client_name}: {e}")
            return ""
    
    def get_statistics(self) -> Dict[str, Any]:
        """Get comprehensive VPN statistics"""
        status = self.get_server_status()
        
        return {
            "server_status": asdict(status),
            "connected_clients": self.get_connected_clients(),
            "network_config": {
                "vpn_network": str(self.vpn_network),
                "server_ip": self.vpn_server_ip,
                "server_port": self.vpn_port
            },
            "integration_status": {
                "dhcp_integrated": self.dhcp_server is not None,
                "nat_integrated": self.nat_engine is not None,
                "firewall_integrated": self.firewall is not None,
                "router_integrated": self.router is not None
            }
        }

# Global OpenVPN manager instance
openvpn_manager = None

def initialize_openvpn_manager(config: Dict[str, Any]) -> OpenVPNManager:
    """Initialize the OpenVPN manager"""
    global openvpn_manager
    openvpn_manager = OpenVPNManager(config)
    return openvpn_manager

def get_openvpn_manager() -> Optional[OpenVPNManager]:
    """Get the global OpenVPN manager instance"""
    return openvpn_manager