tweek 0.1.0__py3-none-any.whl

tweek/proxy/interceptor.py

@@ -0,0 +1,313 @@
+"""
+LLM API Interceptor - Screens requests and responses to LLM APIs.
+
+This module provides the core interception logic for the Tweek proxy,
+analyzing LLM API traffic for security threats.
+"""
+
+import json
+import re
+import uuid
+from dataclasses import dataclass, field
+from typing import Optional, Any
+from enum import Enum
+
+
+class LLMProvider(Enum):
+    """Supported LLM API providers."""
+    ANTHROPIC = "anthropic"
+    OPENAI = "openai"
+    GOOGLE = "google"
+    BEDROCK = "bedrock"
+    UNKNOWN = "unknown"
+
+
+@dataclass
+class InterceptionResult:
+    """Result of screening an LLM API request or response."""
+    allowed: bool
+    provider: LLMProvider
+    reason: Optional[str] = None
+    blocked_tools: list[str] = field(default_factory=list)
+    warnings: list[str] = field(default_factory=list)
+    matched_patterns: list[str] = field(default_factory=list)
+
+
+@dataclass
+class ToolCall:
+    """Represents a tool call extracted from an LLM response."""
+    id: str
+    name: str
+    input: dict[str, Any]
+    provider: LLMProvider
+
+
+class LLMAPIInterceptor:
+    """
+    Intercepts and screens LLM API traffic.
+
+    Analyzes both requests (prompts) and responses (tool calls)
+    for security threats using Tweek's pattern matching engine.
+    """
+
+    # API endpoints to monitor
+    MONITORED_HOSTS = {
+        "api.anthropic.com": LLMProvider.ANTHROPIC,
+        "api.openai.com": LLMProvider.OPENAI,
+        "generativelanguage.googleapis.com": LLMProvider.GOOGLE,
+    }
+
+    # Bedrock uses regional endpoints
+    BEDROCK_PATTERN = re.compile(r"bedrock-runtime\.[\w-]+\.amazonaws\.com")
+
+    def __init__(self, pattern_matcher=None, security_logger=None):
+        """
+        Initialize the interceptor.
+
+        Args:
+            pattern_matcher: Tweek PatternMatcher instance for screening
+            security_logger: SecurityLogger for audit logging
+        """
+        self.pattern_matcher = pattern_matcher
+        self.security_logger = security_logger
+
+    def _new_correlation_id(self) -> str:
+        """Generate a new correlation ID for a screening pass."""
+        return uuid.uuid4().hex[:12]
+
+    def identify_provider(self, host: str) -> LLMProvider:
+        """Identify the LLM provider from the request host."""
+        if host in self.MONITORED_HOSTS:
+            return self.MONITORED_HOSTS[host]
+
+        if self.BEDROCK_PATTERN.match(host):
+            return LLMProvider.BEDROCK
+
+        return LLMProvider.UNKNOWN
+
+    def should_intercept(self, host: str) -> bool:
+        """Check if this host should be intercepted."""
+        return self.identify_provider(host) != LLMProvider.UNKNOWN
+
+    def extract_tool_calls_anthropic(self, response: dict) -> list[ToolCall]:
+        """Extract tool calls from Anthropic API response."""
+        tool_calls = []
+
+        content = response.get("content", [])
+        for block in content:
+            if block.get("type") == "tool_use":
+                tool_calls.append(ToolCall(
+                    id=block.get("id", ""),
+                    name=block.get("name", ""),
+                    input=block.get("input", {}),
+                    provider=LLMProvider.ANTHROPIC,
+                ))
+
+        return tool_calls
+
+    def extract_tool_calls_openai(self, response: dict) -> list[ToolCall]:
+        """Extract tool calls from OpenAI API response."""
+        tool_calls = []
+
+        choices = response.get("choices", [])
+        for choice in choices:
+            message = choice.get("message", {})
+            for tc in message.get("tool_calls", []):
+                func = tc.get("function", {})
+                try:
+                    args = json.loads(func.get("arguments", "{}"))
+                except json.JSONDecodeError:
+                    args = {"_raw": func.get("arguments", "")}
+
+                tool_calls.append(ToolCall(
+                    id=tc.get("id", ""),
+                    name=func.get("name", ""),
+                    input=args,
+                    provider=LLMProvider.OPENAI,
+                ))
+
+        return tool_calls
+
+    def extract_tool_calls(self, response: dict, provider: LLMProvider) -> list[ToolCall]:
+        """Extract tool calls from an LLM API response."""
+        if provider == LLMProvider.ANTHROPIC:
+            return self.extract_tool_calls_anthropic(response)
+        elif provider == LLMProvider.OPENAI:
+            return self.extract_tool_calls_openai(response)
+        else:
+            return []
+
+    def screen_tool_call(self, tool_call: ToolCall) -> InterceptionResult:
+        """Screen a single tool call for security threats."""
+        if not self.pattern_matcher:
+            return InterceptionResult(allowed=True, provider=tool_call.provider)
+
+        # Build command string for pattern matching
+        # Handle common tool patterns
+        command = ""
+        tool_name = tool_call.name.lower()
+
+        if tool_name in ("bash", "shell", "execute", "run_command"):
+            command = tool_call.input.get("command", "")
+        elif tool_name in ("read", "read_file", "cat"):
+            command = f"cat {tool_call.input.get('path', tool_call.input.get('file_path', ''))}"
+        elif tool_name in ("write", "write_file"):
+            path = tool_call.input.get('path', tool_call.input.get('file_path', ''))
+            command = f"write to {path}"
+        elif tool_name in ("fetch", "web_fetch", "curl", "http"):
+            url = tool_call.input.get('url', '')
+            command = f"curl {url}"
+        else:
+            # Generic: serialize input for pattern matching
+            command = json.dumps(tool_call.input)
+
+        # Run through pattern matcher
+        matches = self.pattern_matcher.match(command)
+
+        if matches:
+            # Log the blocked attempt
+            if self.security_logger:
+                self.security_logger.log_event(
+                    event_type="proxy_block",
+                    tool=tool_call.name,
+                    command=command[:500],  # Truncate for logging
+                    patterns=matches,
+                )
+
+            return InterceptionResult(
+                allowed=False,
+                provider=tool_call.provider,
+                reason=f"Blocked by patterns: {', '.join(matches)}",
+                blocked_tools=[tool_call.name],
+                matched_patterns=matches,
+            )
+
+        return InterceptionResult(allowed=True, provider=tool_call.provider)
+
+    def screen_response(self, response_body: bytes, provider: LLMProvider) -> InterceptionResult:
+        """
+        Screen an LLM API response for dangerous tool calls.
+
+        Args:
+            response_body: Raw response body bytes
+            provider: The LLM provider
+
+        Returns:
+            InterceptionResult with screening decision
+        """
+        correlation_id = self._new_correlation_id()
+
+        try:
+            response = json.loads(response_body)
+        except json.JSONDecodeError:
+            # Can't parse, allow through
+            return InterceptionResult(allowed=True, provider=provider)
+
+        tool_calls = self.extract_tool_calls(response, provider)
+
+        if not tool_calls:
+            return InterceptionResult(allowed=True, provider=provider)
+
+        # Screen each tool call
+        blocked_tools = []
+        all_patterns = []
+        all_warnings = []
+
+        for tc in tool_calls:
+            result = self.screen_tool_call(tc)
+            if not result.allowed:
+                blocked_tools.extend(result.blocked_tools)
+                all_patterns.extend(result.matched_patterns)
+            all_warnings.extend(result.warnings)
+
+        if blocked_tools:
+            self._log_proxy_event(
+                "block", blocked_tools, all_patterns, provider, correlation_id
+            )
+            return InterceptionResult(
+                allowed=False,
+                provider=provider,
+                reason=f"Blocked dangerous tool calls: {', '.join(blocked_tools)}",
+                blocked_tools=blocked_tools,
+                matched_patterns=all_patterns,
+                warnings=all_warnings,
+            )
+
+        return InterceptionResult(
+            allowed=True,
+            provider=provider,
+            warnings=all_warnings,
+        )
+
+    def _log_proxy_event(
+        self,
+        decision: str,
+        blocked_tools: list,
+        matched_patterns: list,
+        provider: LLMProvider,
+        correlation_id: str,
+    ):
+        """Log a proxy screening event to the security logger."""
+        try:
+            from tweek.logging.security_log import get_logger, SecurityEvent, EventType
+            get_logger().log(SecurityEvent(
+                event_type=EventType.PROXY_EVENT,
+                tool_name="http_proxy",
+                decision=decision,
+                decision_reason=f"Blocked tools: {', '.join(blocked_tools)}" if blocked_tools else None,
+                metadata={
+                    "provider": provider.value,
+                    "blocked_tools": blocked_tools,
+                    "matched_patterns": matched_patterns,
+                },
+                correlation_id=correlation_id,
+                source="http_proxy",
+            ))
+        except Exception:
+            pass
+
+    def screen_request(self, request_body: bytes, provider: LLMProvider) -> InterceptionResult:
+        """
+        Screen an LLM API request for prompt injection attempts.
+
+        Args:
+            request_body: Raw request body bytes
+            provider: The LLM provider
+
+        Returns:
+            InterceptionResult with screening decision
+        """
+        if not self.pattern_matcher:
+            return InterceptionResult(allowed=True, provider=provider)
+
+        try:
+            request = json.loads(request_body)
+        except json.JSONDecodeError:
+            return InterceptionResult(allowed=True, provider=provider)
+
+        # Extract user messages for screening
+        messages = []
+
+        if provider == LLMProvider.ANTHROPIC:
+            messages = request.get("messages", [])
+        elif provider == LLMProvider.OPENAI:
+            messages = request.get("messages", [])
+
+        # Screen each user message
+        warnings = []
+
+        for msg in messages:
+            if msg.get("role") == "user":
+                content = msg.get("content", "")
+                if isinstance(content, str):
+                    # Check for prompt injection patterns
+                    # Note: This is advisory only, we don't block requests
+                    matches = self.pattern_matcher.match_prompt_injection(content)
+                    if matches:
+                        warnings.append(f"Potential prompt injection: {matches}")
+
+        return InterceptionResult(
+            allowed=True,  # Requests are allowed, just warned
+            provider=provider,
+            warnings=warnings,
+        )

tweek/proxy/server.py

@@ -0,0 +1,315 @@
+"""
+Proxy Server Management - Start, stop, and manage the Tweek proxy.
+
+This module handles the lifecycle of the mitmproxy-based security proxy.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import signal
+import subprocess
+import time
+from pathlib import Path
+from typing import Optional
+import logging
+
+logger = logging.getLogger("tweek.proxy")
+
+# Default configuration
+DEFAULT_PORT = 9877
+DEFAULT_WEB_PORT = 9878  # mitmproxy web interface
+PID_FILE = Path.home() / ".tweek" / "proxy" / "proxy.pid"
+LOG_FILE = Path.home() / ".tweek" / "proxy" / "proxy.log"
+CA_DIR = Path.home() / ".tweek" / "proxy" / "certs"
+
+
+def ensure_directories():
+    """Ensure proxy directories exist."""
+    PID_FILE.parent.mkdir(parents=True, exist_ok=True)
+    CA_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def get_addon_script_path() -> Path:
+    """Get the path to the mitmproxy addon script."""
+    return Path(__file__).parent / "addon.py"
+
+
+def is_proxy_running() -> tuple[bool, Optional[int]]:
+    """Check if the proxy is running and return its PID."""
+    if not PID_FILE.exists():
+        return False, None
+
+    try:
+        pid = int(PID_FILE.read_text().strip())
+
+        # Check if process exists
+        os.kill(pid, 0)
+        return True, pid
+    except (ValueError, ProcessLookupError, PermissionError):
+        # Clean up stale PID file
+        PID_FILE.unlink(missing_ok=True)
+        return False, None
+
+
+def start_proxy(
+    port: int = DEFAULT_PORT,
+    web_port: Optional[int] = None,
+    block_mode: bool = True,
+    log_only: bool = False,
+    foreground: bool = False,
+) -> tuple[bool, str]:
+    """
+    Start the Tweek proxy server.
+
+    Args:
+        port: Port for the proxy to listen on
+        web_port: Port for mitmproxy web interface (None to disable)
+        block_mode: If True, block dangerous responses
+        log_only: If True, log only without blocking
+        foreground: If True, run in foreground (for debugging)
+
+    Returns:
+        Tuple of (success, message)
+    """
+    # Check if already running
+    running, pid = is_proxy_running()
+    if running:
+        return False, f"Proxy already running (PID {pid})"
+
+    # Check for mitmproxy
+    try:
+        import mitmproxy
+    except ImportError:
+        return False, "mitmproxy not installed. Run: pip install tweek[proxy]"
+
+    ensure_directories()
+
+    # Build mitmproxy command
+    addon_path = get_addon_script_path()
+
+    cmd = [
+        sys.executable, "-m", "mitmproxy.tools.main",
+        "--mode", "regular",
+        "--listen-port", str(port),
+        "--set", f"confdir={CA_DIR}",
+        "--scripts", str(addon_path),
+        "--quiet",  # Reduce noise
+    ]
+
+    if web_port:
+        cmd.extend(["--web-port", str(web_port)])
+    else:
+        cmd.append("--no-web-open-browser")
+
+    # Set environment for addon configuration
+    env = os.environ.copy()
+    env["TWEEK_PROXY_BLOCK_MODE"] = "1" if block_mode else "0"
+    env["TWEEK_PROXY_LOG_ONLY"] = "1" if log_only else "0"
+
+    if foreground:
+        # Run in foreground for debugging
+        try:
+            subprocess.run(cmd, env=env)
+            return True, "Proxy stopped"
+        except KeyboardInterrupt:
+            return True, "Proxy stopped by user"
+    else:
+        # Run in background
+        with open(LOG_FILE, "a") as log:
+            process = subprocess.Popen(
+                cmd,
+                env=env,
+                stdout=log,
+                stderr=subprocess.STDOUT,
+                start_new_session=True,
+            )
+
+        # Save PID
+        PID_FILE.write_text(str(process.pid))
+
+        # Wait a moment to check if it started successfully
+        time.sleep(1)
+
+        running, _ = is_proxy_running()
+        if running:
+            return True, f"Proxy started on port {port} (PID {process.pid})"
+        else:
+            return False, f"Proxy failed to start. Check {LOG_FILE} for details"
+
+
+def stop_proxy() -> tuple[bool, str]:
+    """
+    Stop the Tweek proxy server.
+
+    Returns:
+        Tuple of (success, message)
+    """
+    running, pid = is_proxy_running()
+
+    if not running:
+        return False, "Proxy is not running"
+
+    try:
+        os.kill(pid, signal.SIGTERM)
+
+        # Wait for graceful shutdown
+        for _ in range(10):
+            time.sleep(0.5)
+            try:
+                os.kill(pid, 0)
+            except ProcessLookupError:
+                break
+        else:
+            # Force kill if still running
+            os.kill(pid, signal.SIGKILL)
+
+        PID_FILE.unlink(missing_ok=True)
+        return True, f"Proxy stopped (PID {pid})"
+
+    except ProcessLookupError:
+        PID_FILE.unlink(missing_ok=True)
+        return True, "Proxy was already stopped"
+    except PermissionError:
+        return False, f"Permission denied stopping PID {pid}"
+
+
+def get_proxy_info() -> dict:
+    """Get detailed proxy status information."""
+    running, pid = is_proxy_running()
+
+    info = {
+        "running": running,
+        "pid": pid,
+        "pid_file": str(PID_FILE),
+        "log_file": str(LOG_FILE),
+        "ca_dir": str(CA_DIR),
+        "ca_cert": str(CA_DIR / "mitmproxy-ca-cert.pem"),
+        "default_port": DEFAULT_PORT,
+    }
+
+    # Check if CA cert exists
+    ca_cert = CA_DIR / "mitmproxy-ca-cert.pem"
+    info["ca_cert_exists"] = ca_cert.exists()
+
+    return info
+
+
+def install_ca_certificate() -> tuple[bool, str]:
+    """
+    Install the Tweek proxy CA certificate in the system trust store.
+
+    Returns:
+        Tuple of (success, message)
+    """
+    import platform
+
+    ca_cert = CA_DIR / "mitmproxy-ca-cert.pem"
+
+    if not ca_cert.exists():
+        # Generate CA cert by starting proxy briefly
+        ensure_directories()
+        try:
+            import mitmproxy
+            from mitmproxy.certs import CertStore
+
+            # This will generate the CA if it doesn't exist
+            store = CertStore.from_store(str(CA_DIR), "mitmproxy", 2048)
+            # CA is now generated
+        except Exception as e:
+            return False, f"Failed to generate CA certificate: {e}"
+
+    if not ca_cert.exists():
+        return False, "CA certificate not found. Start the proxy first to generate it."
+
+    system = platform.system()
+
+    if system == "Darwin":
+        # macOS: Add to System Keychain
+        cmd = [
+            "sudo", "security", "add-trusted-cert",
+            "-d", "-r", "trustRoot",
+            "-k", "/Library/Keychains/System.keychain",
+            str(ca_cert)
+        ]
+        try:
+            subprocess.run(cmd, check=True)
+            return True, "CA certificate installed in macOS System Keychain"
+        except subprocess.CalledProcessError as e:
+            return False, f"Failed to install CA certificate: {e}"
+
+    elif system == "Linux":
+        # Linux: Copy to /usr/local/share/ca-certificates
+        dest = Path("/usr/local/share/ca-certificates/tweek-proxy.crt")
+        try:
+            subprocess.run(["sudo", "cp", str(ca_cert), str(dest)], check=True)
+            subprocess.run(["sudo", "update-ca-certificates"], check=True)
+            return True, "CA certificate installed in system trust store"
+        except subprocess.CalledProcessError as e:
+            return False, f"Failed to install CA certificate: {e}"
+
+    else:
+        return False, f"Automatic CA installation not supported on {system}. Please install manually: {ca_cert}"
+
+
+def get_proxy_env_vars(port: int = DEFAULT_PORT) -> dict[str, str]:
+    """
+    Get environment variables to configure applications to use the proxy.
+
+    Returns:
+        Dictionary of environment variables
+    """
+    proxy_url = f"http://127.0.0.1:{port}"
+
+    return {
+        "HTTP_PROXY": proxy_url,
+        "HTTPS_PROXY": proxy_url,
+        "http_proxy": proxy_url,
+        "https_proxy": proxy_url,
+        # For Node.js apps that don't respect standard env vars
+        "NODE_TLS_REJECT_UNAUTHORIZED": "0",  # Required for self-signed CA
+    }
+
+
+def generate_wrapper_script(
+    app_command: str,
+    port: int = DEFAULT_PORT,
+    output_path: Optional[Path] = None,
+) -> str:
+    """
+    Generate a wrapper script to run an application through the Tweek proxy.
+
+    Args:
+        app_command: The command to wrap (e.g., "npm start")
+        port: Proxy port
+        output_path: If provided, write script to this path
+
+    Returns:
+        The wrapper script content
+    """
+    env_vars = get_proxy_env_vars(port)
+
+    script = f"""#!/bin/bash
+# Tweek Proxy Wrapper Script
+# This script runs an application through the Tweek security proxy.
+
+# Ensure Tweek proxy is running
+if ! pgrep -f "tweek.*proxy" > /dev/null; then
+    echo "Starting Tweek proxy..."
+    tweek proxy start
+    sleep 2
+fi
+
+# Set proxy environment variables
+{chr(10).join(f'export {k}="{v}"' for k, v in env_vars.items())}
+
+# Run the application
+{app_command}
+"""
+
+    if output_path:
+        output_path.write_text(script)
+        output_path.chmod(0o755)
+
+    return script

tweek/sandbox/__init__.py

@@ -0,0 +1,71 @@
+"""
+Tweek Sandbox - Cross-platform command sandboxing.
+
+Provides isolated execution environments:
+- macOS: sandbox-exec (built-in)
+- Linux: firejail (optional, recommended) or bubblewrap
+- Windows: Not available
+"""
+
+from tweek.platform import PLATFORM, Platform, IS_MACOS, IS_LINUX, get_sandbox_tool
+
+# Import platform-specific implementations
+SANDBOX_AVAILABLE = False
+SANDBOX_TOOL = None
+
+if IS_MACOS:
+    try:
+        from .profile_generator import ProfileGenerator, SkillManifest
+        from .executor import SandboxExecutor, ExecutionResult
+        SANDBOX_AVAILABLE = True
+        SANDBOX_TOOL = "sandbox-exec"
+    except ImportError:
+        ProfileGenerator = None
+        SkillManifest = None
+        SandboxExecutor = None
+        ExecutionResult = None
+
+elif IS_LINUX:
+    try:
+        from .linux import LinuxSandbox, prompt_install_firejail, get_sandbox
+        _sandbox = get_sandbox()
+        if _sandbox and _sandbox.available:
+            SANDBOX_AVAILABLE = True
+            SANDBOX_TOOL = _sandbox.tool
+    except ImportError:
+        LinuxSandbox = None
+        prompt_install_firejail = None
+        get_sandbox = None
+
+# Keep macOS imports available for backwards compatibility
+try:
+    from .profile_generator import ProfileGenerator, SkillManifest
+    from .executor import SandboxExecutor, ExecutionResult
+except ImportError:
+    pass
+
+
+def get_sandbox_status() -> dict:
+    """Get sandbox availability status for current platform."""
+    tool = get_sandbox_tool()
+
+    return {
+        "available": SANDBOX_AVAILABLE,
+        "tool": SANDBOX_TOOL or tool,
+        "platform": PLATFORM.value,
+    }
+
+
+__all__ = [
+    "ProfileGenerator",
+    "SkillManifest",
+    "SandboxExecutor",
+    "ExecutionResult",
+    "SANDBOX_AVAILABLE",
+    "SANDBOX_TOOL",
+    "get_sandbox_status",
+]
+
+# Add Linux-specific exports if available
+if IS_LINUX:
+    __all__.extend(["LinuxSandbox", "prompt_install_firejail", "get_sandbox"])