golem-vm-provider 0.1.0__py3-none-any.whl

provider/vm/port_manager.py

@@ -0,0 +1,196 @@
+import os
+import json
+import logging
+import asyncio
+from pathlib import Path
+from typing import Optional, Set, List, Dict
+from threading import Lock
+
+from ..config import settings
+from ..network.port_verifier import PortVerifier, PortVerificationResult
+from ..utils.port_display import PortVerificationDisplay
+
+logger = logging.getLogger(__name__)
+
+class PortManager:
+    """Manages port allocation and verification for VM SSH proxying."""
+    
+    def __init__(
+        self,
+        start_port: int = 50800,
+        end_port: int = 50900,
+        state_file: Optional[str] = None,
+        port_check_servers: Optional[List[str]] = None,
+        discovery_port: Optional[int] = None
+    ):
+        """Initialize the port manager.
+        
+        Args:
+            start_port: Beginning of port range
+            end_port: End of port range (exclusive)
+            state_file: Path to persist port assignments
+            port_check_servers: List of URLs for port checking services
+        """
+        self.start_port = start_port
+        self.end_port = end_port
+        self.state_file = state_file or os.path.expanduser("~/.golem/provider/ports.json")
+        self.lock = Lock()
+        self._used_ports: dict[str, int] = {}  # vm_id -> port
+        self.verified_ports: Set[int] = set()
+        
+        # Initialize port verifier with default servers
+        self.port_check_servers = port_check_servers or [
+            # "http://portcheck1.golem.network:7466",
+            # "http://portcheck2.golem.network:7466",
+            # "http://portcheck3.golem.network:7466",
+            "http://localhost:9000"  # Fallback for local development
+        ]
+        self.discovery_port = discovery_port or settings.PORT
+        self.port_verifier = PortVerifier(
+            self.port_check_servers,
+            discovery_port=self.discovery_port
+        )
+        
+        self._load_state()
+    
+    async def initialize(self) -> bool:
+        """Initialize port manager with verification.
+        
+        Returns:
+            bool: True if required ports were verified successfully
+        """
+        from ..config import settings
+        
+        display = PortVerificationDisplay(
+            provider_port=self.discovery_port,
+            port_range_start=self.start_port,
+            port_range_end=self.end_port
+        )
+        display.print_header()
+        
+        # Verify all ports together (discovery port and SSH ports)
+        all_ports = [self.discovery_port] + list(range(self.start_port, self.end_port))
+        logger.info(f"Verifying all ports: discovery port {self.discovery_port} and SSH ports {self.start_port}-{self.end_port}...")
+        
+        results = await self.port_verifier.verify_ports(all_ports)
+        
+        # Display discovery port status with animation
+        discovery_result = results[self.discovery_port]
+        await display.print_discovery_status(discovery_result)
+        
+        if not discovery_result.accessible:
+            logger.error(f"Failed to verify discovery port: {discovery_result.error}")
+            # Print summary before returning
+            display.print_summary(discovery_result, {})
+            return False
+            
+        # Display SSH ports status with animation
+        ssh_results = {port: result for port, result in results.items() if port != self.discovery_port}
+        await display.print_ssh_status(ssh_results)
+        
+        # Store verified ports
+        self.verified_ports = {port for port, result in ssh_results.items() if result.accessible}
+        
+        # Only show critical issues and quick fix if there are problems
+        if not discovery_result.accessible or not self.verified_ports:
+            display.print_critical_issues(discovery_result, ssh_results)
+            display.print_quick_fix(discovery_result, ssh_results)
+        
+        # Print precise summary of current status
+        display.print_summary(discovery_result, ssh_results)
+        
+        if not self.verified_ports:
+            logger.error("No SSH ports were verified as accessible")
+            return False
+            
+        logger.info(f"Successfully verified {len(self.verified_ports)} SSH ports")
+        return True
+        
+    def _load_state(self) -> None:
+        """Load port assignments from state file."""
+        try:
+            state_path = Path(self.state_file)
+            if state_path.exists():
+                with open(state_path, 'r') as f:
+                    self._used_ports = json.load(f)
+                logger.info(f"Loaded port assignments for {len(self._used_ports)} VMs")
+            else:
+                state_path.parent.mkdir(parents=True, exist_ok=True)
+                self._save_state()
+        except Exception as e:
+            logger.error(f"Failed to load port state: {e}")
+            self._used_ports = {}
+    
+    def _save_state(self) -> None:
+        """Save current port assignments to state file."""
+        try:
+            with open(self.state_file, 'w') as f:
+                json.dump(self._used_ports, f)
+        except Exception as e:
+            logger.error(f"Failed to save port state: {e}")
+    
+    def _get_used_ports(self) -> Set[int]:
+        """Get set of currently used ports."""
+        return set(self._used_ports.values())
+    
+    def allocate_port(self, vm_id: str) -> Optional[int]:
+        """Allocate a verified port for a VM.
+        
+        Args:
+            vm_id: Unique identifier for the VM
+            
+        Returns:
+            Allocated port number or None if allocation failed
+        """
+        with self.lock:
+            # Check if VM already has a port
+            if vm_id in self._used_ports:
+                port = self._used_ports[vm_id]
+                if port in self.verified_ports:
+                    return port
+                else:
+                    # Previously allocated port is no longer verified
+                    self._used_ports.pop(vm_id)
+            
+            used_ports = self._get_used_ports()
+            
+            # Find first available verified port
+            for port in sorted(self.verified_ports):
+                if port not in used_ports:
+                    self._used_ports[vm_id] = port
+                    self._save_state()
+                    logger.info(f"Allocated verified port {port} for VM {vm_id}")
+                    return port
+            
+            logger.error("No verified ports available for allocation")
+            return None
+    
+    def deallocate_port(self, vm_id: str) -> None:
+        """Release a port allocation for a VM.
+        
+        Args:
+            vm_id: Unique identifier for the VM
+        """
+        with self.lock:
+            if vm_id in self._used_ports:
+                port = self._used_ports.pop(vm_id)
+                self._save_state()
+                logger.info(f"Deallocated port {port} for VM {vm_id}")
+    
+    def get_port(self, vm_id: str) -> Optional[int]:
+        """Get currently allocated port for a VM.
+        
+        Args:
+            vm_id: Unique identifier for the VM
+            
+        Returns:
+            Port number or None if VM has no allocation
+        """
+        return self._used_ports.get(vm_id)
+    
+    def cleanup(self) -> None:
+        """Remove all port allocations."""
+        with self.lock:
+            self._used_ports.clear()
+            self._save_state()
+            logger.info("Cleared all port allocations")

provider/vm/proxy_manager.py

@@ -0,0 +1,239 @@
+import os
+import json
+import asyncio
+import logging
+from pathlib import Path
+from typing import Optional, Dict
+from asyncio import Task, Transport, Protocol
+
+from .port_manager import PortManager
+
+logger = logging.getLogger(__name__)
+
+class SSHProxyProtocol(Protocol):
+    """Protocol for handling SSH proxy connections."""
+    
+    def __init__(self, target_host: str, target_port: int):
+        self.target_host = target_host
+        self.target_port = target_port
+        self.transport: Optional[Transport] = None
+        self.target_transport: Optional[Transport] = None
+        self.target_protocol: Optional['SSHTargetProtocol'] = None
+        self.buffer = bytearray()
+    
+    def connection_made(self, transport: Transport) -> None:
+        """Called when connection is established."""
+        self.transport = transport
+        asyncio.create_task(self.connect_to_target())
+    
+    async def connect_to_target(self) -> None:
+        """Establish connection to target."""
+        try:
+            loop = asyncio.get_running_loop()
+            self.target_protocol = SSHTargetProtocol(self)
+            target_transport, _ = await loop.create_connection(
+                lambda: self.target_protocol,
+                self.target_host,
+                self.target_port
+            )
+            self.target_transport = target_transport
+            # If we have buffered data, send it now
+            if self.buffer:
+                self.target_transport.write(self.buffer)
+                self.buffer.clear()
+        except Exception as e:
+            logger.error(f"Failed to connect to target {self.target_host}:{self.target_port}: {e}")
+            if self.transport:
+                self.transport.close()
+    
+    def data_received(self, data: bytes) -> None:
+        """Forward received data to target."""
+        if self.target_transport and not self.target_transport.is_closing():
+            self.target_transport.write(data)
+        else:
+            # Buffer data until target connection is established
+            self.buffer.extend(data)
+    
+    def connection_lost(self, exc: Optional[Exception]) -> None:
+        """Handle connection loss."""
+        if exc:
+            logger.error(f"Client connection lost with error: {exc}")
+        if self.target_transport and not self.target_transport.is_closing():
+            self.target_transport.close()
+
+class SSHTargetProtocol(Protocol):
+    """Protocol for handling target SSH connections."""
+    
+    def __init__(self, client_protocol: SSHProxyProtocol):
+        self.client_protocol = client_protocol
+        self.transport: Optional[Transport] = None
+    
+    def connection_made(self, transport: Transport) -> None:
+        """Called when connection is established."""
+        self.transport = transport
+    
+    def data_received(self, data: bytes) -> None:
+        """Forward received data to client."""
+        if (self.client_protocol.transport and 
+            not self.client_protocol.transport.is_closing()):
+            self.client_protocol.transport.write(data)
+    
+    def connection_lost(self, exc: Optional[Exception]) -> None:
+        """Handle connection loss."""
+        if exc:
+            logger.error(f"Target connection lost with error: {exc}")
+        if (self.client_protocol.transport and 
+            not self.client_protocol.transport.is_closing()):
+            self.client_protocol.transport.close()
+
+class ProxyServer:
+    """Manages a single proxy server instance."""
+    
+    def __init__(self, listen_port: int, target_host: str, target_port: int = 22):
+        """Initialize proxy server.
+        
+        Args:
+            listen_port: Port to listen on
+            target_host: Target host to forward to
+            target_port: Target port (default: 22 for SSH)
+        """
+        self.listen_port = listen_port
+        self.target_host = target_host
+        self.target_port = target_port
+        self.server: Optional[asyncio.AbstractServer] = None
+    
+    async def start(self) -> None:
+        """Start the proxy server."""
+        loop = asyncio.get_running_loop()
+        
+        try:
+            self.server = await loop.create_server(
+                lambda: SSHProxyProtocol(self.target_host, self.target_port),
+                '0.0.0.0',  # Listen on all interfaces
+                self.listen_port
+            )
+            logger.info(f"Proxy server listening on port {self.listen_port}")
+        except Exception as e:
+            logger.error(f"Failed to start proxy server on port {self.listen_port}: {e}")
+            raise
+    
+    async def stop(self) -> None:
+        """Stop the proxy server."""
+        if self.server:
+            self.server.close()
+            await self.server.wait_closed()
+            logger.info(f"Proxy server on port {self.listen_port} stopped")
+
+class PythonProxyManager:
+    """Manages proxy servers for VM SSH access."""
+    
+    def __init__(
+        self,
+        port_manager: Optional[PortManager] = None,
+        state_file: Optional[str] = None
+    ):
+        """Initialize the proxy manager.
+        
+        Args:
+            port_manager: Port allocation manager
+            state_file: Path to persist proxy state
+        """
+        self.port_manager = port_manager or PortManager()
+        self.state_file = state_file or os.path.expanduser("~/.golem/provider/proxy_state.json")
+        self._proxies: Dict[str, ProxyServer] = {}  # vm_id -> ProxyServer
+        self._load_state()
+    
+    def _load_state(self) -> None:
+        """Load proxy state from file."""
+        try:
+            state_path = Path(self.state_file)
+            if state_path.exists():
+                with open(state_path, 'r') as f:
+                    state = json.load(f)
+                    # We only need to restore port allocations
+                    # Actual proxy servers will be recreated as needed
+                    logger.info(f"Loaded proxy state for {len(state)} VMs")
+        except Exception as e:
+            logger.error(f"Failed to load proxy state: {e}")
+    
+    def _save_state(self) -> None:
+        """Save current proxy state to file."""
+        try:
+            state = {
+                vm_id: {
+                    'port': proxy.listen_port,
+                    'target': proxy.target_host
+                }
+                for vm_id, proxy in self._proxies.items()
+            }
+            os.makedirs(os.path.dirname(self.state_file), exist_ok=True)
+            with open(self.state_file, 'w') as f:
+                json.dump(state, f)
+        except Exception as e:
+            logger.error(f"Failed to save proxy state: {e}")
+    
+    async def add_vm(self, vm_id: str, vm_ip: str, port: Optional[int] = None) -> bool:
+        """Add proxy configuration for a new VM.
+        
+        Args:
+            vm_id: Unique identifier for the VM
+            vm_ip: IP address of the VM
+            port: Optional specific port to use, if not provided one will be allocated
+            
+        Returns:
+            True if proxy configuration was successful, False otherwise
+        """
+        try:
+            # Use provided port or allocate one
+            if port is None:
+                port = self.port_manager.allocate_port(vm_id)
+                if port is None:
+                    logger.error(f"Failed to allocate port for VM {vm_id}")
+                    return False
+            
+            # Create and start proxy server
+            proxy = ProxyServer(port, vm_ip)
+            await proxy.start()
+            
+            self._proxies[vm_id] = proxy
+            self._save_state()
+            
+            logger.info(f"Started proxy for VM {vm_id} on port {port}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to configure proxy for VM {vm_id}: {e}")
+            # Only deallocate if we allocated the port ourselves
+            if port is None and 'port' in locals():
+                self.port_manager.deallocate_port(vm_id)
+            return False
+    
+    async def remove_vm(self, vm_id: str) -> None:
+        """Remove proxy configuration for a VM.
+        
+        Args:
+            vm_id: Unique identifier for the VM
+        """
+        try:
+            if vm_id in self._proxies:
+                proxy = self._proxies.pop(vm_id)
+                await proxy.stop()
+                self.port_manager.deallocate_port(vm_id)
+                self._save_state()
+                logger.info(f"Removed proxy for VM {vm_id}")
+        except Exception as e:
+            logger.error(f"Failed to remove proxy for VM {vm_id}: {e}")
+    
+    def get_port(self, vm_id: str) -> Optional[int]:
+        """Get allocated port for a VM."""
+        return self.port_manager.get_port(vm_id)
+    
+    async def cleanup(self) -> None:
+        """Remove all proxy configurations."""
+        try:
+            for vm_id in list(self._proxies.keys()):
+                await self.remove_vm(vm_id)
+            self._save_state()
+            logger.info("Cleaned up all proxy configurations")
+        except Exception as e:
+            logger.error(f"Failed to cleanup proxy configurations: {e}")