aiptx 2.0.2__py3-none-any.whl

aipt_v2/runtime/local.py

@@ -0,0 +1,346 @@
+"""
+Local Runtime for AIPT v2
+=========================
+
+Non-sandboxed local execution runtime.
+Implements AbstractRuntime interface for local command execution.
+
+WARNING: This runtime provides NO isolation. Only use for trusted commands.
+"""
+
+import asyncio
+import uuid
+import os
+from dataclasses import dataclass, field
+from typing import Optional, Dict, Tuple, List
+from pathlib import Path
+from datetime import datetime
+
+from aipt_v2.utils.logging import logger
+
+
+@dataclass
+class LocalSandboxInfo:
+    """Information about a local 'sandbox' (execution context)."""
+
+    sandbox_id: str
+    created_at: datetime = field(default_factory=datetime.now)
+    working_dir: str = field(default_factory=lambda: os.getcwd())
+    is_local: bool = True
+    url: str = "local"
+
+
+class LocalRuntime:
+    """
+    Local runtime for command execution without sandboxing.
+
+    This runtime executes commands directly on the local system.
+    It provides a consistent interface with sandboxed runtimes
+    but does NOT provide any isolation.
+
+    Usage:
+        runtime = LocalRuntime()
+        sandbox = await runtime.create_sandbox()
+        stdout, stderr, code = await runtime.execute(sandbox.sandbox_id, "ls -la")
+        await runtime.destroy_sandbox(sandbox.sandbox_id)
+    """
+
+    def __init__(self, working_dir: Optional[str] = None):
+        """
+        Initialize local runtime.
+
+        Args:
+            working_dir: Default working directory for commands
+        """
+        self.working_dir = working_dir or os.getcwd()
+        self._sandboxes: Dict[str, LocalSandboxInfo] = {}
+        self._processes: Dict[str, asyncio.subprocess.Process] = {}
+
+    async def create_sandbox(
+        self,
+        image: Optional[str] = None,
+        working_dir: Optional[str] = None,
+        **kwargs
+    ) -> LocalSandboxInfo:
+        """
+        Create a local execution context.
+
+        Note: For local runtime, this just creates a tracking ID.
+        No actual isolation is provided.
+
+        Args:
+            image: Ignored for local runtime
+            working_dir: Working directory for commands
+            **kwargs: Additional options (ignored)
+
+        Returns:
+            LocalSandboxInfo with sandbox details
+        """
+        sandbox_id = f"local-{uuid.uuid4().hex[:8]}"
+        work_dir = working_dir or self.working_dir
+
+        # Ensure working directory exists
+        Path(work_dir).mkdir(parents=True, exist_ok=True)
+
+        info = LocalSandboxInfo(
+            sandbox_id=sandbox_id,
+            working_dir=work_dir,
+        )
+
+        self._sandboxes[sandbox_id] = info
+        logger.info("Created local execution context", sandbox_id=sandbox_id, working_dir=work_dir)
+
+        return info
+
+    async def get_sandbox_url(self, sandbox_id: str) -> str:
+        """
+        Get sandbox URL.
+
+        Args:
+            sandbox_id: Sandbox identifier
+
+        Returns:
+            URL string ('local' for local runtime)
+
+        Raises:
+            ValueError: If sandbox not found
+        """
+        if sandbox_id not in self._sandboxes:
+            raise ValueError(f"Unknown sandbox: {sandbox_id}")
+        return "local"
+
+    async def destroy_sandbox(self, sandbox_id: str) -> None:
+        """
+        Destroy sandbox and cleanup resources.
+
+        Args:
+            sandbox_id: Sandbox identifier
+        """
+        if sandbox_id in self._sandboxes:
+            # Kill any running processes
+            if sandbox_id in self._processes:
+                try:
+                    self._processes[sandbox_id].kill()
+                    await self._processes[sandbox_id].wait()
+                except ProcessLookupError:
+                    pass
+                del self._processes[sandbox_id]
+
+            del self._sandboxes[sandbox_id]
+            logger.info("Destroyed local sandbox", sandbox_id=sandbox_id)
+
+    async def execute(
+        self,
+        sandbox_id: str,
+        command: str,
+        timeout: int = 300,
+        env: Optional[Dict[str, str]] = None,
+    ) -> Tuple[str, str, int]:
+        """
+        Execute command in sandbox.
+
+        Args:
+            sandbox_id: Sandbox identifier
+            command: Shell command to execute
+            timeout: Command timeout in seconds
+            env: Additional environment variables
+
+        Returns:
+            Tuple of (stdout, stderr, exit_code)
+
+        Raises:
+            ValueError: If sandbox not found
+        """
+        if sandbox_id not in self._sandboxes:
+            raise ValueError(f"Unknown sandbox: {sandbox_id}")
+
+        sandbox = self._sandboxes[sandbox_id]
+
+        # Prepare environment
+        cmd_env = os.environ.copy()
+        if env:
+            cmd_env.update(env)
+
+        logger.debug(
+            "Executing command",
+            sandbox_id=sandbox_id,
+            command=command[:100] + "..." if len(command) > 100 else command,
+            timeout=timeout,
+        )
+
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                command,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+                cwd=sandbox.working_dir,
+                env=cmd_env,
+            )
+
+            self._processes[sandbox_id] = proc
+
+            try:
+                stdout_bytes, stderr_bytes = await asyncio.wait_for(
+                    proc.communicate(),
+                    timeout=timeout
+                )
+
+                stdout = stdout_bytes.decode("utf-8", errors="replace")
+                stderr = stderr_bytes.decode("utf-8", errors="replace")
+                exit_code = proc.returncode or 0
+
+                logger.debug(
+                    "Command completed",
+                    sandbox_id=sandbox_id,
+                    exit_code=exit_code,
+                    stdout_len=len(stdout),
+                    stderr_len=len(stderr),
+                )
+
+                return stdout, stderr, exit_code
+
+            except asyncio.TimeoutError:
+                logger.warning(
+                    "Command timed out",
+                    sandbox_id=sandbox_id,
+                    command=command[:50],
+                    timeout=timeout,
+                )
+                proc.kill()
+                await proc.wait()
+                return "", f"Command timed out after {timeout} seconds", 124
+
+        except Exception as e:
+            logger.error(
+                "Command execution failed",
+                sandbox_id=sandbox_id,
+                command=command[:50],
+                error=str(e),
+                exc_info=True,
+            )
+            return "", str(e), 1
+
+        finally:
+            if sandbox_id in self._processes:
+                del self._processes[sandbox_id]
+
+    async def execute_background(
+        self,
+        sandbox_id: str,
+        command: str,
+        env: Optional[Dict[str, str]] = None,
+    ) -> str:
+        """
+        Execute command in background.
+
+        Args:
+            sandbox_id: Sandbox identifier
+            command: Shell command to execute
+            env: Additional environment variables
+
+        Returns:
+            Process ID as string
+        """
+        if sandbox_id not in self._sandboxes:
+            raise ValueError(f"Unknown sandbox: {sandbox_id}")
+
+        sandbox = self._sandboxes[sandbox_id]
+
+        cmd_env = os.environ.copy()
+        if env:
+            cmd_env.update(env)
+
+        proc = await asyncio.create_subprocess_shell(
+            command,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            cwd=sandbox.working_dir,
+            env=cmd_env,
+        )
+
+        proc_id = f"{sandbox_id}-{proc.pid}"
+        self._processes[proc_id] = proc
+
+        logger.info(
+            "Started background process",
+            sandbox_id=sandbox_id,
+            proc_id=proc_id,
+            command=command[:50],
+        )
+
+        return proc_id
+
+    async def check_tool_available(self, tool_name: str) -> bool:
+        """
+        Check if a tool is available locally.
+
+        Args:
+            tool_name: Name of the tool to check
+
+        Returns:
+            True if tool is available
+        """
+        try:
+            proc = await asyncio.create_subprocess_shell(
+                f"which {tool_name}",
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+            stdout, _ = await proc.communicate()
+            return proc.returncode == 0 and len(stdout.strip()) > 0
+        except Exception:
+            return False
+
+    async def get_available_tools(self, tool_list: List[str]) -> Dict[str, bool]:
+        """
+        Check availability of multiple tools.
+
+        Args:
+            tool_list: List of tool names to check
+
+        Returns:
+            Dict mapping tool names to availability
+        """
+        results = {}
+        for tool in tool_list:
+            results[tool] = await self.check_tool_available(tool)
+        return results
+
+    def get_sandbox_info(self, sandbox_id: str) -> Optional[LocalSandboxInfo]:
+        """
+        Get information about a sandbox.
+
+        Args:
+            sandbox_id: Sandbox identifier
+
+        Returns:
+            LocalSandboxInfo or None if not found
+        """
+        return self._sandboxes.get(sandbox_id)
+
+    async def cleanup_all(self) -> int:
+        """
+        Cleanup all sandboxes and processes.
+
+        Returns:
+            Number of sandboxes cleaned up
+        """
+        count = len(self._sandboxes)
+
+        for sandbox_id in list(self._sandboxes.keys()):
+            await self.destroy_sandbox(sandbox_id)
+
+        logger.info("Cleaned up all local sandboxes", count=count)
+        return count
+
+
+# Singleton instance for convenience
+_default_runtime: Optional[LocalRuntime] = None
+
+
+def get_local_runtime() -> LocalRuntime:
+    """Get or create the default local runtime instance."""
+    global _default_runtime
+    if _default_runtime is None:
+        _default_runtime = LocalRuntime()
+    return _default_runtime

aipt_v2/runtime/tool_server.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+import signal
+import sys
+from multiprocessing import Process, Queue
+from typing import Any
+
+import uvicorn
+from fastapi import Depends, FastAPI, HTTPException, status
+from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
+from pydantic import BaseModel, ValidationError
+
+
+SANDBOX_MODE = os.getenv("AIPT_SANDBOX_MODE", "false").lower() == "true"
+if not SANDBOX_MODE:
+    raise RuntimeError("Tool server should only run in sandbox mode (AIPT_SANDBOX_MODE=true)")
+
+parser = argparse.ArgumentParser(description="Start AIPTx tool server")
+parser.add_argument("--token", required=True, help="Authentication token")
+parser.add_argument("--host", default="0.0.0.0", help="Host to bind to")  # nosec
+parser.add_argument("--port", type=int, required=True, help="Port to bind to")
+
+args = parser.parse_args()
+EXPECTED_TOKEN = args.token
+
+app = FastAPI()
+security = HTTPBearer()
+
+security_dependency = Depends(security)
+
+agent_processes: dict[str, dict[str, Any]] = {}
+agent_queues: dict[str, dict[str, Queue[Any]]] = {}
+
+
+def verify_token(credentials: HTTPAuthorizationCredentials) -> str:
+    if not credentials or credentials.scheme != "Bearer":
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication scheme. Bearer token required.",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    if credentials.credentials != EXPECTED_TOKEN:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication token",
+            headers={"WWW-Authenticate": "Bearer"},
+        )
+
+    return credentials.credentials
+
+
+class ToolExecutionRequest(BaseModel):
+    agent_id: str
+    tool_name: str
+    kwargs: dict[str, Any]
+
+
+class ToolExecutionResponse(BaseModel):
+    result: Any | None = None
+    error: str | None = None
+
+
+def agent_worker(_agent_id: str, request_queue: Queue[Any], response_queue: Queue[Any]) -> None:
+    null_handler = logging.NullHandler()
+
+    root_logger = logging.getLogger()
+    root_logger.handlers = [null_handler]
+    root_logger.setLevel(logging.CRITICAL)
+
+    from aipt_v2.tools.argument_parser import ArgumentConversionError, convert_arguments
+    from aipt_v2.tools.registry import get_tool_by_name
+
+    while True:
+        try:
+            request = request_queue.get()
+
+            if request is None:
+                break
+
+            tool_name = request["tool_name"]
+            kwargs = request["kwargs"]
+
+            try:
+                tool_func = get_tool_by_name(tool_name)
+                if not tool_func:
+                    response_queue.put({"error": f"Tool '{tool_name}' not found"})
+                    continue
+
+                converted_kwargs = convert_arguments(tool_func, kwargs)
+                result = tool_func(**converted_kwargs)
+
+                response_queue.put({"result": result})
+
+            except (ArgumentConversionError, ValidationError) as e:
+                response_queue.put({"error": f"Invalid arguments: {e}"})
+            except (RuntimeError, ValueError, ImportError) as e:
+                response_queue.put({"error": f"Tool execution error: {e}"})
+
+        except (RuntimeError, ValueError, ImportError) as e:
+            response_queue.put({"error": f"Worker error: {e}"})
+
+
+def ensure_agent_process(agent_id: str) -> tuple[Queue[Any], Queue[Any]]:
+    if agent_id not in agent_processes:
+        request_queue: Queue[Any] = Queue()
+        response_queue: Queue[Any] = Queue()
+
+        process = Process(
+            target=agent_worker, args=(agent_id, request_queue, response_queue), daemon=True
+        )
+        process.start()
+
+        agent_processes[agent_id] = {"process": process, "pid": process.pid}
+        agent_queues[agent_id] = {"request": request_queue, "response": response_queue}
+
+    return agent_queues[agent_id]["request"], agent_queues[agent_id]["response"]
+
+
+@app.post("/execute", response_model=ToolExecutionResponse)
+async def execute_tool(
+    request: ToolExecutionRequest, credentials: HTTPAuthorizationCredentials = security_dependency
+) -> ToolExecutionResponse:
+    verify_token(credentials)
+
+    request_queue, response_queue = ensure_agent_process(request.agent_id)
+
+    request_queue.put({"tool_name": request.tool_name, "kwargs": request.kwargs})
+
+    try:
+        loop = asyncio.get_event_loop()
+        response = await loop.run_in_executor(None, response_queue.get)
+
+        if "error" in response:
+            return ToolExecutionResponse(error=response["error"])
+        return ToolExecutionResponse(result=response.get("result"))
+
+    except (RuntimeError, ValueError, OSError) as e:
+        return ToolExecutionResponse(error=f"Worker error: {e}")
+
+
+@app.post("/register_agent")
+async def register_agent(
+    agent_id: str, credentials: HTTPAuthorizationCredentials = security_dependency
+) -> dict[str, str]:
+    verify_token(credentials)
+
+    ensure_agent_process(agent_id)
+    return {"status": "registered", "agent_id": agent_id}
+
+
+@app.get("/health")
+async def health_check() -> dict[str, Any]:
+    return {
+        "status": "healthy",
+        "sandbox_mode": str(SANDBOX_MODE),
+        "environment": "sandbox" if SANDBOX_MODE else "main",
+        "auth_configured": "true" if EXPECTED_TOKEN else "false",
+        "active_agents": len(agent_processes),
+        "agents": list(agent_processes.keys()),
+    }
+
+
+def cleanup_all_agents() -> None:
+    for agent_id in list(agent_processes.keys()):
+        try:
+            agent_queues[agent_id]["request"].put(None)
+            process = agent_processes[agent_id]["process"]
+
+            process.join(timeout=1)
+
+            if process.is_alive():
+                process.terminate()
+                process.join(timeout=1)
+
+            if process.is_alive():
+                process.kill()
+
+        except (BrokenPipeError, EOFError, OSError) as e:
+            logging.getLogger(__name__).debug(f"Expected error during process cleanup: {e}")
+        except (RuntimeError, ValueError) as e:
+            logging.getLogger(__name__).debug(f"Error during agent cleanup: {e}")
+
+
+def signal_handler(_signum: int, _frame: Any) -> None:
+    signal.signal(signal.SIGPIPE, signal.SIG_IGN) if hasattr(signal, "SIGPIPE") else None
+    cleanup_all_agents()
+    sys.exit(0)
+
+
+if hasattr(signal, "SIGPIPE"):
+    signal.signal(signal.SIGPIPE, signal.SIG_IGN)
+
+signal.signal(signal.SIGTERM, signal_handler)
+signal.signal(signal.SIGINT, signal_handler)
+
+if __name__ == "__main__":
+    try:
+        uvicorn.run(app, host=args.host, port=args.port, log_level="info")
+    finally:
+        cleanup_all_agents()

aipt_v2/scanners/__init__.py

@@ -0,0 +1,28 @@
+"""
+AIPT Scanners Module
+
+Integrations with popular security scanning tools:
+- Nuclei - Template-based vulnerability scanner
+- Nmap - Network scanner
+- Nikto - Web server scanner
+- SQLMap - SQL injection scanner
+- Dirb/Gobuster - Directory enumeration
+"""
+
+from .base import BaseScanner, ScanResult
+from .nuclei import NucleiScanner, NucleiConfig
+from .nmap import NmapScanner, NmapConfig
+from .nikto import NiktoScanner
+from .web import WebScanner, WebScanConfig
+
+__all__ = [
+    "BaseScanner",
+    "ScanResult",
+    "NucleiScanner",
+    "NucleiConfig",
+    "NmapScanner",
+    "NmapConfig",
+    "NiktoScanner",
+    "WebScanner",
+    "WebScanConfig",
+]