test_handlers.py

"""Tests for AI guardrails handlers."""

from typing import Any
from unittest.mock import MagicMock, patch

import pytest
import typer

from cycode.cli.apps.ai_guardrails.scan.handlers import (
    handle_before_mcp_execution,
    handle_before_read_file,
    handle_before_submit_prompt,
)
from cycode.cli.apps.ai_guardrails.scan.payload import AIHookPayload
from cycode.cli.apps.ai_guardrails.scan.types import AIHookOutcome, BlockReason


@pytest.fixture
def mock_ctx() -> MagicMock:
    """Create a mock Typer context."""
    ctx = MagicMock(spec=typer.Context)
    ctx.obj = {
        'ai_security_client': MagicMock(),
        'scan_type': 'secret',
    }
    return ctx


@pytest.fixture
def mock_payload() -> AIHookPayload:
    """Create a mock AIHookPayload."""
    return AIHookPayload(
        event_name='prompt',
        conversation_id='test-conv-id',
        generation_id='test-gen-id',
        ide_user_email='test@example.com',
        model='gpt-4',
        ide_provider='cursor',
        ide_version='1.0.0',
        prompt='Test prompt',
    )


@pytest.fixture
def default_policy() -> dict[str, Any]:
    """Create a default policy dict."""
    return {
        'mode': 'block',
        'fail_open': True,
        'secrets': {'max_bytes': 200000},
        'prompt': {'enabled': True, 'action': 'block'},
        'file_read': {'enabled': True, 'action': 'block', 'scan_content': True, 'deny_globs': []},
        'mcp': {'enabled': True, 'action': 'block', 'scan_arguments': True},
    }


# Tests for handle_before_submit_prompt


def test_handle_before_submit_prompt_disabled(
    mock_ctx: MagicMock, mock_payload: AIHookPayload, default_policy: dict[str, Any]
) -> None:
    """Test that disabled prompt scanning allows the prompt."""
    default_policy['prompt']['enabled'] = False

    result = handle_before_submit_prompt(mock_ctx, mock_payload, default_policy)

    assert result == {'continue': True}
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    mock_ctx.obj['ai_security_client'].create_conversation.assert_not_called()


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_submit_prompt_no_secrets(
    mock_scan: MagicMock, mock_ctx: MagicMock, mock_payload: AIHookPayload, default_policy: dict[str, Any]
) -> None:
    """Test that prompt with no secrets is allowed."""
    mock_scan.return_value = (None, 'scan-id-123')

    result = handle_before_submit_prompt(mock_ctx, mock_payload, default_policy)

    assert result == {'continue': True}
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    mock_ctx.obj['ai_security_client'].create_conversation.assert_not_called()
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    # outcome is arg[2], scan_id and block_reason are kwargs
    assert call_args.args[2] == AIHookOutcome.ALLOWED
    assert call_args.kwargs['scan_id'] == 'scan-id-123'
    assert call_args.kwargs['block_reason'] is None


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_submit_prompt_with_secrets_blocked(
    mock_scan: MagicMock, mock_ctx: MagicMock, mock_payload: AIHookPayload, default_policy: dict[str, Any]
) -> None:
    """Test that prompt with secrets is blocked."""
    mock_scan.return_value = ('Found 1 secret: API key', 'scan-id-456')

    result = handle_before_submit_prompt(mock_ctx, mock_payload, default_policy)

    assert result['continue'] is False
    assert 'Found 1 secret: API key' in result['user_message']
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.BLOCKED
    assert call_args.kwargs['block_reason'] == BlockReason.SECRETS_IN_PROMPT


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_submit_prompt_with_secrets_warned(
    mock_scan: MagicMock, mock_ctx: MagicMock, mock_payload: AIHookPayload, default_policy: dict[str, Any]
) -> None:
    """Test that prompt with secrets in warn mode is allowed."""
    default_policy['prompt']['action'] = 'warn'
    mock_scan.return_value = ('Found 1 secret: API key', 'scan-id-789')

    result = handle_before_submit_prompt(mock_ctx, mock_payload, default_policy)

    assert result == {'continue': True}
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.WARNED


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_submit_prompt_scan_failure_fail_open(
    mock_scan: MagicMock, mock_ctx: MagicMock, mock_payload: AIHookPayload, default_policy: dict[str, Any]
) -> None:
    """Test that scan failure with fail_open=True allows the prompt."""
    mock_scan.side_effect = RuntimeError('Scan failed')
    default_policy['fail_open'] = True

    with pytest.raises(RuntimeError):
        handle_before_submit_prompt(mock_ctx, mock_payload, default_policy)

    # Event should be tracked even on exception
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.ALLOWED
    # block_reason is set for tracking even when fail_open allows the action
    assert call_args.kwargs['block_reason'] == BlockReason.SCAN_FAILURE


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_submit_prompt_scan_failure_fail_closed(
    mock_scan: MagicMock, mock_ctx: MagicMock, mock_payload: AIHookPayload, default_policy: dict[str, Any]
) -> None:
    """Test that scan failure with fail_open=False blocks the prompt."""
    mock_scan.side_effect = RuntimeError('Scan failed')
    default_policy['fail_open'] = False

    with pytest.raises(RuntimeError):
        handle_before_submit_prompt(mock_ctx, mock_payload, default_policy)

    # Event should be tracked even on exception
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.BLOCKED
    assert call_args.kwargs['block_reason'] == BlockReason.SCAN_FAILURE


# Tests for handle_before_read_file


def test_handle_before_read_file_disabled(mock_ctx: MagicMock, default_policy: dict[str, Any]) -> None:
    """Test that disabled file read scanning allows the file."""
    default_policy['file_read']['enabled'] = False
    payload = AIHookPayload(
        event_name='file_read',
        ide_provider='cursor',
        file_path='/path/to/file.txt',
    )

    result = handle_before_read_file(mock_ctx, payload, default_policy)

    assert result == {'permission': 'allow'}


@patch('cycode.cli.apps.ai_guardrails.scan.handlers.is_denied_path')
def test_handle_before_read_file_sensitive_path(
    mock_is_denied: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that sensitive path is blocked."""
    mock_is_denied.return_value = True
    payload = AIHookPayload(
        event_name='file_read',
        ide_provider='cursor',
        file_path='/path/to/.env',
    )

    result = handle_before_read_file(mock_ctx, payload, default_policy)

    assert result['permission'] == 'deny'
    assert '.env' in result['user_message']
    mock_ctx.obj['ai_security_client'].create_event.assert_called_once()
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.BLOCKED
    assert call_args.kwargs['block_reason'] == BlockReason.SENSITIVE_PATH
    assert call_args.kwargs['file_path'] == '/path/to/.env'


@patch('cycode.cli.apps.ai_guardrails.scan.handlers.is_denied_path')
@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_path_for_secrets')
def test_handle_before_read_file_no_secrets(
    mock_scan: MagicMock, mock_is_denied: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that file with no secrets is allowed."""
    mock_is_denied.return_value = False
    mock_scan.return_value = (None, 'scan-id-123')
    payload = AIHookPayload(
        event_name='file_read',
        ide_provider='cursor',
        file_path='/path/to/file.txt',
    )

    result = handle_before_read_file(mock_ctx, payload, default_policy)

    assert result == {'permission': 'allow'}
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.ALLOWED
    assert call_args.kwargs['file_path'] == '/path/to/file.txt'


@patch('cycode.cli.apps.ai_guardrails.scan.handlers.is_denied_path')
@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_path_for_secrets')
def test_handle_before_read_file_with_secrets(
    mock_scan: MagicMock, mock_is_denied: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that file with secrets is blocked."""
    mock_is_denied.return_value = False
    mock_scan.return_value = ('Found 1 secret: password', 'scan-id-456')
    payload = AIHookPayload(
        event_name='file_read',
        ide_provider='cursor',
        file_path='/path/to/file.txt',
    )

    result = handle_before_read_file(mock_ctx, payload, default_policy)

    assert result['permission'] == 'deny'
    assert 'Found 1 secret: password' in result['user_message']
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.BLOCKED
    assert call_args.kwargs['block_reason'] == BlockReason.SECRETS_IN_FILE
    assert call_args.kwargs['file_path'] == '/path/to/file.txt'


@patch('cycode.cli.apps.ai_guardrails.scan.handlers.is_denied_path')
@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_path_for_secrets')
def test_handle_before_read_file_scan_disabled(
    mock_scan: MagicMock, mock_is_denied: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that file is allowed when content scanning is disabled."""
    mock_is_denied.return_value = False
    default_policy['file_read']['scan_content'] = False
    payload = AIHookPayload(
        event_name='file_read',
        ide_provider='cursor',
        file_path='/path/to/file.txt',
    )

    result = handle_before_read_file(mock_ctx, payload, default_policy)

    assert result == {'permission': 'allow'}
    mock_scan.assert_not_called()


# Tests for handle_before_mcp_execution


def test_handle_before_mcp_execution_disabled(mock_ctx: MagicMock, default_policy: dict[str, Any]) -> None:
    """Test that disabled MCP scanning allows the execution."""
    default_policy['mcp']['enabled'] = False
    payload = AIHookPayload(
        event_name='mcp_execution',
        ide_provider='cursor',
        mcp_tool_name='test_tool',
        mcp_arguments={'arg1': 'value1'},
    )

    result = handle_before_mcp_execution(mock_ctx, payload, default_policy)

    assert result == {'permission': 'allow'}


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_mcp_execution_no_secrets(
    mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that MCP execution with no secrets is allowed."""
    mock_scan.return_value = (None, 'scan-id-123')
    payload = AIHookPayload(
        event_name='mcp_execution',
        ide_provider='cursor',
        mcp_tool_name='test_tool',
        mcp_arguments={'arg1': 'value1'},
    )

    result = handle_before_mcp_execution(mock_ctx, payload, default_policy)

    assert result == {'permission': 'allow'}
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.ALLOWED


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_mcp_execution_with_secrets_blocked(
    mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that MCP execution with secrets is blocked."""
    mock_scan.return_value = ('Found 1 secret: token', 'scan-id-456')
    payload = AIHookPayload(
        event_name='mcp_execution',
        ide_provider='cursor',
        mcp_tool_name='test_tool',
        mcp_arguments={'arg1': 'secret_token_12345'},
    )

    result = handle_before_mcp_execution(mock_ctx, payload, default_policy)

    assert result['permission'] == 'deny'
    assert 'Found 1 secret: token' in result['user_message']
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.BLOCKED
    assert call_args.kwargs['block_reason'] == BlockReason.SECRETS_IN_MCP_ARGS


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_mcp_execution_with_secrets_warned(
    mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that MCP execution with secrets in warn mode asks permission."""
    mock_scan.return_value = ('Found 1 secret: token', 'scan-id-789')
    default_policy['mcp']['action'] = 'warn'
    payload = AIHookPayload(
        event_name='mcp_execution',
        ide_provider='cursor',
        mcp_tool_name='test_tool',
        mcp_arguments={'arg1': 'secret_token_12345'},
    )

    result = handle_before_mcp_execution(mock_ctx, payload, default_policy)

    assert result['permission'] == 'ask'
    assert 'Found 1 secret: token' in result['user_message']
    call_args = mock_ctx.obj['ai_security_client'].create_event.call_args
    assert call_args.args[2] == AIHookOutcome.WARNED


@patch('cycode.cli.apps.ai_guardrails.scan.handlers._scan_text_for_secrets')
def test_handle_before_mcp_execution_scan_disabled(
    mock_scan: MagicMock, mock_ctx: MagicMock, default_policy: dict[str, Any]
) -> None:
    """Test that MCP execution is allowed when argument scanning is disabled."""
    default_policy['mcp']['scan_arguments'] = False
    payload = AIHookPayload(
        event_name='mcp_execution',
        ide_provider='cursor',
        mcp_tool_name='test_tool',
        mcp_arguments={'arg1': 'value1'},
    )

    result = handle_before_mcp_execution(mock_ctx, payload, default_policy)

    assert result == {'permission': 'allow'}
    mock_scan.assert_not_called()