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 = "/var/log/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
            
            # Create log directory
            os.makedirs("/var/log/openvpn", exist_ok=True)
            
            # Start OpenVPN server
            cmd = [
                "/usr/bin/sudo", "openvpn",
                "--config", self.server_config_path,
                "--daemon", "openvpn-server",
                "--log", "/var/log/openvpn/openvpn.log",
                "--status", self.status_log_path, "10"
            ]
            
            result = subprocess.run(cmd, capture_output=True, text=True)
            
            if result.returncode == 0:
                self.is_running = True
                self.start_time = time.time()
                logger.info("OpenVPN server started successfully with traffic routing")
                
                # 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
            else:
                logger.error(f"Failed to start OpenVPN server: {result.stderr}")
                return False
                
        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)
            
            # Kill OpenVPN process
            result = subprocess.run(["/usr/bin/sudo", "pkill", "-f", "openvpn.*server"], 
                                  capture_output=True, text=True)
            
            # Stop traffic router
            if self.traffic_router and self.traffic_router.is_running:
                self.traffic_router.stop()
            
            self.is_running = False
            self.start_time = None
            self.clients.clear()
            
            logger.info("OpenVPN server and traffic routing stopped")
            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:
            if not os.path.exists(self.status_log_path):
                return
            
            with open(self.status_log_path, 'r') as f:
                content = f.read()
            
            # Parse OpenVPN status log
            lines = content.split('\n')
            client_section = False
            routing_section = False
            
            current_clients = {}
            previous_clients = set(self.clients.keys())
            
            for line in lines:
                line = line.strip()
                
                if line.startswith("OpenVPN CLIENT LIST"):
                    client_section = True
                    continue
                elif line.startswith("ROUTING TABLE"):
                    client_section = False
                    routing_section = True
                    continue
                elif line.startswith("GLOBAL STATS"):
                    routing_section = False
                    continue
                
                if client_section and line and not line.startswith("Updated,"):
                    # Parse client line: Common Name,Real Address,Bytes Received,Bytes Sent,Connected Since
                    parts = line.split(',')
                    if len(parts) >= 5:
                        common_name = parts[0]
                        real_address = parts[1]
                        bytes_received = int(parts[2]) if parts[2].isdigit() else 0
                        bytes_sent = int(parts[3]) if parts[3].isdigit() else 0
                        connected_since = parts[4]
                        
                        # Get VPN IP from routing table (will be parsed later)
                        vpn_ip = "unknown"
                        
                        # Check if this is an existing client
                        routed_through_vpn = False
                        if common_name in self.clients:
                            routed_through_vpn = self.clients[common_name].routed_through_vpn
                        
                        client = VPNClient(
                            client_id=common_name,
                            common_name=common_name,
                            ip_address=vpn_ip,
                            connected_at=time.time(),  # Simplified for now
                            bytes_received=bytes_received,
                            bytes_sent=bytes_sent,
                            routed_through_vpn=routed_through_vpn
                        )
                        current_clients[common_name] = client
                
                elif routing_section and line and not line.startswith("Virtual Address,"):
                    # Parse routing line: Virtual Address,Common Name,Real Address,Last Ref
                    parts = line.split(',')
                    if len(parts) >= 2:
                        vpn_ip = parts[0]
                        common_name = parts[1]
                        
                        if common_name in current_clients:
                            current_clients[common_name].ip_address = vpn_ip
            
            # Handle new clients - set up traffic routing
            new_clients = set(current_clients.keys()) - previous_clients
            for client_id in new_clients:
                client = current_clients[client_id]
                if client.ip_address != "unknown" and self.traffic_router:
                    # Add client route for free data access
                    if self.traffic_router.add_client_route(client_id, client.ip_address):
                        client.routed_through_vpn = True
                        logger.info(f"Added traffic routing for new VPN client: {client_id} ({client.ip_address})")
            
            # Handle disconnected clients - clean up routing
            disconnected_clients = previous_clients - set(current_clients.keys())
            for client_id in disconnected_clients:
                if self.traffic_router:
                    self.traffic_router.remove_client_route(client_id)
                    logger.info(f"Removed traffic routing for disconnected VPN client: {client_id}")
            
            # Update clients dictionary
            self.clients = current_clients
            
            # Integrate with DHCP server if available
            if self.dhcp_server:
                self._sync_with_dhcp()
            
        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"
            try:
                with open(ca_cert_path, 'r') as f:
                    ca_cert = f.read()
            except FileNotFoundError:
                # Fallback to embedded certificate for development
                ca_cert = """-----BEGIN CERTIFICATE-----
MIIDMzCCAhugAwIBAgIUNO765P4t/yD/PnIFTMVs0Q32TJYwDQYJKoZIhvcNAQEL
BQAwDjEMMAoGA1UEAwwDeWVzMB4XDTI1MDgwMjAxMjkzNVoXDTM1MDczMTAxMjkz
NVowDjEMMAoGA1UEAwwDeWVzMIIBIjANBgkqhkiG9w0BAQEFAAOCAQ8AMIIBCgKC
AQEAtwhMGXouHnHBRd2RhdrW8sOMgqt4wDXZC0J+4UMjOX6Y7t2O1Sgw/sWhwFPk
QF/cMoQIvsucklPogcnzzGtv9zDkAXyVyCC27UYbg8JfWZK3ZMrt6dfEmYf4KKXm
D6PLn9guxzBB63dhEWx/7fd6H9C/rK/u0rOh15DQRnfEI468cmXS5uNg8ke/73+y
Gzb6q7ZOFByBAwM0hW0lStBaIIcxouFrIK8B72O8H+6t10K1GvgiBhKvM3cc8dpN
y4qvRoN/o+eXarZG7G9dfm9OFgdd9LoXPTTbO+ftFPKOq4F41PnMd2Zcyk7P3GCr
3oK7NbISxZ5efLpy45lgSpqKBwIDAQABo4GIMIGFMB0GA1UdDgQWBBQIi0Er30cV
Qzi+U/LPV4Lf3yvGIzBJBgNVHSMEQjBAgBQIi0Er30cVQzi+U/LPV4Lf3yvGI6ES
pBAwDjEMMAoGA1UEAwwDeWVzghQ07vrk/i3/IP8+cgVMxWzRDfZMljAMBgNVHRME
BTADAQH/MAsGA1UdDwQEAwIBBjANBgkqhkiG9w0BAQsFAAOCAQEAHzfSFbi1G7WC
vMSOqSv4/jlBExnz/AlLUBHhgDomIdLK8Pb3tyCD5IYkmi0NT5x6DORcOV2ow1JZ
o4BL7OVV+fhz3VKXEpG+s3gq5j2m+raqLtu6QKBGg7SIUZ4MLjggvAcPjsK+n8sK
86sAUFVTccBxJlKBShAUPSNihyWwxB4PQFvwhefNQSoID1kAB2Fzf1beMX6Gp6Lj
ldI6e63lpYtIbp4+2F5SxJ/hGTUx+nWbOAHPvhBfhN6sEu9G1C5KPR0cm+xxOpZ9
lA7y4Dea7pyVybR/b7lFquE3TReXCoLx79UNNSv8erIlsy1jh9yXDnTCk8SN1dpO
YwJ9U0AHXA==
-----END CERTIFICATE-----"""

            # Sample client certificate (in production, generate unique per client)
            client_cert = f"""-----BEGIN CERTIFICATE-----
MIIDSzCCAjOgAwIBAgIU{client_name}1234567890abcdefghijklmnopqrstuvwxyz
ABCDEFGHIJKLMNOPQRSTUVWXYZ1234567890abcdefghijklmnopqrstuvwxyz1234567
890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12345
67890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz123
4567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1
234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxy
z1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuv
wxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrs
tuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmno
pqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghij
klmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcd
efghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567
890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxy
z1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnop
qrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcde
fghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12
34567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijk
lmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12
34567890abcdefghijklmnopqrstuvwxyz1234567890abc
defghijklmnopqrstuvwxyz1234567890abcdefg
hijklmnopqrstuvwxyz1234567890
abcdefghijklmnopqr
stuvwxyz12345
67890abc
def
-----END CERTIFICATE-----"""

            # Sample client private key (in production, generate unique per client)
            client_key = f"""-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQC{client_name}1234567
890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz123456
7890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12345
67890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234
567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz123
4567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12
34567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1
234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxy
z1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuv
wxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrs
tuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmno
pqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghij
klmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcd
efghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567
890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxy
z1234567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijklmnop
qrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz1234567890abcde
fghijklmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12
34567890abcdefghijklmnopqrstuvwxyz1234567890abcdefghijk
lmnopqrstuvwxyz1234567890abcdefghijklmnopqrstuvwxyz12
34567890abcdefghijklmnopqrstuvwxyz1234567890abc
defghijklmnopqrstuvwxyz1234567890abcdefg
hijklmnopqrstuvwxyz1234567890
abcdefghijklmnopqr
stuvwxyz12345
67890abc
defghijk
lmnopqr
stuv
-----END PRIVATE KEY-----"""

            # 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