request-vm-on-golem 0.1.0__py3-none-any.whl

requestor/config.py

@@ -0,0 +1,67 @@
+from pathlib import Path
+from typing import Optional, Dict
+from pydantic import BaseSettings, Field
+
+class RequestorConfig(BaseSettings):
+    """Configuration settings for the requestor node."""
+    
+    class Config:
+        env_prefix = "GOLEM_REQUESTOR_"
+    
+    # Environment
+    environment: str = Field(
+        default="development",
+        description="Environment mode: 'development' or 'production'"
+    )
+    
+    # Discovery Service
+    discovery_url: str = Field(
+        default="http://localhost:7465",
+        description="URL of the discovery service"
+    )
+    
+    # Base Directory
+    base_dir: Path = Field(
+        default_factory=lambda: Path.home() / ".golem",
+        description="Base directory for all Golem requestor files"
+    )
+    
+    # SSH Settings
+    ssh_key_dir: Path = Field(
+        default=None,
+        description="Directory for SSH keys. Defaults to {base_dir}/ssh"
+    )
+    
+    # Database Settings
+    db_path: Path = Field(
+        default=None,
+        description="Path to SQLite database. Defaults to {base_dir}/vms.db"
+    )
+
+    def __init__(self, **kwargs):
+        # Set dependent paths before validation
+        if 'ssh_key_dir' not in kwargs:
+            base_dir = kwargs.get('base_dir', Path.home() / ".golem")
+            kwargs['ssh_key_dir'] = base_dir / "ssh"
+        if 'db_path' not in kwargs:
+            base_dir = kwargs.get('base_dir', Path.home() / ".golem")
+            kwargs['db_path'] = base_dir / "vms.db"
+        super().__init__(**kwargs)
+
+    def get_provider_url(self, ip_address: str) -> str:
+        """Get provider API URL.
+        
+        Args:
+            ip_address: The IP address of the provider. In development mode, this is ignored
+                       and localhost is always used.
+        
+        Returns:
+            The complete provider URL with protocol and port.
+        """
+        # In development mode, always use localhost regardless of the ip_address parameter
+        if self.environment == "development":
+            return "http://localhost:7466"
+        # In production mode, use the actual provider IP address
+        return f"http://{ip_address}:7466"
+
+config = RequestorConfig()

requestor/db/__init__.py

@@ -0,0 +1,5 @@
+"""Database management module."""
+
+from .sqlite import Database
+
+__all__ = ['Database']

requestor/db/sqlite.py

@@ -0,0 +1,120 @@
+import aiosqlite
+from pathlib import Path
+from typing import Optional, Dict, List
+import json
+
+class Database:
+    def __init__(self, db_path: Path):
+        self.db_path = db_path
+        self.db_path.parent.mkdir(parents=True, exist_ok=True)
+
+    async def init(self):
+        """Initialize database and handle migrations."""
+        async with aiosqlite.connect(self.db_path) as db:
+            # Check if table exists
+            async with db.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='vms'") as cursor:
+                table_exists = await cursor.fetchone() is not None
+
+            if not table_exists:
+                # Create new table without ssh_key_name
+                await db.execute("""
+                    CREATE TABLE vms (
+                        name TEXT PRIMARY KEY,
+                        provider_ip TEXT NOT NULL,
+                        vm_id TEXT NOT NULL,
+                        config TEXT NOT NULL,
+                        status TEXT NOT NULL DEFAULT 'running',
+                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                    )
+                """)
+            else:
+                # Check if ssh_key_name column exists
+                async with db.execute("PRAGMA table_info(vms)") as cursor:
+                    columns = await cursor.fetchall()
+                    has_ssh_key_name = any(col[1] == 'ssh_key_name' for col in columns)
+
+                if has_ssh_key_name:
+                    # Migrate existing data
+                    await db.execute("""
+                        CREATE TABLE vms_new (
+                            name TEXT PRIMARY KEY,
+                            provider_ip TEXT NOT NULL,
+                            vm_id TEXT NOT NULL,
+                            config TEXT NOT NULL,
+                            status TEXT NOT NULL DEFAULT 'running',
+                            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+                        )
+                    """)
+                    await db.execute("""
+                        INSERT INTO vms_new (name, provider_ip, vm_id, config, status, created_at)
+                        SELECT name, provider_ip, vm_id, config, status, created_at FROM vms
+                    """)
+                    await db.execute("DROP TABLE vms")
+                    await db.execute("ALTER TABLE vms_new RENAME TO vms")
+
+            await db.commit()
+
+    async def save_vm(
+        self,
+        name: str,
+        provider_ip: str,
+        vm_id: str,
+        config: Dict,
+        status: str = 'running'
+    ) -> None:
+        """Save VM details."""
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute(
+                """
+                INSERT INTO vms (name, provider_ip, vm_id, config, status)
+                VALUES (?, ?, ?, ?, ?)
+                """,
+                (name, provider_ip, vm_id, json.dumps(config), status)
+            )
+            await db.commit()
+
+    async def get_vm(self, name: str) -> Optional[Dict]:
+        """Get VM details."""
+        async with aiosqlite.connect(self.db_path) as db:
+            db.row_factory = aiosqlite.Row
+            async with db.execute(
+                "SELECT * FROM vms WHERE name = ?",
+                (name,)
+            ) as cursor:
+                row = await cursor.fetchone()
+                if row:
+                    vm = dict(row)
+                    vm['config'] = json.loads(vm['config'])
+                    return vm
+                return None
+
+    async def delete_vm(self, name: str) -> None:
+        """Delete VM details."""
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute(
+                "DELETE FROM vms WHERE name = ?",
+                (name,)
+            )
+            await db.commit()
+
+    async def update_vm_status(self, name: str, status: str) -> None:
+        """Update VM status."""
+        async with aiosqlite.connect(self.db_path) as db:
+            await db.execute(
+                "UPDATE vms SET status = ? WHERE name = ?",
+                (status, name)
+            )
+            await db.commit()
+
+    async def list_vms(self) -> List[Dict]:
+        """List all VMs."""
+        async with aiosqlite.connect(self.db_path) as db:
+            db.row_factory = aiosqlite.Row
+            async with db.execute("SELECT * FROM vms") as cursor:
+                rows = await cursor.fetchall()
+                vms = []
+                for row in rows:
+                    vm = dict(row)
+                    vm['config'] = json.loads(vm['config'])
+                    vms.append(vm)
+                return vms

requestor/errors.py

@@ -0,0 +1,29 @@
+class RequestorError(Exception):
+    """Base class for requestor errors."""
+    pass
+
+class ProviderError(RequestorError):
+    """Provider communication error."""
+    pass
+
+class DiscoveryError(RequestorError):
+    """Discovery service error."""
+    pass
+
+class SSHError(RequestorError):
+    """SSH-related error."""
+    pass
+
+class ConfigError(RequestorError):
+    """Configuration error."""
+    pass
+
+class DatabaseError(RequestorError):
+    """Database operation error."""
+    pass
+
+class VMError(RequestorError):
+    """VM operation error."""
+    def __init__(self, message: str, vm_id: str = None):
+        self.vm_id = vm_id
+        super().__init__(message)

requestor/provider/__init__.py

@@ -0,0 +1,5 @@
+"""Provider communication module."""
+
+from .client import ProviderClient
+
+__all__ = ['ProviderClient']

requestor/provider/client.py

@@ -0,0 +1,94 @@
+import aiohttp
+from typing import Dict, Optional
+from pathlib import Path
+
+class ProviderClient:
+    def __init__(self, provider_url: str):
+        self.provider_url = provider_url
+        self.session = None
+
+    async def __aenter__(self):
+        self.session = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        if self.session:
+            await self.session.close()
+
+    async def create_vm(
+        self,
+        name: str,
+        cpu: int,
+        memory: int,
+        storage: int,
+        ssh_key: str
+    ) -> Dict:
+        """Create a VM on the provider."""
+        async with self.session.post(
+            f"{self.provider_url}/api/v1/vms",
+            json={
+                "name": name,
+                "resources": {
+                    "cpu": cpu,
+                    "memory": memory,
+                    "storage": storage
+                },
+                "ssh_key": ssh_key
+            }
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise Exception(f"Failed to create VM: {error_text}")
+            return await response.json()
+
+    async def add_ssh_key(self, vm_id: str, key: str) -> None:
+        """Add SSH key to VM."""
+        async with self.session.post(
+            f"{self.provider_url}/api/v1/vms/{vm_id}/ssh-keys",
+            json={
+                "key": key,
+                "name": "default"
+            }
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise Exception(f"Failed to add SSH key: {error_text}")
+
+    async def start_vm(self, vm_id: str) -> Dict:
+        """Start a VM."""
+        async with self.session.post(
+            f"{self.provider_url}/api/v1/vms/{vm_id}/start"
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise Exception(f"Failed to start VM: {error_text}")
+            return await response.json()
+
+    async def stop_vm(self, vm_id: str) -> Dict:
+        """Stop a VM."""
+        async with self.session.post(
+            f"{self.provider_url}/api/v1/vms/{vm_id}/stop"
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise Exception(f"Failed to stop VM: {error_text}")
+            return await response.json()
+
+    async def destroy_vm(self, vm_id: str) -> None:
+        """Destroy a VM."""
+        async with self.session.delete(
+            f"{self.provider_url}/api/v1/vms/{vm_id}"
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise Exception(f"Failed to destroy VM: {error_text}")
+
+    async def get_vm_access(self, vm_id: str) -> Dict:
+        """Get VM access information."""
+        async with self.session.get(
+            f"{self.provider_url}/api/v1/vms/{vm_id}/access"
+        ) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise Exception(f"Failed to get VM access info: {error_text}")
+            return await response.json()

requestor/run.py

@@ -0,0 +1,51 @@
+#!/usr/bin/env python3
+import os
+import sys
+from pathlib import Path
+from dotenv import load_dotenv
+
+from requestor.utils.logging import setup_logger
+from requestor.cli.commands import cli
+
+# Configure logging with debug mode from environment variable
+logger = setup_logger(__name__)
+
+def check_requirements():
+    """Check if all requirements are met."""
+    # Check required directories
+    ssh_key_dir = os.environ.get(
+        'GOLEM_REQUESTOR_SSH_KEY_DIR',
+        str(Path.home() / '.golem' / 'requestor' / 'ssh')
+    )
+    
+    try:
+        # Create and secure directories
+        path = Path(ssh_key_dir)
+        path.mkdir(parents=True, exist_ok=True)
+        path.chmod(0o700)  # Secure permissions for SSH keys
+    except Exception as e:
+        logger.error(f"Failed to create required directories: {e}")
+        return False
+        
+    return True
+
+def main():
+    """Run the requestor CLI."""
+    try:
+        # Load environment variables from .env file
+        env_path = Path(__file__).parent / '.env'
+        load_dotenv(dotenv_path=env_path)
+
+        # Check requirements
+        if not check_requirements():
+            logger.error("Requirements check failed")
+            sys.exit(1)
+
+        # Run CLI
+        cli()
+    except Exception as e:
+        logger.error(f"Failed to start requestor CLI: {e}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

requestor/ssh/__init__.py

@@ -0,0 +1,4 @@
+"""SSH management module for VM on Golem requestor."""
+from .manager import SSHKeyManager
+
+__all__ = ['SSHKeyManager']

requestor/ssh/manager.py

@@ -0,0 +1,161 @@
+"""SSH key management for VM on Golem requestor."""
+import os
+import asyncio
+import logging
+import sys
+from pathlib import Path
+from typing import Tuple, Optional, Union, NamedTuple
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from cryptography.hazmat.backends import default_backend
+
+# Configure basic logging
+logging.basicConfig(
+    level=logging.DEBUG,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    stream=sys.stderr
+)
+logger = logging.getLogger(__name__)
+
+class KeyPair(NamedTuple):
+    """Represents an SSH key pair with both private and public keys.
+    
+    Attributes:
+        private_key: Path to the private key file
+        public_key: Path to the public key file
+        private_key_content: Content of the private key file
+        public_key_content: Content of the public key file
+    """
+    private_key: Path
+    public_key: Path
+    private_key_content: str
+    public_key_content: str
+
+class SSHKeyManager:
+    """Manages SSH keys for VM connections."""
+    
+    def __init__(self, golem_dir: Union[str, Path] = None):
+        """Initialize SSH key manager.
+        
+        Args:
+            golem_dir: Optional custom directory for Golem SSH keys. If None, uses default from config.
+        """
+        # Set up Golem SSH directory
+        if golem_dir is None:
+            from ..config import config
+            self.golem_dir = Path(config.ssh_key_dir)
+        else:
+            self.golem_dir = Path(golem_dir)
+
+        # Define key paths
+        self.system_key_path = Path.home() / '.ssh' / 'id_rsa'
+        self.golem_key_path = self.golem_dir / 'golem_id_rsa'  # Single reusable key
+        
+        # Create Golem directory if needed
+        self.golem_dir.mkdir(parents=True, exist_ok=True)
+        # Secure directory permissions (on Unix-like systems)
+        if os.name == 'posix':
+            os.chmod(self.golem_dir, 0o700)
+
+    async def get_key_pair(self) -> KeyPair:
+        """Get the SSH key pair to use.
+        
+        Returns system SSH key if available, otherwise returns/creates Golem key.
+        """
+        # Try to use system SSH key first
+        logger.debug("Checking for system SSH key at %s", self.system_key_path)
+        system_pub_key = self.system_key_path.parent / 'id_rsa.pub'
+        
+        if self.system_key_path.exists() and system_pub_key.exists():
+            logger.info("Using existing system SSH key")
+            try:
+                return KeyPair(
+                    private_key=self.system_key_path,
+                    public_key=system_pub_key,
+                    private_key_content=self.system_key_path.read_text().strip(),
+                    public_key_content=system_pub_key.read_text().strip()
+                )
+            except (PermissionError, OSError) as e:
+                logger.warning("Could not read system SSH key: %s", e)
+                # Fall through to use Golem key
+        
+        # Use/create Golem key if system key unavailable
+        logger.debug("Using Golem SSH key at %s", self.golem_key_path)
+        if not self.golem_key_path.exists():
+            logger.info("No existing Golem SSH key found, generating new key pair")
+            await self._generate_key_pair()
+        
+        golem_pub_key = Path(str(self.golem_key_path) + '.pub')
+        return KeyPair(
+            private_key=self.golem_key_path,
+            public_key=golem_pub_key,
+            private_key_content=self.golem_key_path.read_text().strip(),
+            public_key_content=golem_pub_key.read_text().strip()
+        )
+
+    async def get_public_key_content(self) -> str:
+        """Get the content of the public key file."""
+        key_pair = await self.get_key_pair()
+        return key_pair.public_key_content
+
+    async def get_key_content(self) -> KeyPair:
+        """Get both the paths and contents of the key pair."""
+        logger.debug("Getting key content")
+        key_pair = await self.get_key_pair()
+        logger.debug("Got key pair with paths: private=%s, public=%s", 
+                    key_pair.private_key, key_pair.public_key)
+        return key_pair
+
+    @classmethod
+    async def generate_key_pair(cls) -> KeyPair:
+        """Generate a new RSA key pair for Golem VMs and return their contents."""
+        logger.info("Generating new SSH key pair")
+        manager = cls()
+        await manager._generate_key_pair()
+        logger.debug("Key pair generated, getting content")
+        return await manager.get_key_content()
+
+    async def _generate_key_pair(self):
+        """Generate a new RSA key pair for Golem VMs."""
+        logger.debug("Generating new RSA key pair")
+        try:
+            # Generate private key
+            private_key = rsa.generate_private_key(
+                public_exponent=65537,
+                key_size=2048,
+                backend=default_backend()
+            )
+            logger.debug("Generated private key")
+
+            # Save private key
+            private_pem = private_key.private_bytes(
+                encoding=serialization.Encoding.PEM,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption()
+            )
+            logger.debug("Saving private key to %s", self.golem_key_path)
+            self.golem_key_path.write_bytes(private_pem)
+            if os.name == 'posix':
+                os.chmod(self.golem_key_path, 0o600)  # Secure key permissions on Unix-like systems
+
+            # Save public key
+            logger.debug("Generating public key")
+            public_key = private_key.public_key()
+            public_pem = public_key.public_bytes(
+                encoding=serialization.Encoding.OpenSSH,
+                format=serialization.PublicFormat.OpenSSH
+            )
+            pub_key_path = Path(str(self.golem_key_path) + '.pub')
+            logger.debug("Saving public key to %s", pub_key_path)
+            pub_key_path.write_bytes(public_pem)
+            if os.name == 'posix':
+                os.chmod(pub_key_path, 0o644)  # Public key can be readable on Unix-like systems
+            logger.info("Successfully generated and saved SSH key pair")
+        except Exception as e:
+            logger.error("Failed to generate key pair: %s", str(e))
+            raise
+
+    async def get_private_key_content(self, force_golem_key: bool = False) -> Optional[str]:
+        """Get the content of the private key file."""
+        key_pair = await self.get_key_pair(force_golem_key)
+        return key_pair.private_key_content

requestor/utils/logging.py

@@ -0,0 +1,109 @@
+import logging
+import colorlog
+import sys
+import os
+from typing import Optional
+from enum import Enum
+
+# Import standard logging levels
+from logging import DEBUG, INFO, WARNING, ERROR, CRITICAL
+
+class LogLevel(Enum):
+    """Custom log levels for enhanced CLI experience."""
+    COMMAND = 15  # For CLI commands
+    PROCESS = 25  # For ongoing processes
+    SUCCESS = 35  # For successful operations
+    DETAIL = 22   # For additional details
+
+# Add custom levels to logging
+logging.addLevelName(LogLevel.COMMAND.value, 'COMMAND')
+logging.addLevelName(LogLevel.PROCESS.value, 'PROCESS')
+logging.addLevelName(LogLevel.SUCCESS.value, 'SUCCESS')
+logging.addLevelName(LogLevel.DETAIL.value, 'DETAIL')
+
+def command(self, message, *args, **kwargs):
+    """Log CLI command with a distinctive style."""
+    if self.isEnabledFor(LogLevel.COMMAND.value):
+        self._log(LogLevel.COMMAND.value, message, args, **kwargs)
+
+def process(self, message, *args, **kwargs):
+    """Log ongoing process with a progress indicator."""
+    if self.isEnabledFor(LogLevel.PROCESS.value):
+        self._log(LogLevel.PROCESS.value, f"⚡ {message}", args, **kwargs)
+
+def success(self, message, *args, **kwargs):
+    """Log successful operation with a checkmark."""
+    if self.isEnabledFor(LogLevel.SUCCESS.value):
+        self._log(LogLevel.SUCCESS.value, f"✨ {message}", args, **kwargs)
+
+def detail(self, message, *args, **kwargs):
+    """Log additional details with an arrow."""
+    if self.isEnabledFor(LogLevel.DETAIL.value):
+        self._log(LogLevel.DETAIL.value, f"  → {message}", args, **kwargs)
+
+# Add methods to Logger class
+logging.Logger.command = command
+logging.Logger.process = process
+logging.Logger.success = success
+logging.Logger.detail = detail
+
+def setup_logger(name: Optional[str] = None) -> logging.Logger:
+    """Setup and return a colored logger optimized for CLI experience.
+    
+    Args:
+        name: Logger name (optional)
+        
+    Returns:
+        Configured logger instance with fancy formatting
+    """
+    logger = logging.getLogger(name or __name__)
+    logger.handlers = []  # Clear existing handlers
+    
+    # Check DEBUG environment variable
+    debug = os.getenv('DEBUG', '').lower() in ('1', 'true', 'yes')
+    
+    # Prevent duplicate logs by removing root handlers
+    root = logging.getLogger()
+    root.handlers = []
+    
+    # Fancy handler for CLI output
+    fancy_handler = colorlog.StreamHandler(sys.stderr)  # Use stderr for logs
+    fancy_formatter = colorlog.ColoredFormatter(
+        "%(log_color)s%(message)s%(reset)s",
+        reset=True,
+        log_colors={
+            'DEBUG': 'blue',
+            'INFO': 'white',
+            'COMMAND': 'bold_cyan',
+            'DETAIL': 'cyan',
+            'PROCESS': 'yellow',
+            'WARNING': 'yellow',
+            'SUCCESS': 'bold_green',
+            'ERROR': 'bold_red',
+            'CRITICAL': 'bold_red,bg_white',
+        },
+        secondary_log_colors={},
+        style='%'
+    )
+    fancy_handler.setFormatter(fancy_formatter)
+    fancy_handler.addFilter(
+        lambda record: record.levelno != DEBUG or debug
+    )
+    logger.addHandler(fancy_handler)
+    logger.propagate = False  # Prevent propagation to avoid duplicates
+    
+    if debug:
+        logger.setLevel(DEBUG)
+        # Enable debug logging for other libraries
+        logging.getLogger('asyncio').setLevel(DEBUG)
+        logging.getLogger('aiosqlite').setLevel(DEBUG)
+    else:
+        logger.setLevel(INFO)
+        # Suppress debug logs from other libraries
+        logging.getLogger('asyncio').setLevel(WARNING)
+        logging.getLogger('aiosqlite').setLevel(WARNING)
+    
+    return logger
+
+# Create default logger
+logger = setup_logger('golem.requestor')