core/outline_server.py

"""
Enhanced Outline VPN Server Implementation
Core server component with offline access and internet sharing support
"""

import asyncio
import logging
import json
import os
import psutil
from typing import Dict, Optional, List, Tuple
from .shadowsocks_protocol import ShadowsocksProtocol
from .nat_engine import NATEngine
from .traffic_router import TrafficRouter
from .outline_config import OutlineManager, OutlineConfig
from .ikev2_server import IKEv2Server as IPsecManager
from .session_tracker import SessionTracker
from .port_manager import PortManager

logger = logging.getLogger(__name__)

class OutlineServer:
    def __init__(self, config: Dict):
        self.config = config
        self.outline_manager = OutlineManager()
        self.ipsec_manager = IPsecManager(config["server"]["host"], logger)
        
        # Initialize authentication
        from .vpn_auth import VPNAuthManager
        from .database_init import init_db
        
        # Initialize database
        if not init_db():
            raise RuntimeError("Failed to initialize database")
            
        self.auth_manager = VPNAuthManager()
        
        # Initialize port manager
        self.port_manager = PortManager()
        
        # Initialize components
        self.nat_engine = NATEngine(logger)
        self.traffic_router = TrafficRouter({
            "vpn_host": "0.0.0.0", # Listen on all interfaces
            "vpn_port": config["server"]["port"],
            "virtual_network": config["server"]["virtual_network"]
        }, logger)
        
        # Initialize session tracker for offline support
        self.session_tracker = SessionTracker()
        
        self.sessions = {}
        self.is_running = False
        self.bound_ports: List[Tuple[int, str]] = []  # List of (port, service) tuples
        
    async def setup_internet_sharing(self):
        """Configure internet sharing for VPN clients"""
        try:
            interface = self.config["server"].get("interface", "eth0")
            await self.nat_engine.setup_nat(interface)
            logger.info(f"Internet sharing enabled on interface {interface}")
        except Exception as e:
            logger.error(f"Failed to setup internet sharing: {e}")
            raise
    async def start(self):
        """Start the VPN server"""
        if self.is_running:
            logger.warning("Outline VPN server is already running")
            return

        try:
            # Clean up any existing connections first
            await self.cleanup_existing_connections()
            
            # Setup network
            await self.setup_internet_sharing()
            
            # Start components in order with proper error handling
            try:
                await self.nat_engine.start()
            except Exception as e:
                logger.error(f"Failed to start NAT engine: {e}")
                raise

            try:
                await self.traffic_router.start()
            except Exception as e:
                await self.nat_engine.stop()
                logger.error(f"Failed to start traffic router: {e}")
                raise

            try:
                await self.ipsec_manager.start()
            except Exception as e:
                await self.nat_engine.stop()
                await self.traffic_router.stop()
                logger.error(f"Failed to start IKEv2 service: {e}")
                raise
            
            self.is_running = True
            logger.info(f"Outline VPN server started successfully on port {self.config['server']['port']}")
                
        except Exception as e:
            logger.error(f"Failed to start server: {e}")
            await self.stop()  # Ensure cleanup on failure
            raise

    async def cleanup_existing_connections(self):
        """Clean up any existing connections before starting"""
        import psutil
        current_pid = os.getpid()
        
        # Find and clean up any existing VPN processes
        for proc in psutil.process_iter(['pid', 'name', 'connections']):
            try:
                if proc.pid != current_pid:
                    for conn in proc.connections():
                        if conn.laddr.port == self.config['server']['port']:
                            logger.warning(f"Found existing process using port {self.config['server']['port']}, terminating...")
                            proc.terminate()
                            proc.wait(timeout=5)
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.TimeoutExpired):
                continue

    async def stop(self):
        """Stop the VPN server"""
        self.is_running = False
        await self.traffic_router.stop()
        await self.nat_engine.stop()
        await self.ipsec_manager.stop()
        
        # Close all active sessions
        for session in self.sessions.values():
            await session.close()
        self.sessions.clear()
        
        logger.info("VPN server stopped")

    async def _handle_client(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
        """Handle new client connections"""
        peer = writer.get_extra_info('peername')
        logger.info(f"New connection from {peer}")
        
        try:
            # First packet contains access key
            first_packet = await reader.read(64)
            access_key = self._extract_access_key(first_packet)
            
            # Validate access key
            user = self.outline_manager.get_user_by_key(access_key)
            if not user:
                logger.warning(f"Invalid access key from {peer}")
                writer.close()
                return
                
            # Create protocol handler
            protocol = ShadowsocksProtocol(access_key)
            self.sessions[user.user_id] = protocol
            
            # Handle the connection
            await protocol.handle_connection(reader, writer)
            
        except Exception as e:
            logger.error(f"Error handling client {peer}: {e}")
        finally:
            if not writer.is_closing():
                writer.close()
                await writer.wait_closed()

    def _extract_access_key(self, packet: bytes) -> str:
        """Extract access key from initial packet"""
        return packet[:32].hex()

    def get_stats(self) -> Dict:
        """Get server statistics"""
        return {
            "is_running": self.is_running,
            "active_sessions": len(self.sessions),
            "traffic_stats": self.traffic_router.get_stats(),
            "nat_stats": self.nat_engine.get_stats()
        }