cycode 3.15.3.dev8__py3-none-any.whl → 3.15.4.dev2__py3-none-any.whl

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

@@ -1,12 +1,16 @@
-"""
-Hook handlers for AI IDE events.
+"""Hook handlers for AI IDE events.
+
+Each handler receives a unified payload and policy, applies the scan + policy
+logic, and returns a canonical ``HookDecision``. ``scan_command`` translates
+that decision into the IDE-specific JSON response via ``IDE.build_hook_response``.
 
-Each handler receives a unified payload from an IDE, applies policy rules,
-and returns a response that either allows or blocks the action.
+Handlers are agent-agnostic by design — adding a new IDE doesn't require
+touching any handler in this module.
 """
 
 import json
 import os
+from dataclasses import dataclass
 from multiprocessing.pool import ThreadPool
 from multiprocessing.pool import TimeoutError as PoolTimeoutError
 from typing import Callable, Optional
@@ -14,9 +18,9 @@ from typing import Callable, Optional
 import typer
 
 from cycode.cli.apps.ai_guardrails.consts import PolicyMode
+from cycode.cli.apps.ai_guardrails.ides.base import HookDecision
 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
@@ -30,21 +34,17 @@ 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.
+HandlerFn = Callable[[typer.Context, AIHookPayload, dict], HookDecision]
 
-    Scans prompt text for secrets before it's sent to the AI model.
-    Returns {"continue": False} to block, {"continue": True} to allow.
-    """
+
+def handle_before_submit_prompt(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> HookDecision:
+    """Scan prompt text for secrets before it's sent to the AI model."""
     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={})
     if not get_policy_value(prompt_config, 'enabled', default=True):
         ai_client.create_event(payload, AiHookEventType.PROMPT, AIHookOutcome.ALLOWED)
-        return response_builder.allow_prompt()
+        return HookDecision.allow(AiHookEventType.PROMPT)
 
     mode = get_policy_value(policy, 'mode', default=PolicyMode.BLOCK)
     prompt = payload.prompt or ''
@@ -66,9 +66,9 @@ def handle_before_submit_prompt(ctx: typer.Context, payload: AIHookPayload, poli
             if action == PolicyMode.BLOCK and mode == PolicyMode.BLOCK:
                 outcome = AIHookOutcome.BLOCKED
                 user_message = f'{violation_summary}. Remove secrets before sending.'
-                return response_builder.deny_prompt(user_message)
+                return HookDecision.deny(AiHookEventType.PROMPT, user_message)
             outcome = AIHookOutcome.WARNED
-        return response_builder.allow_prompt()
+        return HookDecision.allow(AiHookEventType.PROMPT)
     except Exception as e:
         outcome = (
             AIHookOutcome.ALLOWED if get_policy_value(policy, 'fail_open', default=True) else AIHookOutcome.BLOCKED
@@ -87,21 +87,14 @@ def handle_before_submit_prompt(ctx: typer.Context, payload: AIHookPayload, poli
         )
 
 
-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.
-    """
+def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> HookDecision:
+    """Block sensitive paths and scan file content for secrets."""
     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={})
     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()
+        return HookDecision.allow(AiHookEventType.FILE_READ)
 
     mode = get_policy_value(policy, 'mode', default=PolicyMode.BLOCK)
     file_path = payload.file_path or ''
@@ -113,20 +106,19 @@ def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy:
     error_message = None
 
     try:
-        # Check path-based denylist first
         is_sensitive_path = is_denied_path(file_path, policy)
         if is_sensitive_path:
             block_reason = BlockReason.SENSITIVE_PATH
             if mode == PolicyMode.BLOCK and action == PolicyMode.BLOCK:
                 outcome = AIHookOutcome.BLOCKED
                 user_message = f'Cycode blocked sending {file_path} to the AI (sensitive path policy).'
-                return response_builder.deny_permission(
+                return HookDecision.deny(
+                    AiHookEventType.FILE_READ,
                     user_message,
                     'This file path is classified as sensitive; do not read/send it to the model.',
                 )
-            # Warn mode - if content scan is enabled, emit a separate event for the
+            # Warn mode: if content scan is enabled, emit a separate event for the
             # sensitive path so the finally block can independently track the scan result.
-            # If content scan is disabled, a single event (from finally) is enough.
             outcome = AIHookOutcome.WARNED
             if get_policy_value(file_read_config, 'scan_content', default=True):
                 ai_client.create_event(
@@ -136,11 +128,9 @@ def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy:
                     block_reason=BlockReason.SENSITIVE_PATH,
                     file_path=payload.file_path,
                 )
-                # Reset for the content scan result tracked by the finally block
                 block_reason = None
                 outcome = AIHookOutcome.ALLOWED
 
-        # 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:
@@ -148,27 +138,28 @@ def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy:
                 if mode == PolicyMode.BLOCK and action == PolicyMode.BLOCK:
                     outcome = AIHookOutcome.BLOCKED
                     user_message = f'Cycode blocked reading {file_path}. {violation_summary}'
-                    return response_builder.deny_permission(
+                    return HookDecision.deny(
+                        AiHookEventType.FILE_READ,
                         user_message,
                         'Secrets detected; do not send this file to the model.',
                     )
-                # Warn mode - ask user for permission
                 outcome = AIHookOutcome.WARNED
                 user_message = f'Cycode detected secrets in {file_path}. {violation_summary}'
-                return response_builder.ask_permission(
+                return HookDecision.ask(
+                    AiHookEventType.FILE_READ,
                     user_message,
                     'Possible secrets detected; proceed with caution.',
                 )
 
-        # If path was sensitive but content scan found no secrets (or scan disabled), still warn
         if is_sensitive_path:
             user_message = f'Cycode flagged {file_path} as sensitive. Allow reading?'
-            return response_builder.ask_permission(
+            return HookDecision.ask(
+                AiHookEventType.FILE_READ,
                 user_message,
                 'This file path is classified as sensitive; proceed with caution.',
             )
 
-        return response_builder.allow_permission()
+        return HookDecision.allow(AiHookEventType.FILE_READ)
     except Exception as e:
         outcome = (
             AIHookOutcome.ALLOWED if get_policy_value(policy, 'fail_open', default=True) else AIHookOutcome.BLOCKED
@@ -188,31 +179,44 @@ def handle_before_read_file(ctx: typer.Context, payload: AIHookPayload, policy:
         )
 
 
-def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> dict:
-    """
-    Handle beforeMCPExecution hook.
+@dataclass(frozen=True)
+class _ArgScanFeature:
+    """Configuration for a "scan some text and decide" event.
 
-    Scans tool arguments for secrets before MCP tool execution.
-    Returns {"permission": "deny"} to block, {"permission": "ask"} to warn,
-    {"permission": "allow"} to allow.
+    MCP execution and command exec share identical scan-and-decide logic;
+    only the policy key, event type, and user-facing messages differ.
     """
+
+    policy_key: str  # 'mcp' or 'command_exec'
+    scan_key: str  # 'scan_arguments' or 'scan_command'
+    event_type: AiHookEventType
+    block_reason: BlockReason
+    deny_message: Callable[[str], str]
+    deny_agent_message: str
+    ask_message: Callable[[str], str]
+    ask_agent_message: str
+
+
+def _handle_arg_scan(
+    ctx: typer.Context,
+    payload: AIHookPayload,
+    policy: dict,
+    feature: _ArgScanFeature,
+    scan_text: str,
+) -> HookDecision:
+    """Shared scan + decision flow for MCP_EXECUTION and COMMAND_EXEC events."""
     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={})
-    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()
+    feature_config = get_policy_value(policy, feature.policy_key, default={})
+    if not get_policy_value(feature_config, 'enabled', default=True):
+        ai_client.create_event(payload, feature.event_type, AIHookOutcome.ALLOWED)
+        return HookDecision.allow(feature.event_type)
 
     mode = get_policy_value(policy, 'mode', default=PolicyMode.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=PolicyMode.BLOCK)
+    clipped = truncate_utf8(scan_text, max_bytes)
+    action = get_policy_value(feature_config, 'action', default=PolicyMode.BLOCK)
 
     scan_id = None
     block_reason = None
@@ -220,24 +224,25 @@ def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, poli
     error_message = None
 
     try:
-        if get_policy_value(mcp_config, 'scan_arguments', default=True):
+        if get_policy_value(feature_config, feature.scan_key, default=True):
             violation_summary, scan_id = _scan_text_for_secrets(ctx, clipped, timeout_ms)
             if violation_summary:
-                block_reason = BlockReason.SECRETS_IN_MCP_ARGS
+                block_reason = feature.block_reason
                 if mode == PolicyMode.BLOCK and action == PolicyMode.BLOCK:
                     outcome = AIHookOutcome.BLOCKED
-                    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.',
+                    return HookDecision.deny(
+                        feature.event_type,
+                        feature.deny_message(violation_summary),
+                        feature.deny_agent_message,
                     )
                 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 HookDecision.ask(
+                    feature.event_type,
+                    feature.ask_message(violation_summary),
+                    feature.ask_agent_message,
                 )
 
-        return response_builder.allow_permission()
+        return HookDecision.allow(feature.event_type)
     except Exception as e:
         outcome = (
             AIHookOutcome.ALLOWED if get_policy_value(policy, 'fail_open', default=True) else AIHookOutcome.BLOCKED
@@ -248,7 +253,7 @@ def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, poli
     finally:
         ai_client.create_event(
             payload,
-            AiHookEventType.MCP_EXECUTION,
+            feature.event_type,
             outcome,
             scan_id=scan_id,
             block_reason=block_reason,
@@ -256,16 +261,32 @@ def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, poli
         )
 
 
-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.
+def handle_before_mcp_execution(ctx: typer.Context, payload: AIHookPayload, policy: dict) -> HookDecision:
+    """Scan MCP tool arguments for secrets before execution."""
+    tool = payload.mcp_tool_name or 'unknown'
+    args = payload.mcp_arguments or {}
+    args_text = args if isinstance(args, str) else json.dumps(args)
+    return _handle_arg_scan(
+        ctx,
+        payload,
+        policy,
+        _ArgScanFeature(
+            policy_key='mcp',
+            scan_key='scan_arguments',
+            event_type=AiHookEventType.MCP_EXECUTION,
+            block_reason=BlockReason.SECRETS_IN_MCP_ARGS,
+            deny_message=lambda v: f'Cycode blocked MCP tool call "{tool}". {v}',
+            deny_agent_message='Do not pass secrets to tools. Use secret references (name/id) instead.',
+            ask_message=lambda v: f'{v} in MCP tool call "{tool}". Allow execution?',
+            ask_agent_message='Possible secrets detected in tool arguments; proceed with caution.',
+        ),
+        scan_text=args_text,
+    )
 
-    Args:
-        event_type: Canonical event type string (from AiHookEventType enum)
 
-    Returns:
-        Handler function or None if event type is not recognized
-    """
-    handlers = {
+def get_handler_for_event(event_type: str) -> Optional[HandlerFn]:
+    """Look up the handler for a canonical event type."""
+    handlers: dict[str, HandlerFn] = {
         AiHookEventType.PROMPT.value: handle_before_submit_prompt,
         AiHookEventType.FILE_READ.value: handle_before_read_file,
         AiHookEventType.MCP_EXECUTION.value: handle_before_mcp_execution,
@@ -275,32 +296,24 @@ def get_handler_for_event(event_type: str) -> Optional[Callable[[typer.Context,
 
 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.obj['sync'] = True
+    ctx.obj['scan_type'] = ScanTypeOption.SECRET
+    ctx.obj['severity_threshold'] = SeverityOption.INFO
     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.
+    """Run a scan on documents, returning (violation_summary, scan_id).
 
-    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).
+    Raises on scan failure / timeout so the fail-open policy can take over.
     """
     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
     )
@@ -324,7 +337,6 @@ def _perform_scan(
 
     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
@@ -333,12 +345,7 @@ def _perform_scan(
 
 
 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.
-    """
+    """Scan text content for secrets using Cycode CLI."""
     if not text:
         return None, None
 
@@ -349,12 +356,7 @@ def _scan_text_for_secrets(ctx: typer.Context, text: str, timeout_ms: int) -> tu
 
 
 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.
-    """
+    """Scan a file path for secrets."""
     if not file_path or not os.path.isfile(file_path):
         return None, None
 
@@ -363,7 +365,6 @@ def _scan_path_for_secrets(ctx: typer.Context, file_path: str, policy: dict) ->
     with open(file_path, encoding='utf-8', errors='replace') as f:
         content = f.read(max_bytes)
 
-    # Get timeout from policy
     timeout_ms = get_policy_value(policy, 'secrets', 'timeout_ms', default=30000)
     timeout_seconds = timeout_ms / 1000.0