xenfra 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

xenfra/engine.py

@@ -1,264 +1,292 @@
-import digitalocean
+# src/xenfra/engine.py
+
 import os
 import time
-from dotenv import load_dotenv
 from pathlib import Path
-from fabric import Connection
+import digitalocean
+import fabric
+from dotenv import load_dotenv
+from sqlmodel import Session
 
-# Internal imports
-from .utils import get_project_context
-from .recipes import generate_stack
-from .dockerizer import detect_framework, generate_deployment_assets
+# Xenfra modules
+from . import dockerizer
+from . import recipes
+from .db.models import Project
+from .db.session import get_session
 
+class DeploymentError(Exception):
+    """Custom exception for deployment failures."""
+    def __init__(self, message, stage="Unknown"):
+        self.message = message
+        self.stage = stage
+        super().__init__(f"Deployment failed at stage '{stage}': {message}")
 
 class InfraEngine:
     """
-    The Core SDK.
-    This class handles all communication with DigitalOcean.
-    It is designed to be used by a CLI or an AI Agent.
+    The InfraEngine is the core of Xenfra. It handles all interactions
+    with the cloud provider and orchestrates the deployment lifecycle.
     """
-
-    def __init__(self):
-        # Load secrets immediately
+    def __init__(self, token: str = None, db_session: Session = None):
+        """
+        Initializes the engine and validates the API token.
+        """
         load_dotenv()
-        self.token = os.getenv("DIGITAL_OCEAN_TOKEN")
-        
-        if not self.token:
-            raise ValueError("❌ FATAL: No DIGITAL_OCEAN_TOKEN found in .env or environment.")
-            
-        self.manager = digitalocean.Manager(token=self.token)
+        self.token = token or os.getenv("DIGITAL_OCEAN_TOKEN")
+        self.db_session = db_session or next(get_session())
 
-    def get_user_info(self):
-        """Returns the current user's email and account status."""
+        if not self.token:
+            raise ValueError("DigitalOcean API token not found. Please set the DIGITAL_OCEAN_TOKEN environment variable.")
         try:
-            account = self.manager.get_account()
-            return {"email": account.email, "status": account.status, "limit": account.droplet_limit}
+            self.manager = digitalocean.Manager(token=self.token)
+            self.get_user_info()
         except Exception as e:
-            raise ConnectionError(f"Could not connect to DigitalOcean: {e}")
-
-    def _ensure_ssh_key(self):
-        """
-        Private method.
-        1. Dynamically finds local public SSH keys in ~/.ssh/.
-        2. Checks if any of them exist on DigitalOcean.
-        3. If not, uploads the first one found.
-        Returns: The SSHKey object.
-        """
-        # 1. Find Local Keys
-        ssh_dir = Path.home() / ".ssh"
-        pub_keys = list(ssh_dir.glob("*.pub"))
-
-        if not pub_keys:
-            raise FileNotFoundError(f"No public SSH keys found in {ssh_dir}. Run 'ssh-keygen' first.")
-
-        # 2. Check Remote Keys against Local Keys
-        remote_keys = self.manager.get_all_sshkeys()
-        remote_key_map = {k.public_key.strip(): k for k in remote_keys}
-
-        first_local_key_content = None
-        for key_path in pub_keys:
-            local_pub_key = key_path.read_text().strip()
-            if not first_local_key_content:
-                first_local_key_content = local_pub_key
-            
-            if local_pub_key in remote_key_map:
-                print(f"✨ Found existing SSH key on DigitalOcean: {key_path.name}")
-                return remote_key_map[local_pub_key]
-
-        # 3. If no matches, upload the first key found
-        if not first_local_key_content:
-             raise FileNotFoundError(f"Could not read any public keys in {ssh_dir}.")
-
-        print(f"✨ No matching SSH key found on DigitalOcean. Uploading new key from {pub_keys[0].name}...")
-        new_key = digitalocean.SSHKey(token=self.token,
-                                      name=f"xenfra-{pub_keys[0].name}-{int(time.time())}",
-                                      public_key=first_local_key_content)
-        new_key.create()
-        return new_key
+            raise ConnectionError(f"Failed to connect to DigitalOcean: {e}")
 
     def _get_connection(self, ip_address: str):
-        """Creates a standardized Fabric connection."""
-        return Connection(
+        """Establishes a Fabric connection to the server."""
+        private_key_path = str(Path.home() / ".ssh" / "id_rsa")
+        if not Path(private_key_path).exists():
+            raise DeploymentError("No private SSH key found at ~/.ssh/id_rsa.", stage="Setup")
+
+        return fabric.Connection(
             host=ip_address,
             user="root",
-            connect_kwargs={
-                "timeout": 10
-            }
+            connect_kwargs={"key_filename": [private_key_path]},
         )
 
-    def _is_setup_complete(self, ip_address: str) -> bool:
-        """Checks for the setup completion marker file on the remote server."""
-        try:
-            c = self._get_connection(ip_address)
-            if c.run("test -f /root/setup_complete", warn=True, hide=True).ok:
-                return True
-        except Exception:
-            # This can happen if the server is not ready for SSH
-            return False
-        return False
-
-    def _auto_heal_server(self, ip_address: str):
-        """
-        Diagnoses issues and attempts to fix them remotely.
-        This is the 'nuclear option' for a stalled setup.
-        """
-        print(f"\n[ENGINE] 🚑 INITIATING AUTO-HEAL PROTOCOL for {ip_address}...")
-        c = self._get_connection(ip_address)
-
-        # A. STOP SERVICES (Prevent Respawn)
-        print("[ENGINE] 🛑 Stopping background update services...")
-        c.run("systemctl stop unattended-upgrades.service", warn=True, hide=True)
-        c.run("systemctl stop apt-daily.service", warn=True, hide=True)
-        c.run("systemctl stop apt-daily-upgrade.service", warn=True, hide=True)
-
-        # B. KILL PROCESSES
-        print("[ENGINE] 🔪 Killing stuck processes...")
-        c.run("pkill apt", warn=True, hide=True)
-        c.run("pkill apt-get", warn=True, hide=True)
-        c.run("pkill dpkg", warn=True, hide=True)
-
-        # C. REMOVE LOCKS
-        print("[ENGINE] PWN: Removing lock files...")
-        c.run("rm -f /var/lib/dpkg/lock*", warn=True, hide=True)
-        c.run("rm -f /var/lib/apt/lists/lock", warn=True, hide=True)
-        c.run("rm -f /var/cache/apt/archives/lock", warn=True, hide=True)
-
-        # D. REPAIR DPKG
-        print("[ENGINE] 🔧 Repairing package database...")
-        c.run("dpkg --configure -a", warn=True, hide=True)
-
-        # E. RE-MARK AS COMPLETE (to unblock the health check)
-        c.run("touch /root/setup_complete", warn=True, hide=True)
-        print("[ENGINE] ✅ Auto-Heal Complete.")
-
-    def upload_code(self, ip_address: str) -> bool:
-        """Uses Fabric to upload project files to the /root/app directory."""
-        print(f"\n[ENGINE] 🚁 INITIATING CODE AIRLIFT to {ip_address}...")
-        
-        ignored = {'.git', '.venv', 'venv', '__pycache__', '.pytest_cache', '.DS_Store', '.env'}
+    def get_user_info(self):
+        """Retrieves user account information."""
+        return self.manager.get_account()
+
+    def list_servers(self):
+        """Retrieves a list of all Droplets."""
+        return self.manager.get_all_droplets()
+
+    def destroy_server(self, droplet_id: int, db_session: Session = None):
+        """Destroys a Droplet by its ID and removes it from the local DB."""
+        session = db_session or self.db_session
         
-        try:
-            c = self._get_connection(ip_address)
-            
-            print("[ENGINE] 📤 Uploading Project Files...")
-            count = 0
-            for root, dirs, files in os.walk("."):
-                dirs[:] = [d for d in dirs if d not in ignored]
-                
-                for file in files:
-                    if file.endswith(".pyc"): continue
-                    
-                    local_path = os.path.join(root, file)
-                    rel_path = os.path.relpath(local_path, ".")
-                    remote_path = f"/root/app/{rel_path}"
-                    
-                    remote_dir = os.path.dirname(remote_path)
-                    c.run(f"mkdir -p {remote_dir}", hide=True)
-                    
-                    c.put(local_path, remote_path)
-                    count += 1
+        # Find the project in the local DB
+        project_to_delete = session.query(Project).filter(Project.droplet_id == droplet_id).first()
+
+        # Destroy the droplet on DigitalOcean
+        droplet = digitalocean.Droplet(token=self.token, id=droplet_id)
+        droplet.destroy()
 
-            print(f"[ENGINE] ✅ Airlift Complete: {count} files transferred.")
-            return True
+        # If it was in our DB, delete it
+        if project_to_delete:
+            session.delete(project_to_delete)
+            session.commit()
 
-        except Exception as e:
-            print(f"[ENGINE] ❌ Upload failed: {e}")
-            print("[ENGINE] 🔑 Ensure your SSH Agent is running and keys are added.")
-            return False
+    def list_projects_from_db(self, db_session: Session = None):
+        """Lists all projects from the local database."""
+        session = db_session or self.db_session
+        return session.query(Project).all()
 
-    def list_servers(self):
-        """Returns a list of all active servers."""
-        return self.manager.get_all_droplets()
+    def sync_with_provider(self, db_session: Session = None):
+        """Reconciles the local database with the live state from DigitalOcean."""
+        session = db_session or self.db_session
+        
+        # 1. Get live and local states
+        live_droplets = self.manager.get_all_droplets(tag_name="xenfra")
+        local_projects = self.list_projects_from_db(session)
+
+        live_map = {d.id: d for d in live_droplets}
+        local_map = {p.droplet_id: p for p in local_projects}
+
+        # 2. Reconcile
+        # Add new servers found on DO to our DB
+        for droplet_id, droplet in live_map.items():
+            if droplet_id not in local_map:
+                new_project = Project(
+                    droplet_id=droplet.id, name=droplet.name, ip_address=droplet.ip_address,
+                    status=droplet.status, region=droplet.region['slug'], size=droplet.size_slug
+                )
+                session.add(new_project)
+
+        # Remove servers from our DB that no longer exist on DO
+        for project_id, project in local_map.items():
+            if project_id not in live_map:
+                session.delete(project)
+        
+        session.commit()
+        return self.list_projects_from_db(session)
 
-    def deploy_server(self, name, region="blr1", size="s-1vcpu-1gb", image="ubuntu-24-04-x64", logger=None):
+    def stream_logs(self, droplet_id: int, db_session: Session = None):
         """
-        Provisions a new server, runs setup, and deploys the application.
-        This is a stateful, blocking method that orchestrates the entire process.
+        Verifies a server exists and streams its logs in real-time.
         """
-        def log(msg):
-            if logger:
-                logger(f"[ENGINE] {msg}")
+        session = db_session or self.db_session
+        
+        # 1. Find project in local DB
+        project = session.query(Project).filter(Project.droplet_id == droplet_id).first()
+        if not project:
+            raise DeploymentError(f"Project with Droplet ID {droplet_id} not found in local database.", stage="Log Streaming")
+
+        # 2. Just-in-Time Verification
+        try:
+            droplet = self.manager.get_droplet(droplet_id)
+        except digitalocean.baseapi.DataReadError as e:
+            if e.response.status_code == 404:
+                # The droplet doesn't exist, so remove it from our DB
+                session.delete(project)
+                session.commit()
+                raise DeploymentError(f"Server '{project.name}' (ID: {droplet_id}) no longer exists on DigitalOcean. It has been removed from your local list.", stage="Log Streaming")
             else:
-                print(f"[ENGINE] {msg}")
+                raise e
+        
+        # 3. Stream logs
+        ip_address = droplet.ip_address
+        with self._get_connection(ip_address) as conn:
+            conn.run("cd /root/app && docker-compose logs -f app", pty=True)
+
 
-        log(f"🚀 Starting deployment for '{name}'...")
 
+    def _ensure_ssh_key(self, logger):
+        """Ensures a local public SSH key is on DigitalOcean."""
+        pub_key_path = Path.home() / ".ssh" / "id_rsa.pub"
+        if not pub_key_path.exists():
+            raise DeploymentError("No SSH key found at ~/.ssh/id_rsa.pub. Please generate one.", stage="Setup")
+        
+        with open(pub_key_path) as f:
+            pub_key_content = f.read()
+
+        existing_keys = self.manager.get_all_sshkeys()
+        for key in existing_keys:
+            if key.public_key.strip() == pub_key_content.strip():
+                logger("   - Found existing SSH key on DigitalOcean.")
+                return key
+
+        logger("   - No matching SSH key found. Creating a new one on DigitalOcean...")
+        key = digitalocean.SSHKey(token=self.token, name="xenfra-cli-key", public_key=pub_key_content)
+        key.create()
+        return key
+
+    def deploy_server(self, name: str, region: str, size: str, image: str, logger: callable, db_session: Session = None, **kwargs):
+        """A stateful, blocking orchestrator for deploying a new server."""
+        droplet = None
+        session = db_session or self.db_session
         try:
-            # 1. Project Context & Asset Generation
-            log("🔎 Detecting project context...")
-            context = get_project_context()
-            framework, _, _ = detect_framework()
-            is_dockerized = False
-            if framework:
-                log(f"✨ Web framework '{framework}' detected.")
-                generate_deployment_assets(context)
-                is_dockerized = True
-
-            # 2. Get SSH Key
-            ssh_key = self._ensure_ssh_key()
-
-            # 3. Generate Setup Script
-            log("📦 Generating cloud-init setup script...")
-            setup_script = generate_stack(context, is_dockerized=is_dockerized)
-
-            # 4. Create Droplet
-            log(f"☁️  Provisioning droplet '{name}'...")
-            droplet = digitalocean.Droplet(token=self.token,
-                                           name=name,
-                                           region=region,
-                                           image=image,
-                                           size_slug=size,
-                                           ssh_keys=[ssh_key],
-                                           user_data=setup_script)
+            # === 1. SETUP STAGE ===
+            logger("\n[bold blue]PHASE 1: SETUP[/bold blue]")
+            ssh_key = self._ensure_ssh_key(logger)
+
+            # === 2. ASSET GENERATION STAGE ===
+            logger("\n[bold blue]PHASE 2: GENERATING DEPLOYMENT ASSETS[/bold blue]")
+            context = {**kwargs} # Pass db config, etc.
+            files = dockerizer.generate_templated_assets(context)
+            for file in files:
+                logger(f"   - Generated {file}")
+
+            # === 3. CLOUD-INIT STAGE ===
+            logger("\n[bold blue]PHASE 3: CREATING SERVER SETUP SCRIPT[/bold blue]")
+            cloud_init_script = recipes.generate_stack(context)
+            logger("   - Generated cloud-init script.")
+
+            # === 4. DROPLET CREATION STAGE ===
+            logger("\n[bold blue]PHASE 4: PROVISIONING SERVER[/bold blue]")
+            droplet = digitalocean.Droplet(
+                token=self.token, name=name, region=region, image=image, size_slug=size,
+                ssh_keys=[ssh_key], userdata=cloud_init_script, tags=["xenfra"]
+            )
             droplet.create()
+            logger(f"   - Droplet '{name}' creation initiated (ID: {droplet.id}). Waiting for it to become active...")
 
-            # 5. Wait for Droplet to become active and get IP
-            log("⏳ Waiting for droplet to become active and get IP address...")
-            while not droplet.ip_address:
+            # === 5. POLLING STAGE ===
+            logger("\n[bold blue]PHASE 5: WAITING FOR SERVER SETUP[/bold blue]")
+            while True:
                 droplet.load()
-                time.sleep(5)
-            log(f"🌍 Droplet is active at {droplet.ip_address}")
-
-            # 6. Health Check & Auto-Heal Loop
-            log("🩺 Performing health checks on the new server...")
-            attempts = 0
-            max_retries = 18 # ~90 seconds
-            is_healthy = False
-            while attempts < max_retries:
-                if self._is_setup_complete(droplet.ip_address):
-                    log("✅ Server setup complete and healthy.")
-                    is_healthy = True
+                if droplet.status == 'active':
+                    logger("   - Droplet is active. Waiting for SSH to be available...")
+                    break
+                time.sleep(10)
+
+            ip_address = droplet.ip_address
+
+            # Retry SSH connection
+            conn = None
+            max_retries = 12 # 2-minute timeout for SSH
+            for i in range(max_retries):
+                try:
+                    logger(f"   - Attempting SSH connection ({i+1}/{max_retries})...")
+                    conn = self._get_connection(ip_address)
+                    conn.open() # Explicitly open the connection
+                    logger("   - SSH connection established.")
                     break
-                log(f"   - Still waiting for setup to complete... (attempt {attempts+1}/{max_retries})")
+                except Exception as e:
+                    if i < max_retries - 1:
+                        logger(f"   - SSH connection failed. Retrying in 10s...")
+                        time.sleep(10)
+                    else:
+                        raise DeploymentError(f"Failed to establish SSH connection: {e}", stage="Polling")
+
+            if not conn or not conn.is_connected:
+                raise DeploymentError("Could not establish SSH connection.", stage="Polling")
+
+            with conn:
+                for i in range(30): # 5-minute timeout for cloud-init
+                    if conn.run("test -f /root/setup_complete", warn=True).ok:
+                        logger("   - Cloud-init setup complete.")
+                        break
+                    time.sleep(10)
+                else:
+                    raise DeploymentError("Server setup script failed to complete in time.", stage="Polling")
+
+            # === 6. CODE UPLOAD STAGE ===
+            logger("\n[bold blue]PHASE 6: UPLOADING APPLICATION CODE[/bold blue]")
+            with self._get_connection(ip_address) as conn:
+                fabric.transfer.Transfer(conn).upload(".", "/root/app", exclude=[".git", ".venv", "__pycache__"])
+            logger("   - Code upload complete.")
+
+            # === 7. FINAL DEPLOY STAGE ===
+            logger("\n[bold blue]PHASE 7: BUILDING AND DEPLOYING CONTAINERS[/bold blue]")
+            with self._get_connection(ip_address) as conn:
+                result = conn.run("cd /root/app && docker-compose up -d --build", hide=True)
+                if result.failed:
+                    raise DeploymentError(f"docker-compose failed: {result.stderr}", stage="Deploy")
+            logger("   - Docker containers are building in the background...")
+
+            # === 8. VERIFICATION STAGE ===
+            logger("\n[bold blue]PHASE 8: VERIFYING DEPLOYMENT[/bold blue]")
+            app_port = context.get('port', 8000)
+            for i in range(24): # 2-minute timeout for health checks
+                logger(f"   - Health check attempt {i+1}/24...")
+                with self._get_connection(ip_address) as conn:
+                    # Check if container is running
+                    ps_result = conn.run("cd /root/app && docker-compose ps", hide=True)
+                    if "running" not in ps_result.stdout:
+                        time.sleep(5)
+                        continue
+                    
+                    # Check if application is responsive
+                    curl_result = conn.run(f"curl -s --fail http://localhost:{app_port}/", warn=True)
+                    if curl_result.ok:
+                        logger("[bold green]   - Health check passed! Application is live.[/bold green]")
+                        
+                        # === 9. PERSISTENCE STAGE ===
+                        logger("\n[bold blue]PHASE 9: SAVING DEPLOYMENT TO DATABASE[/bold blue]")
+                        project = Project(
+                            droplet_id=droplet.id,
+                            name=droplet.name,
+                            ip_address=ip_address,
+                            status=droplet.status,
+                            region=droplet.region['slug'],
+                            size=droplet.size_slug,
+                        )
+                        session.add(project)
+                        session.commit()
+                        logger("   - Deployment saved.")
+
+                        return droplet # Return the full droplet object
                 time.sleep(5)
-                attempts += 1
-            
-            if not is_healthy:
-                log("⚠️ Server setup stalled. Initiating auto-heal protocol...")
-                self._auto_heal_server(droplet.ip_address)
-                if not self._is_setup_complete(droplet.ip_address):
-                     raise Exception("Deployment failed: Auto-heal could not recover the server.")
-                log("✅ Server recovered successfully after auto-heal.")
-
-            # 7. Upload Code
-            self.upload_code(droplet.ip_address)
-
-            # 8. Run Docker Compose if applicable
-            if is_dockerized:
-                log("🚀 Starting application with Docker Compose...")
-                c = self._get_connection(droplet.ip_address)
-                c.run("cd /root/app && docker compose up -d", hide=True)
-
-            log("🎉 Deployment successful!")
-            return {"status": "success", "ip": droplet.ip_address, "name": name}
+            else:
+                # On failure, get logs and destroy droplet
+                with self._get_connection(ip_address) as conn:
+                    logs = conn.run("cd /root/app && docker-compose logs", hide=True).stdout
+                raise DeploymentError(f"Application failed to become healthy in time. Logs:\n{logs}", stage="Verification")
 
         except Exception as e:
-            log(f"❌ CRITICAL DEPLOYMENT FAILURE: {e}")
-            # Optionally, trigger a cleanup action here
-            raise e
-
-    def destroy_server(self, droplet_id):
-        """Destroys a server by ID. IRREVERSIBLE."""
-        droplet = self.manager.get_droplet(droplet_id)
-        droplet.destroy()
-        return True
+            if droplet:
+                logger(f"[bold red]Deployment failed. Cleaning up created server '{droplet.name}'...[/bold red]")
+                droplet.destroy()
+            raise e