decodingtrust-agent-sdk 0.1.0__py3-none-any.whl

utils/reset_helpers.py

@@ -0,0 +1,318 @@
+import asyncio
+import subprocess
+import urllib.request
+import urllib.error
+from pathlib import Path
+from typing import Dict, Any, Optional
+
+from .template_helpers import render_template
+
+
+# Check if sudo is needed for docker commands
+_USE_SUDO = False
+
+
+async def reset_via_endpoints(
+    env_name: str,
+    ports: Dict[str, int],
+    env_config: Dict[str, Any],
+    timeout: int = 30,
+    max_retries: int = 10,
+    retry_delay: float = 2.0,
+) -> None:
+    """
+    Reset an environment via HTTP API endpoints.
+
+    Args:
+        env_name: Name of the environment (e.g., "salesforce", "gmail")
+        ports: Dict mapping port variable names to allocated port numbers
+        env_config: Full env.yaml config dict
+        timeout: HTTP request timeout in seconds
+        max_retries: Maximum number of retry attempts per endpoint
+        retry_delay: Delay between retries in seconds
+
+    Raises:
+        RuntimeError: If any endpoint call fails after all retries
+    """
+    environments = env_config.get("environments", {})
+    env_def = environments.get(env_name, {})
+    reset_endpoints = env_def.get("reset_endpoints", {})
+
+    if not reset_endpoints:
+        raise RuntimeError(f"No reset endpoints configured for {env_name}")
+
+    print(f"[RESET] Resetting {env_name} via API endpoints", flush=True)
+
+    for endpoint_name, endpoint_config in reset_endpoints.items():
+        url_template = endpoint_config.get("url", "")
+        method = endpoint_config.get("method", "POST").upper()
+
+        # Resolve ${VAR_NAME} placeholders using ports
+        url = render_template(url_template, ports)
+
+        # Check if URL still has unresolved placeholders
+        if "${" in url:
+            print(f"[RESET] Skipping {endpoint_name}: unresolved placeholders in URL", flush=True)
+            continue
+
+        # Make the HTTP request with retries
+        last_error = None
+
+        for attempt in range(1, max_retries + 1):
+            try:
+                req = urllib.request.Request(url, method=method, data=b"")
+                req.add_header("Content-Type", "application/json")
+
+                with urllib.request.urlopen(req, timeout=timeout) as response:
+                    status = response.status
+                    if 200 <= status < 300:
+                        print(f"[RESET] {env_name}/{endpoint_name}: OK (HTTP {status})", flush=True)
+                        break
+                    else:
+                        last_error = RuntimeError(f"Reset endpoint '{endpoint_name}' returned HTTP {status}")
+                        if attempt < max_retries:
+                            print(f"[RESET] {endpoint_name} returned HTTP {status}, retrying ({attempt}/{max_retries})...", flush=True)
+                            await asyncio.sleep(retry_delay)
+                        else:
+                            raise last_error
+            except urllib.error.HTTPError as e:
+                last_error = RuntimeError(f"Reset endpoint '{endpoint_name}' failed: HTTP {e.code}")
+                if attempt < max_retries:
+                    print(f"[RESET] {endpoint_name} failed with HTTP {e.code}, retrying ({attempt}/{max_retries})...", flush=True)
+                    await asyncio.sleep(retry_delay)
+                else:
+                    raise last_error
+            except urllib.error.URLError as e:
+                last_error = RuntimeError(f"Reset endpoint '{endpoint_name}' failed: {e.reason}")
+                if attempt < max_retries:
+                    print(f"[RESET] {endpoint_name} failed: {e.reason}, retrying ({attempt}/{max_retries})...", flush=True)
+                    await asyncio.sleep(retry_delay)
+                else:
+                    raise last_error
+
+
+async def reset_via_scripts(
+    env_name: str,
+    project_name: str,
+    compose_file: Path,
+    env_config: Dict[str, Any],
+    timeout: int = 30,
+) -> None:
+    """
+    Reset an environment via docker exec scripts.
+
+    Args:
+        env_name: Name of the environment (e.g., "salesforce", "gmail")
+        project_name: Docker compose project name
+        compose_file: Path to docker-compose.yml file
+        env_config: Full env.yaml config dict
+        timeout: Script execution timeout in seconds
+
+    Raises:
+        RuntimeError: If any script execution fails
+    """
+    environments = env_config.get("environments", {})
+    env_def = environments.get(env_name, {})
+    reset_scripts = env_def.get("reset_scripts", {})
+
+    if not reset_scripts:
+        raise RuntimeError(f"No reset scripts configured for {env_name}")
+
+    print(f"[RESET] Resetting {env_name} via scripts", flush=True)
+
+    for service, script_path in reset_scripts.items():
+        cmd = [
+            "docker", "compose", "-p", project_name,
+            "-f", str(compose_file), "exec", "-T", service,
+            "/bin/sh", "-c", script_path
+        ]
+        if _USE_SUDO:
+            cmd = ["sudo"] + cmd
+
+        proc = await asyncio.create_subprocess_exec(
+            *cmd,
+            cwd=str(compose_file.parent),
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        try:
+            stdout, stderr = await asyncio.wait_for(proc.communicate(), timeout=timeout)
+        except asyncio.TimeoutError:
+            raise RuntimeError(f"Reset script for {service} timed out")
+
+        if proc.returncode != 0:
+            raise RuntimeError(f"Reset script for {service} failed: {stderr.decode()}")
+
+        print(f"[RESET] {env_name}/{service}: OK (via script)", flush=True)
+
+
+async def reset_environment(
+    env_name: str,
+    ports: Dict[str, int],
+    env_config: Dict[str, Any],
+    project_name: Optional[str] = None,
+    compose_file: Optional[Path] = None,
+    endpoint_timeout: int = 30,
+    script_timeout: int = 30,
+    max_retries: int = 10,
+) -> None:
+    """
+    Reset an environment's data state.
+
+    Priority order:
+    1. API endpoints (reset_endpoints) - preferred, safer
+    2. Docker exec scripts (reset_scripts) - fallback if endpoints fail
+
+    Args:
+        env_name: Name of the environment (e.g., "salesforce", "gmail")
+        ports: Dict mapping port variable names to allocated port numbers
+        env_config: Full env.yaml config dict
+        project_name: Docker compose project name (required for script fallback)
+        compose_file: Path to docker-compose.yml (required for script fallback)
+        endpoint_timeout: HTTP request timeout in seconds
+        script_timeout: Script execution timeout in seconds
+        max_retries: Maximum retry attempts for endpoints
+
+    Raises:
+        RuntimeError: If reset fails (both endpoints and scripts failed)
+    """
+    environments = env_config.get("environments", {})
+    env_def = environments.get(env_name, {})
+    reset_endpoints = env_def.get("reset_endpoints", {})
+    reset_scripts = env_def.get("reset_scripts", {})
+
+    if not reset_endpoints and not reset_scripts:
+        return  # No reset configured, skip silently
+
+    # Try endpoints first (prioritized)
+    if reset_endpoints:
+        try:
+            await reset_via_endpoints(
+                env_name=env_name,
+                ports=ports,
+                env_config=env_config,
+                timeout=endpoint_timeout,
+                max_retries=max_retries,
+            )
+            return
+        except RuntimeError as e:
+            # If endpoints failed and we have scripts, fall back to scripts
+            if reset_scripts and project_name and compose_file:
+                print(f"[RESET] Endpoint reset failed ({e}), falling back to scripts", flush=True)
+            elif not reset_scripts:
+                raise
+            else:
+                # Scripts configured but no project_name/compose_file provided
+                print(f"[RESET] Endpoint reset failed, scripts available but project info not provided", flush=True)
+                raise
+
+    # Use scripts (either as primary or fallback)
+    if reset_scripts:
+        if not project_name or not compose_file:
+            raise RuntimeError(f"Reset scripts configured for {env_name} but project_name/compose_file not provided")
+        await reset_via_scripts(
+            env_name=env_name,
+            project_name=project_name,
+            compose_file=compose_file,
+            env_config=env_config,
+            timeout=script_timeout,
+        )
+
+
+async def reset_all_environments(
+    ports: Dict[str, int],
+    env_config: Dict[str, Any],
+    project_name_func: Optional[callable] = None,
+    compose_file_func: Optional[callable] = None,
+    endpoint_timeout: int = 5,
+    max_retries: int = 3,
+) -> int:
+    """
+    Reset all environments that have allocated ports.
+
+    This function identifies which environments have ports allocated and resets them.
+
+    Args:
+        ports: Dict mapping port variable names to allocated port numbers
+               (e.g., {"SALESFORCE_API_PORT": 8080, "GMAIL_AUTH_PORT": 8030})
+        env_config: Full env.yaml config dict
+        project_name_func: Optional callable(env_name) -> project_name for script fallback
+        compose_file_func: Optional callable(env_name) -> compose_file path for script fallback
+        endpoint_timeout: HTTP request timeout in seconds (shorter for batch reset)
+        max_retries: Maximum retry attempts (fewer for batch reset)
+
+    Returns:
+        Number of environments successfully reset
+    """
+    environments = env_config.get("environments", {})
+    reset_count = 0
+
+    # Build a mapping of port vars to environment names
+    port_to_env: Dict[str, str] = {}
+    for env_name, env_def in environments.items():
+        env_ports = env_def.get("ports", {})
+        for port_var in env_ports.keys():
+            port_to_env[port_var] = env_name
+
+    # Find which environments have allocated ports
+    envs_to_reset = set()
+    for port_var in ports.keys():
+        if port_var in port_to_env:
+            envs_to_reset.add(port_to_env[port_var])
+
+    if not envs_to_reset:
+        return 0
+
+    # Reset each environment
+    for env_name in envs_to_reset:
+        env_def = environments.get(env_name, {})
+        reset_endpoints = env_def.get("reset_endpoints", {})
+        reset_scripts = env_def.get("reset_scripts", {})
+
+        if not reset_endpoints and not reset_scripts:
+            continue
+
+        # Get project info for script fallback if functions provided
+        project_name = project_name_func(env_name) if project_name_func else None
+        compose_file = compose_file_func(env_name) if compose_file_func else None
+
+        try:
+            await reset_environment(
+                env_name=env_name,
+                ports=ports,
+                env_config=env_config,
+                project_name=project_name,
+                compose_file=compose_file,
+                endpoint_timeout=endpoint_timeout,
+                max_retries=max_retries,
+            )
+            reset_count += 1
+        except Exception as e:
+            print(f"[RESET] Failed to reset {env_name}: {e}", flush=True)
+
+    return reset_count
+
+
+def reset_all_environments_sync(
+    ports: Dict[str, int],
+    env_config: Dict[str, Any],
+    project_name_func: Optional[callable] = None,
+    compose_file_func: Optional[callable] = None,
+    endpoint_timeout: int = 5,
+    max_retries: int = 3,
+) -> int:
+    """
+    Synchronous wrapper for reset_all_environments.
+
+    Useful for calling from synchronous code (like query_victim.exec).
+    """
+    return asyncio.get_event_loop().run_until_complete(
+        reset_all_environments(
+            ports=ports,
+            env_config=env_config,
+            project_name_func=project_name_func,
+            compose_file_func=compose_file_func,
+            endpoint_timeout=endpoint_timeout,
+            max_retries=max_retries,
+        )
+    )

utils/resource_manager.py

@@ -0,0 +1,370 @@
+import grp
+import os
+import socket
+import subprocess
+import threading
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Dict, List, Optional, Set
+import random
+
+
+DEFAULT_PORT_START = 8000
+DEFAULT_PORT_END = 20000
+
+
+def _needs_sudo_for_docker() -> bool:
+    """Check if we need sudo to run docker commands."""
+    # First, try running docker directly
+    try:
+        result = subprocess.run(
+            ["docker", "ps"],
+            capture_output=True,
+            timeout=5
+        )
+        if result.returncode == 0:
+            return False
+    except (subprocess.SubprocessError, FileNotFoundError):
+        pass
+    
+    # Fallback to group check
+    try:
+        docker_gid = grp.getgrnam("docker").gr_gid
+        if docker_gid in os.getgroups():
+            return False
+    except (KeyError, OSError):
+        pass
+    return True
+
+
+@dataclass
+class TaskResources:
+    """Resources allocated for a single task."""
+    task_id: str
+    # Port allocations: var_name -> port
+    ports: Dict[str, int] = field(default_factory=dict)
+    # Docker compose projects started for this task
+    docker_projects: List[str] = field(default_factory=list)
+    # Docker compose file paths for teardown
+    compose_files: Dict[str, Path] = field(default_factory=dict)
+
+
+class ResourceManager:
+    """
+    Singleton in-process manager for ports and Docker resources.
+    
+    Thread-safe for concurrent access from parallel tasks.
+    """
+    
+    _instance: Optional["ResourceManager"] = None
+    _lock = threading.Lock()
+    
+    def __init__(self):
+        self._tasks: Dict[str, TaskResources] = {}
+        self._used_ports: Set[int] = set()
+        self._port_range = self._get_port_range()
+        self._mutex = threading.Lock()
+    
+    @classmethod
+    def instance(cls) -> "ResourceManager":
+        """Get or create the singleton instance."""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+    
+    @classmethod
+    def reset(cls) -> None:
+        """Reset the singleton (for testing)."""
+        with cls._lock:
+            if cls._instance is not None:
+                # Cleanup all tasks before reset
+                for task_id in list(cls._instance._tasks.keys()):
+                    try:
+                        cls._instance.cleanup_task(task_id)
+                    except Exception:
+                        pass
+            cls._instance = None
+    
+    def _get_port_range(self) -> tuple:
+        """Get port range from environment or use defaults."""
+        env_range = os.getenv("DT_PORT_RANGE")
+        if env_range:
+            try:
+                start_str, end_str = env_range.split("-", 1)
+                return (int(start_str.strip()), int(end_str.strip()))
+            except Exception:
+                pass
+        
+        start = int(os.getenv("DT_PORT_RANGE_START", str(DEFAULT_PORT_START)))
+        end = int(os.getenv("DT_PORT_RANGE_END", str(DEFAULT_PORT_END)))
+        return (start, end)
+    
+    def _is_port_available(self, port: int) -> bool:
+        """Check if a port is available on localhost.
+
+        Performs comprehensive checks:
+        1. Internal check against already-allocated ports
+        2. Check Docker containers for port mappings
+        3. Check if something is listening on IPv4/IPv6
+        4. Attempt to bind to the port on both IPv4 and IPv6
+        """
+        # Check internal tracking first
+        if port in self._used_ports:
+            return False
+
+        # Check if Docker has containers using this port (even if not listening yet)
+        try:
+            cmd = ["docker", "ps", "--format", "{{.Ports}}"]
+            if _needs_sudo_for_docker():
+                cmd = ["sudo"] + cmd
+
+            result = subprocess.run(
+                cmd,
+                capture_output=True,
+                text=True,
+                timeout=2
+            )
+            if result.returncode == 0:
+                # Parse port mappings like "0.0.0.0:8080->8080/tcp, [::]:8080->8080/tcp"
+                for line in result.stdout.splitlines():
+                    if not line.strip():
+                        continue
+
+                    # Split by comma to handle multiple port mappings
+                    for port_mapping in line.split(','):
+                        port_mapping = port_mapping.strip()
+
+                        # Check for host port mappings: "0.0.0.0:8080->..." or "[::]:8080->..."
+                        # The format is: HOST_IP:HOST_PORT->CONTAINER_PORT/PROTO
+                        if '->' in port_mapping:
+                            # Extract the host port (before ->)
+                            host_part = port_mapping.split('->')[0]
+                            if ':' in host_part:
+                                # Get the port number after the last colon
+                                host_port_str = host_part.split(':')[-1]
+                                try:
+                                    host_port = int(host_port_str)
+                                    if host_port == port:
+                                        return False
+                                except ValueError:
+                                    continue
+        except (subprocess.TimeoutExpired, FileNotFoundError):
+            pass
+
+        # Check if something is already listening (IPv4)
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                sock.settimeout(0.1)
+                result = sock.connect_ex(("127.0.0.1", port))
+                if result == 0:
+                    return False
+        except OSError:
+            pass
+
+        # Check if something is already listening (IPv6)
+        try:
+            with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock:
+                sock.settimeout(0.1)
+                result = sock.connect_ex(("::1", port))
+                if result == 0:
+                    return False
+        except OSError:
+            pass
+
+        # Try to bind on IPv4
+        try:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+                sock.bind(("0.0.0.0", port))
+        except OSError:
+            return False
+
+        # Try to bind on IPv6 (Docker often binds to IPv6)
+        try:
+            with socket.socket(socket.AF_INET6, socket.SOCK_STREAM) as sock:
+                sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
+                sock.bind(("::", port))
+        except (OSError, AttributeError):
+            return False
+
+        return True
+    
+    def _get_or_create_task(self, task_id: str) -> TaskResources:
+        """Get or create task resources."""
+        if task_id not in self._tasks:
+            self._tasks[task_id] = TaskResources(task_id=task_id)
+        return self._tasks[task_id]
+    
+    def allocate_port(self, task_id: str, var_name: str, default: Optional[int] = None) -> int:
+        """
+        Allocate a port for a task.
+        
+        Args:
+            task_id: Unique identifier for the task
+            var_name: Environment variable name (e.g., "GMAIL_AUTH_PORT")
+            default: Default port to try first (optional)
+        
+        Returns:
+            Allocated port number
+        
+        Raises:
+            RuntimeError: If no port could be allocated
+        """
+        with self._mutex:
+            task = self._get_or_create_task(task_id)
+            
+            # If already allocated for this task, return it
+            if var_name in task.ports:
+                return task.ports[var_name]
+            
+            # Try default port first if specified
+            if default is not None:
+                if self._is_port_available(default):
+                    task.ports[var_name] = default
+                    self._used_ports.add(default)
+                    print(f"[PORT] Allocated default port {default} for {task_id}:{var_name}", flush=True)
+                    return default
+                else:
+                    print(f"[PORT] Default port {default} unavailable for {task_id}:{var_name}, trying random", flush=True)
+
+            # Find a free port in range
+            start, end = self._port_range
+            attempts = min(1000, end - start)
+
+            for _ in range(attempts):
+                port = random.randint(start, end)
+                if self._is_port_available(port):
+                    task.ports[var_name] = port
+                    self._used_ports.add(port)
+                    print(f"[PORT] Allocated random port {port} for {task_id}:{var_name}", flush=True)
+                    return port
+            
+            # Sequential scan as fallback
+            for port in range(start, end + 1):
+                if self._is_port_available(port):
+                    task.ports[var_name] = port
+                    self._used_ports.add(port)
+                    return port
+            
+            raise RuntimeError(f"Unable to allocate port for {var_name} in range [{start}, {end}]")
+    
+    def get_port(self, task_id: str, var_name: str) -> Optional[int]:
+        """Get an allocated port for a task (returns None if not allocated)."""
+        with self._mutex:
+            task = self._tasks.get(task_id)
+            if task:
+                return task.ports.get(var_name)
+            return None
+    
+    def get_all_ports(self, task_id: str) -> Dict[str, int]:
+        """Get all allocated ports for a task."""
+        with self._mutex:
+            task = self._tasks.get(task_id)
+            if task:
+                return dict(task.ports)
+            return {}
+    
+    def register_docker_project(
+        self, 
+        task_id: str, 
+        project_name: str, 
+        compose_file: Path
+    ) -> None:
+        """Register a Docker compose project for a task."""
+        with self._mutex:
+            task = self._get_or_create_task(task_id)
+            if project_name not in task.docker_projects:
+                task.docker_projects.append(project_name)
+            task.compose_files[project_name] = compose_file
+    
+    def get_docker_projects(self, task_id: str) -> List[str]:
+        """Get all Docker projects for a task."""
+        with self._mutex:
+            task = self._tasks.get(task_id)
+            if task:
+                return list(task.docker_projects)
+            return []
+    
+    def cleanup_task(self, task_id: str, verbose: bool = True) -> None:
+        """
+        Cleanup all resources for a task.
+        
+        This will:
+        1. Stop and remove all Docker compose projects for this task
+        2. Release all allocated ports
+        """
+        with self._mutex:
+            task = self._tasks.get(task_id)
+            if not task:
+                return
+            
+            # Teardown Docker projects
+            for project_name in task.docker_projects:
+                compose_file = task.compose_files.get(project_name)
+                if compose_file and compose_file.exists():
+                    try:
+                        if verbose:
+                            print(f"[CLEANUP] Stopping Docker project: {project_name}")
+                        subprocess.run(
+                            [
+                                "docker", "compose",
+                                "-p", project_name,
+                                "-f", str(compose_file),
+                                "down", "--remove-orphans"
+                            ],
+                            cwd=str(compose_file.parent),
+                            capture_output=True,
+                            check=False,
+                            timeout=60,
+                        )
+                    except Exception as e:
+                        if verbose:
+                            print(f"[CLEANUP] Error stopping {project_name}: {e}")
+            
+            # Release ports
+            for port in task.ports.values():
+                self._used_ports.discard(port)
+            
+            if verbose and task.ports:
+                print(f"[CLEANUP] Released {len(task.ports)} port(s) for task {task_id}")
+            
+            # Remove task from tracking
+            del self._tasks[task_id]
+    
+    def cleanup_all(self, verbose: bool = True) -> None:
+        """Cleanup all tracked resources."""
+        with self._mutex:
+            task_ids = list(self._tasks.keys())
+        
+        for task_id in task_ids:
+            self.cleanup_task(task_id, verbose=verbose)
+    
+    def snapshot(self) -> Dict[str, dict]:
+        """Get a snapshot of all tracked resources (for debugging)."""
+        with self._mutex:
+            return {
+                task_id: {
+                    "ports": dict(task.ports),
+                    "docker_projects": list(task.docker_projects),
+                }
+                for task_id, task in self._tasks.items()
+            }
+
+
+def generate_task_id(task_dir: Path) -> str:
+    """Generate a unique task ID based on task directory and PID."""
+    task_name = task_dir.name
+    pid = os.getpid()
+    return f"{task_name}_{pid}"
+
+
+def generate_project_name(task_id: str, env_name: Optional[str] = None) -> str:
+    """Generate a Docker compose project name."""
+    base = f"wf_{task_id}".lower()
+    # Sanitize for Docker compose requirements
+    base = "".join(c if c.isalnum() or c in "-_" else "_" for c in base)
+    if env_name:
+        return f"{base}_{env_name}"
+    return base
+