"""Async subprocess wrapper for Ansible playbook execution.""" import asyncio import json import os import re from collections.abc import AsyncGenerator from pathlib import Path from typing import Any from .sessions import DeploymentStatus, Session, sessions # Get the algo project root directory ALGO_DIR = Path(__file__).parent.parent def build_extra_vars(session: Session) -> dict[str, Any]: """Build the extra-vars dict for Ansible from session config.""" config = session.config credentials = session.credentials # Map provider IDs to Ansible provider aliases provider_map = { "digitalocean": "digitalocean", "hetzner": "hetzner", "linode": "linode", "vultr": "vultr", } extra_vars: dict[str, Any] = { "provider": provider_map.get(session.provider, session.provider), } # Add credentials based on provider if session.provider == "digitalocean": extra_vars["do_token"] = credentials.get("do_token", "") elif session.provider == "hetzner": extra_vars["hcloud_token"] = credentials.get("hcloud_token", "") elif session.provider == "linode": extra_vars["linode_token"] = credentials.get("linode_token", "") elif session.provider == "vultr": # Vultr uses a different approach - write temp config file extra_vars["vultr_config"] = credentials.get("vultr_api_key", "") # Add server configuration if config.get("server_name"): extra_vars["server_name"] = config["server_name"] if config.get("region"): extra_vars["region"] = config["region"] # Add users list if config.get("users"): extra_vars["users"] = config["users"] # Add boolean options bool_options = [ ("wireguard_enabled", "wireguard_enabled"), ("ipsec_enabled", "ipsec_enabled"), ("dns_adblocking", "dns_adblocking"), ("ssh_tunneling", "ssh_tunneling"), ("ondemand_cellular", "ondemand_cellular"), ("ondemand_wifi", "ondemand_wifi"), ("store_pki", "store_pki"), ] for config_key, ansible_var in bool_options: if config_key in config: extra_vars[ansible_var] = config[config_key] return extra_vars def strip_ansi(text: str) -> str: """Remove ANSI escape codes from text.""" ansi_escape = re.compile(r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])") return ansi_escape.sub("", text) async def run_playbook(session: Session) -> AsyncGenerator[str, None]: """Run ansible-playbook with streaming output. This function uses asyncio.create_subprocess_exec which is safe from shell injection - arguments are passed as a list, not interpolated. Args: session: The deployment session Yields: Lines of output from the playbook run """ extra_vars = build_extra_vars(session) # Build command as list (safe - no shell interpolation) cmd = [ "uv", "run", "ansible-playbook", "main.yml", "--extra-vars", json.dumps(extra_vars), ] # Set up environment env = os.environ.copy() env["ANSIBLE_FORCE_COLOR"] = "0" # Disable colors for cleaner output env["ANSIBLE_NOCOLOR"] = "1" env["PYTHONUNBUFFERED"] = "1" try: # Start the process using create_subprocess_exec (safe, no shell) process = await asyncio.create_subprocess_exec( *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT, cwd=str(ALGO_DIR), env=env, ) session.process = process session.status = DeploymentStatus.RUNNING await sessions.update(session) # Clear credentials from memory now that we've started session.clear_credentials() # Stream output line by line config_path_pattern = re.compile(r"configs/[\d.]+|configs/localhost") if process.stdout: async for line_bytes in process.stdout: line = strip_ansi(line_bytes.decode("utf-8", errors="replace")).rstrip() if line: session.add_output(line) yield line # Try to capture config path from output if match := config_path_pattern.search(line): session.config_path = str(ALGO_DIR / match.group(0)) # Wait for process to complete await process.wait() session.exit_code = process.returncode if process.returncode == 0: session.status = DeploymentStatus.SUCCESS yield "[SUCCESS] Deployment completed successfully!" else: session.status = DeploymentStatus.FAILED session.error = f"Playbook exited with code {process.returncode}" yield f"[ERROR] Deployment failed with exit code {process.returncode}" except asyncio.CancelledError: session.status = DeploymentStatus.CANCELLED session.error = "Deployment was cancelled" if session.process: session.process.terminate() await session.process.wait() yield "[CANCELLED] Deployment was cancelled" raise except Exception as e: session.status = DeploymentStatus.FAILED session.error = str(e) yield f"[ERROR] {e!s}" finally: session.process = None await sessions.update(session) async def cancel_deployment(session: Session) -> bool: """Cancel a running deployment. Args: session: The session to cancel Returns: True if cancellation was successful """ if session.process and session.status == DeploymentStatus.RUNNING: session.process.terminate() try: await asyncio.wait_for(session.process.wait(), timeout=10.0) except TimeoutError: session.process.kill() await session.process.wait() session.status = DeploymentStatus.CANCELLED session.error = "Deployment was cancelled by user" await sessions.update(session) return True return False