cycode 3.8.9.dev1__py3-none-any.whl → 3.8.11.dev1__py3-none-any.whl

cycode/cli/apps/ai_guardrails/scan/handlers.py

@@ -0,0 +1,341 @@
+"""
+Hook handlers for AI IDE events.
+
+Each handler receives a unified payload from an IDE, applies policy rules,
+and returns a response that either allows or blocks the action.
+"""
+
+import json
+import os
+from multiprocessing.pool import ThreadPool
+from multiprocessing.pool import TimeoutError as PoolTimeoutError
+from typing import Callable, Optional
+
+import typer
+
+from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload
+from cycode.cli.apps.ai_guardrails.scan.policy import get_policy_value
+from cycode.cli.apps.ai_guardrails.scan.response_builders import get_response_builder
+from cycode.cli.apps.ai_guardrails.scan.types import AiHookEventType, AIHookOutcome, BlockReason
+from cycode.cli.apps.ai_guardrails.scan.utils import is_denied_path, truncate_utf8
+from cycode.cli.apps.scan.code_scanner import _get_scan_documents_thread_func
+from cycode.cli.apps.scan.scan_parameters import get_scan_parameters
+from cycode.cli.cli_types import ScanTypeOption, SeverityOption
+from cycode.cli.models import Document
+from cycode.cli.utils.progress_bar import DummyProgressBar, ScanProgressBarSection
+from cycode.cli.utils.scan_utils import build_violation_summary
+from cycode.logger import get_logger
+
+logger = get_logger('AI Guardrails')
+
+
+def handle_before_submit_prompt(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> dict:
+    """
+    Handle beforeSubmitPrompt hook.
+
+    Scans prompt text for secrets before it's sent to the AI model.
+    Returns {"continue": False} to block, {"continue": True} to allow.
+    """
+    ai_client = ctx.obj['ai_security_client']
+    ide = payload.ide_provider
+    response_builder = get_response_builder(ide)
+
+    prompt_config = get_policy_value(policy, 'prompt', default={})
+    ai_client.create_conversation(payload)
+    if not get_policy_value(prompt_config, 'enabled', default=True):
+        ai_client.create_event(payload, AiHookEventType.PROMPT, AIHookOutcome.ALLOWED)
+        return response_builder.allow_prompt()
+
+    mode = get_policy_value(policy, 'mode', default='block')
+    prompt = payload.prompt or ''
+    max_bytes = get_policy_value(policy, 'secrets', 'max_bytes', default=200000)
+    timeout_ms = get_policy_value(policy, 'secrets', 'timeout_ms', default=30000)
+    clipped = truncate_utf8(prompt, max_bytes)
+
+    scan_id = None
+    block_reason = None
+    outcome = AIHookOutcome.ALLOWED
+
+    try:
+        violation_summary, scan_id = _scan_text_for_secrets(ctx, clipped, timeout_ms)
+
+        if (
+            violation_summary
+            and get_policy_value(prompt_config, 'action', default='block') == 'block'
+            and mode == 'block'
+        ):
+            outcome = AIHookOutcome.BLOCKED
+            block_reason = BlockReason.SECRETS_IN_PROMPT
+            user_message = f'{violation_summary}. Remove secrets before sending.'
+            response = response_builder.deny_prompt(user_message)
+        else:
+            if violation_summary:
+                outcome = AIHookOutcome.WARNED
+            response = response_builder.allow_prompt()
+        return response
+    except Exception as e:
+        outcome = (
+            AIHookOutcome.ALLOWED if get_policy_value(policy, 'fail_open', default=True) else AIHookOutcome.BLOCKED
+        )
+        block_reason = BlockReason.SCAN_FAILURE if outcome == AIHookOutcome.BLOCKED else None
+        raise e
+    finally:
+        ai_client.create_event(
+            payload,
+            AiHookEventType.PROMPT,
+            outcome,
+            scan_id=scan_id,
+            block_reason=block_reason,
+        )
+
+
+def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> dict:
+    """
+    Handle beforeReadFile hook.
+
+    Blocks sensitive files (via deny_globs) and scans file content for secrets.
+    Returns {"permission": "deny"} to block, {"permission": "allow"} to allow.
+    """
+    ai_client = ctx.obj['ai_security_client']
+    ide = payload.ide_provider
+    response_builder = get_response_builder(ide)
+
+    file_read_config = get_policy_value(policy, 'file_read', default={})
+    ai_client.create_conversation(payload)
+    if not get_policy_value(file_read_config, 'enabled', default=True):
+        ai_client.create_event(payload, AiHookEventType.FILE_READ, AIHookOutcome.ALLOWED)
+        return response_builder.allow_permission()
+
+    mode = get_policy_value(policy, 'mode', default='block')
+    file_path = payload.file_path or ''
+    action = get_policy_value(file_read_config, 'action', default='block')
+
+    scan_id = None
+    block_reason = None
+    outcome = AIHookOutcome.ALLOWED
+
+    try:
+        # Check path-based denylist first
+        if is_denied_path(file_path, policy) and action == 'block':
+            outcome = AIHookOutcome.BLOCKED
+            block_reason = BlockReason.SENSITIVE_PATH
+            user_message = f'Cycode blocked sending {file_path} to the AI (sensitive path policy).'
+            return response_builder.deny_permission(
+                user_message,
+                'This file path is classified as sensitive; do not read/send it to the model.',
+            )
+
+        # Scan file content if enabled
+        if get_policy_value(file_read_config, 'scan_content', default=True):
+            violation_summary, scan_id = _scan_path_for_secrets(ctx, file_path, policy)
+            if violation_summary and action == 'block' and mode == 'block':
+                outcome = AIHookOutcome.BLOCKED
+                block_reason = BlockReason.SECRETS_IN_FILE
+                user_message = f'Cycode blocked reading {file_path}. {violation_summary}'
+                return response_builder.deny_permission(
+                    user_message,
+                    'Secrets detected; do not send this file to the model.',
+                )
+            if violation_summary:
+                outcome = AIHookOutcome.WARNED
+            return response_builder.allow_permission()
+
+        return response_builder.allow_permission()
+    except Exception as e:
+        outcome = (
+            AIHookOutcome.ALLOWED if get_policy_value(policy, 'fail_open', default=True) else AIHookOutcome.BLOCKED
+        )
+        block_reason = BlockReason.SCAN_FAILURE if outcome == AIHookOutcome.BLOCKED else None
+        raise e
+    finally:
+        ai_client.create_event(
+            payload,
+            AiHookEventType.FILE_READ,
+            outcome,
+            scan_id=scan_id,
+            block_reason=block_reason,
+        )
+
+
+def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> dict:
+    """
+    Handle beforeMCPExecution hook.
+
+    Scans tool arguments for secrets before MCP tool execution.
+    Returns {"permission": "deny"} to block, {"permission": "ask"} to warn,
+    {"permission": "allow"} to allow.
+    """
+    ai_client = ctx.obj['ai_security_client']
+    ide = payload.ide_provider
+    response_builder = get_response_builder(ide)
+
+    mcp_config = get_policy_value(policy, 'mcp', default={})
+    ai_client.create_conversation(payload)
+    if not get_policy_value(mcp_config, 'enabled', default=True):
+        ai_client.create_event(payload, AiHookEventType.MCP_EXECUTION, AIHookOutcome.ALLOWED)
+        return response_builder.allow_permission()
+
+    mode = get_policy_value(policy, 'mode', default='block')
+    tool = payload.mcp_tool_name or 'unknown'
+    args = payload.mcp_arguments or {}
+    args_text = args if isinstance(args, str) else json.dumps(args)
+    max_bytes = get_policy_value(policy, 'secrets', 'max_bytes', default=200000)
+    timeout_ms = get_policy_value(policy, 'secrets', 'timeout_ms', default=30000)
+    clipped = truncate_utf8(args_text, max_bytes)
+    action = get_policy_value(mcp_config, 'action', default='block')
+
+    scan_id = None
+    block_reason = None
+    outcome = AIHookOutcome.ALLOWED
+
+    try:
+        if get_policy_value(mcp_config, 'scan_arguments', default=True):
+            violation_summary, scan_id = _scan_text_for_secrets(ctx, clipped, timeout_ms)
+            if violation_summary:
+                if mode == 'block' and action == 'block':
+                    outcome = AIHookOutcome.BLOCKED
+                    block_reason = BlockReason.SECRETS_IN_MCP_ARGS
+                    user_message = f'Cycode blocked MCP tool call "{tool}". {violation_summary}'
+                    return response_builder.deny_permission(
+                        user_message,
+                        'Do not pass secrets to tools. Use secret references (name/id) instead.',
+                    )
+                outcome = AIHookOutcome.WARNED
+                return response_builder.ask_permission(
+                    f'{violation_summary} in MCP tool call "{tool}". Allow execution?',
+                    'Possible secrets detected in tool arguments; proceed with caution.',
+                )
+
+        return response_builder.allow_permission()
+    except Exception as e:
+        outcome = (
+            AIHookOutcome.ALLOWED if get_policy_value(policy, 'fail_open', default=True) else AIHookOutcome.BLOCKED
+        )
+        block_reason = BlockReason.SCAN_FAILURE if outcome == AIHookOutcome.BLOCKED else None
+        raise e
+    finally:
+        ai_client.create_event(
+            payload,
+            AiHookEventType.MCP_EXECUTION,
+            outcome,
+            scan_id=scan_id,
+            block_reason=block_reason,
+        )
+
+
+def get_handler_for_event(event_type: str) -> Optional[Callable[[typer.Context, AIHookPayload, dict], dict]]:
+    """Get the appropriate handler function for a canonical event type.
+
+    Args:
+        event_type: Canonical event type string (from AiHookEventType enum)
+
+    Returns:
+        Handler function or None if event type is not recognized
+    """
+    handlers = {
+        AiHookEventType.PROMPT.value: handle_before_submit_prompt,
+        AiHookEventType.FILE_READ.value: handle_before_read_file,
+        AiHookEventType.MCP_EXECUTION.value: handle_before_mcp_execution,
+    }
+    return handlers.get(event_type)
+
+
+def _setup_scan_context(ctx: typer.Context) -> typer.Context:
+    """Set up minimal context for scan_documents without progress bars or printing."""
+
+    # Set up minimal required context
+    ctx.obj['progress_bar'] = DummyProgressBar([ScanProgressBarSection])
+    ctx.obj['sync'] = True  # Synchronous scan
+    ctx.obj['scan_type'] = ScanTypeOption.SECRET  # AI guardrails always scans for secrets
+    ctx.obj['severity_threshold'] = SeverityOption.INFO  # Report all severities
+
+    # Set command name for scan logic
+    ctx.info_name = 'ai_guardrails'
+
+    return ctx
+
+
+def _perform_scan(
+    ctx: typer.Context, documents: list[Document], scan_parameters: dict, timeout_seconds: float
+) -> tuple[Optional[str], Optional[str]]:
+    """
+    Perform a scan on documents and extract results.
+
+    Returns tuple of (violation_summary, scan_id) if secrets found, (None, scan_id) if clean.
+    Raises exception if scan fails or times out (triggers fail_open policy).
+    """
+    if not documents:
+        return None, None
+
+    # Get the thread function for scanning
+    scan_batch_thread_func = _get_scan_documents_thread_func(
+        ctx, is_git_diff=False, is_commit_range=False, scan_parameters=scan_parameters
+    )
+
+    # Use ThreadPool.apply_async with timeout to abort if scan takes too long
+    # This uses the same ThreadPool mechanism as run_parallel_batched_scan but with timeout support
+    with ThreadPool(processes=1) as pool:
+        result = pool.apply_async(scan_batch_thread_func, (documents,))
+        try:
+            scan_id, error, local_scan_result = result.get(timeout=timeout_seconds)
+        except PoolTimeoutError:
+            logger.debug('Scan timed out after %s seconds', timeout_seconds)
+            raise RuntimeError(f'Scan timed out after {timeout_seconds} seconds') from None
+
+    # Check if scan failed - raise exception to trigger fail_open policy
+    if error:
+        raise RuntimeError(error.message)
+
+    if not local_scan_result:
+        return None, None
+
+    scan_id = local_scan_result.scan_id
+
+    # Check if there are any detections
+    if local_scan_result.detections_count > 0:
+        violation_summary = build_violation_summary([local_scan_result])
+        return violation_summary, scan_id
+
+    return None, scan_id
+
+
+def _scan_text_for_secrets(ctx: typer.Context, text: str, timeout_ms: int) -> tuple[Optional[str], Optional[str]]:
+    """
+    Scan text content for secrets using Cycode CLI.
+
+    Returns tuple of (violation_summary, scan_id) if secrets found, (None, scan_id) if clean.
+    Raises exception on error or timeout.
+    """
+    if not text:
+        return None, None
+
+    document = Document(path='prompt-content.txt', content=text, is_git_diff_format=False)
+    scan_ctx = _setup_scan_context(ctx)
+    timeout_seconds = timeout_ms / 1000.0
+    return _perform_scan(scan_ctx, [document], get_scan_parameters(scan_ctx, None), timeout_seconds)
+
+
+def _scan_path_for_secrets(ctx: typer.Context, file_path: str, policy: dict) -> tuple[Optional[str], Optional[str]]:
+    """
+    Scan a file path for secrets.
+
+    Returns tuple of (violation_summary, scan_id) if secrets found, (None, scan_id) if clean.
+    Raises exception on error or timeout.
+    """
+    if not file_path or not os.path.exists(file_path):
+        return None, None
+
+    with open(file_path, encoding='utf-8', errors='replace') as f:
+        content = f.read()
+
+    # Truncate content based on policy max_bytes
+    max_bytes = get_policy_value(policy, 'secrets', 'max_bytes', default=200000)
+    content = truncate_utf8(content, max_bytes)
+
+    # Get timeout from policy
+    timeout_ms = get_policy_value(policy, 'secrets', 'timeout_ms', default=30000)
+    timeout_seconds = timeout_ms / 1000.0
+
+    document = Document(path=os.path.basename(file_path), content=content, is_git_diff_format=False)
+    scan_ctx = _setup_scan_context(ctx)
+    return _perform_scan(scan_ctx, [document], get_scan_parameters(scan_ctx, (file_path,)), timeout_seconds)

cycode/cli/apps/ai_guardrails/scan/payload.py

@@ -0,0 +1,72 @@
+"""Unified payload object for AI hook events from different tools."""
+
+from dataclasses import dataclass
+from typing import Optional
+
+from cycode.cli.apps.ai_guardrails.scan.types import CURSOR_EVENT_MAPPING
+
+
+@dataclass
+class AIHookPayload:
+    """Unified payload object that normalizes field names from different AI tools."""
+
+    # Event identification
+    event_name: str  # Canonical event type (e.g., 'prompt', 'file_read', 'mcp_execution')
+    conversation_id: Optional[str] = None
+    generation_id: Optional[str] = None
+
+    # User and IDE information
+    ide_user_email: Optional[str] = None
+    model: Optional[str] = None
+    ide_provider: str = None  # e.g., 'cursor', 'claude-code'
+    ide_version: Optional[str] = None
+
+    # Event-specific data
+    prompt: Optional[str] = None  # For prompt events
+    file_path: Optional[str] = None  # For file_read events
+    mcp_server_name: Optional[str] = None  # For mcp_execution events
+    mcp_tool_name: Optional[str] = None  # For mcp_execution events
+    mcp_arguments: Optional[dict] = None  # For mcp_execution events
+
+    @classmethod
+    def from_cursor_payload(cls, payload: dict) -> 'AIHookPayload':
+        """Create AIHookPayload from Cursor IDE payload.
+
+        Maps Cursor-specific event names to canonical event types.
+        """
+        cursor_event_name = payload.get('hook_event_name', '')
+        # Map Cursor event name to canonical type, fallback to original if not found
+        canonical_event = CURSOR_EVENT_MAPPING.get(cursor_event_name, cursor_event_name)
+
+        return cls(
+            event_name=canonical_event,
+            conversation_id=payload.get('conversation_id'),
+            generation_id=payload.get('generation_id'),
+            ide_user_email=payload.get('user_email'),
+            model=payload.get('model'),
+            ide_provider='cursor',
+            ide_version=payload.get('cursor_version'),
+            prompt=payload.get('prompt', ''),
+            file_path=payload.get('file_path') or payload.get('path'),
+            mcp_server_name=payload.get('command'),  # MCP server name
+            mcp_tool_name=payload.get('tool_name') or payload.get('tool'),
+            mcp_arguments=payload.get('arguments') or payload.get('tool_input') or payload.get('input'),
+        )
+
+    @classmethod
+    def from_payload(cls, payload: dict, tool: str = 'cursor') -> 'AIHookPayload':
+        """Create AIHookPayload from any tool's payload.
+
+        Args:
+            payload: The raw payload from the IDE
+            tool: The IDE/tool name (e.g., 'cursor')
+
+        Returns:
+            AIHookPayload instance
+
+        Raises:
+            ValueError: If the tool is not supported
+        """
+        if tool == 'cursor':
+            return cls.from_cursor_payload(payload)
+        raise ValueError(f'Unsupported IDE/tool: {tool}.')

cycode/cli/apps/ai_guardrails/scan/policy.py

@@ -0,0 +1,85 @@
+"""
+Policy loading and configuration management for AI guardrails.
+
+Policies are loaded and merged in order (later overrides earlier):
+1. Built-in defaults (consts.DEFAULT_POLICY)
+2. User-level config (~/.cycode/ai-guardrails.yaml)
+3. Repo-level config (<workspace>/.cycode/ai-guardrails.yaml)
+"""
+
+import json
+from pathlib import Path
+from typing import Any, Optional
+
+import yaml
+
+from cycode.cli.apps.ai_guardrails.scan.consts import DEFAULT_POLICY, POLICY_FILE_NAME
+
+
+def deep_merge(base: dict, override: dict) -> dict:
+    """Deep merge two dictionaries, with override taking precedence."""
+    result = base.copy()
+    for key, value in override.items():
+        if key in result and isinstance(result[key], dict) and isinstance(value, dict):
+            result[key] = deep_merge(result[key], value)
+        else:
+            result[key] = value
+    return result
+
+
+def load_yaml_file(path: Path) -> Optional[dict]:
+    """Load a YAML or JSON config file."""
+    if not path.exists():
+        return None
+    try:
+        content = path.read_text(encoding='utf-8')
+        if path.suffix in ('.yaml', '.yml'):
+            return yaml.safe_load(content)
+        return json.loads(content)
+    except Exception:
+        return None
+
+
+def load_defaults() -> dict:
+    """Load built-in defaults."""
+    return DEFAULT_POLICY.copy()
+
+
+def get_policy_value(policy: dict, *keys: str, default: Any = None) -> Any:
+    """Get a nested value from the policy dict."""
+    current = policy
+    for key in keys:
+        if not isinstance(current, dict):
+            return default
+        current = current.get(key)
+        if current is None:
+            return default
+    return current
+
+
+def load_policy(workspace_root: Optional[str] = None) -> dict:
+    """
+    Load policy by merging configs in order of precedence.
+
+    Merge order: defaults <- user config <- repo config
+
+    Args:
+        workspace_root: Workspace root path for repo-level config lookup.
+    """
+    # Start with defaults
+    policy = load_defaults()
+
+    # Merge user-level config (if exists)
+    user_policy_path = Path.home() / '.cycode' / POLICY_FILE_NAME
+    user_config = load_yaml_file(user_policy_path)
+    if user_config:
+        policy = deep_merge(policy, user_config)
+
+    # Merge repo-level config (if exists) - highest precedence
+    if workspace_root:
+        repo_policy_path = Path(workspace_root) / '.cycode' / POLICY_FILE_NAME
+        repo_config = load_yaml_file(repo_policy_path)
+        if repo_config:
+            policy = deep_merge(policy, repo_config)
+
+    return policy

cycode/cli/apps/ai_guardrails/scan/response_builders.py

@@ -0,0 +1,86 @@
+"""
+Response builders for different AI IDE hooks.
+
+Each IDE has its own response format for hooks. This module provides
+an abstract interface and concrete implementations for each supported IDE.
+"""
+
+from abc import ABC, abstractmethod
+
+
+class IDEResponseBuilder(ABC):
+    """Abstract base class for IDE-specific response builders."""
+
+    @abstractmethod
+    def allow_permission(self) -> dict:
+        """Build response to allow file read or MCP execution."""
+
+    @abstractmethod
+    def deny_permission(self, user_message: str, agent_message: str) -> dict:
+        """Build response to deny file read or MCP execution."""
+
+    @abstractmethod
+    def ask_permission(self, user_message: str, agent_message: str) -> dict:
+        """Build response to ask user for permission (warn mode)."""
+
+    @abstractmethod
+    def allow_prompt(self) -> dict:
+        """Build response to allow prompt submission."""
+
+    @abstractmethod
+    def deny_prompt(self, user_message: str) -> dict:
+        """Build response to deny prompt submission."""
+
+
+class CursorResponseBuilder(IDEResponseBuilder):
+    """Response builder for Cursor IDE hooks.
+
+    Cursor hook response formats:
+    - beforeSubmitPrompt: {"continue": bool, "user_message": str}
+    - beforeReadFile: {"permission": str, "user_message": str, "agent_message": str}
+    - beforeMCPExecution: {"permission": str, "user_message": str, "agent_message": str}
+    """
+
+    def allow_permission(self) -> dict:
+        """Allow file read or MCP execution."""
+        return {'permission': 'allow'}
+
+    def deny_permission(self, user_message: str, agent_message: str) -> dict:
+        """Deny file read or MCP execution."""
+        return {'permission': 'deny', 'user_message': user_message, 'agent_message': agent_message}
+
+    def ask_permission(self, user_message: str, agent_message: str) -> dict:
+        """Ask user for permission (warn mode)."""
+        return {'permission': 'ask', 'user_message': user_message, 'agent_message': agent_message}
+
+    def allow_prompt(self) -> dict:
+        """Allow prompt submission."""
+        return {'continue': True}
+
+    def deny_prompt(self, user_message: str) -> dict:
+        """Deny prompt submission."""
+        return {'continue': False, 'user_message': user_message}
+
+
+# Registry of response builders by IDE name
+_RESPONSE_BUILDERS: dict[str, IDEResponseBuilder] = {
+    'cursor': CursorResponseBuilder(),
+}
+
+
+def get_response_builder(ide: str = 'cursor') -> IDEResponseBuilder:
+    """Get the response builder for a specific IDE.
+
+    Args:
+        ide: The IDE name (e.g., 'cursor', 'claude-code')
+
+    Returns:
+        IDEResponseBuilder instance for the specified IDE
+
+    Raises:
+        ValueError: If the IDE is not supported
+    """
+    builder = _RESPONSE_BUILDERS.get(ide.lower())
+    if not builder:
+        raise ValueError(f'Unsupported IDE: {ide}. Supported IDEs: {list(_RESPONSE_BUILDERS.keys())}')
+    return builder