cleveragents-core/features/steps/security_eval_steps.py

"""Step definitions for SEC1 - eval/exec removal security tests."""

from __future__ import annotations

from typing import Any

import rx
from behave import given, then, when
from behave.runner import Context

from cleveragents.core.exceptions import StreamRoutingError
from cleveragents.reactive.stream_router import (
    ReactiveStreamRouter,
    SimpleToolAgent,
    StreamMessage,
)

# ---------------------------------------------------------------------------
# Scenario: Config with code injection attempt is rejected
# ---------------------------------------------------------------------------


@given("a SimpleToolAgent configured with a code injection payload")
def step_code_injection_agent(context: Context) -> None:
    context.tool_agent = SimpleToolAgent(
        tools=[{"code": "import os; os.system('echo pwned')"}]
    )


@when("I attempt to process content through the injected agent")
def step_attempt_process_injected(context: Context) -> None:
    context.tool_error = None
    try:
        context.tool_result = context.tool_agent.process("harmless input")
    except Exception as exc:  # pylint: disable=broad-except
        context.tool_error = exc


@then("the agent should reject the code block with a routing error")
def step_verify_code_block_rejected(context: Context) -> None:
    assert isinstance(context.tool_error, StreamRoutingError), (
        f"Expected StreamRoutingError, got {type(context.tool_error).__name__}: "
        f"{context.tool_error}"
    )


@then("the error message should mention code blocks are not supported")
def step_verify_code_blocks_message(context: Context) -> None:
    msg = str(context.tool_error).lower()
    assert "code" in msg and "not supported" in msg, (
        f"Error message did not mention code blocks: {context.tool_error}"
    )


# ---------------------------------------------------------------------------
# Scenario: Malicious transform expression does not execute
# ---------------------------------------------------------------------------


@when("I configure a transform with a malicious eval expression")
def step_malicious_eval_transform(context: Context) -> None:
    context.transform_error = None
    try:
        context.stream_router._create_operator(
            {
                "type": "transform",
                "params": {"fn": "__import__('os').system('echo pwned')"},
            }
        )
    except Exception as exc:  # pylint: disable=broad-except
        context.transform_error = exc


@then("the transform should be rejected with a routing error")
def step_verify_transform_rejected(context: Context) -> None:
    assert isinstance(context.transform_error, StreamRoutingError), (
        f"Expected StreamRoutingError, got "
        f"{type(context.transform_error).__name__}: {context.transform_error}"
    )


@then("the error message should mention unknown transform")
def step_verify_unknown_transform_message(context: Context) -> None:
    msg = str(context.transform_error).lower()
    assert "unknown transform" in msg, (
        f"Error message did not mention unknown transform: {context.transform_error}"
    )


# ---------------------------------------------------------------------------
# Scenario: Valid config still works after eval removal
# ---------------------------------------------------------------------------


@given('a SimpleToolAgent configured with a named operation "{op_name}"')
def step_named_operation_agent(context: Context, op_name: str) -> None:
    context.tool_agent = SimpleToolAgent(tools=[{"operation": op_name}])


@when('I process "{content}" through the named operation agent')
def step_process_named_operation(context: Context, content: str) -> None:
    context.tool_result = context.tool_agent.process(content)


@then('the named operation result should be "{expected}"')
def step_verify_named_operation_result(context: Context, expected: str) -> None:
    assert context.tool_result == expected, (
        f"Expected '{expected}', got '{context.tool_result}'"
    )


# ---------------------------------------------------------------------------
# Scenario: Named transform from registry works after eval removal
# ---------------------------------------------------------------------------


@when('I configure a transform with a registered named function "{fn_name}"')
def step_registered_transform(context: Context, fn_name: str) -> None:
    operator = context.stream_router._create_operator(
        {"type": "transform", "params": {"fn": fn_name}}
    )
    captured: list[Any] = []
    rx.just(StreamMessage(content="hello world")).pipe(operator).subscribe(
        captured.append
    )
    context.transform_result = captured[0] if captured else None


@then("the named transform should produce the uppercased result")
def step_verify_named_transform_uppercase(context: Context) -> None:
    assert context.transform_result == "HELLO WORLD", (
        f"Expected 'HELLO WORLD', got '{context.transform_result}'"
    )


# ---------------------------------------------------------------------------
# Scenario: Code block with import attempt is rejected
# ---------------------------------------------------------------------------


@given("a SimpleToolAgent configured with an import injection payload")
def step_import_injection_agent(context: Context) -> None:
    context.tool_agent = SimpleToolAgent(
        tools=[{"code": "__import__('subprocess').check_output(['id'])"}]
    )


# Reuses: step_attempt_process_injected, step_verify_code_block_rejected


# ---------------------------------------------------------------------------
# Scenario: Eval expression mimicking registry name is still rejected
# ---------------------------------------------------------------------------


@when('I configure a transform with expression "{expr}"')
def step_malicious_expression_transform(context: Context, expr: str) -> None:
    context.transform_error = None
    try:
        context.stream_router._create_operator(
            {"type": "transform", "params": {"fn": expr}}
        )
    except Exception as exc:  # pylint: disable=broad-except
        context.transform_error = exc


# Reuses: step_verify_transform_rejected


# ---------------------------------------------------------------------------
# Scenario: Custom registered operation works
# ---------------------------------------------------------------------------


@given('a SimpleToolAgent with a custom registered operation "{op_name}"')
def step_custom_operation_agent(context: Context, op_name: str) -> None:
    SimpleToolAgent.register_operation(
        op_name, lambda content, meta, ctx: str(content)[::-1]
    )
    context.tool_agent = SimpleToolAgent(tools=[{"operation": op_name}])

    def cleanup() -> None:
        SimpleToolAgent._SAFE_OPERATIONS.pop(op_name, None)

    add_cleanup = getattr(context, "add_cleanup", None)
    if callable(add_cleanup):
        add_cleanup(cleanup)


@when('I process "{content}" through the custom operation agent')
def step_process_custom_operation(context: Context, content: str) -> None:
    context.tool_result = context.tool_agent.process(content)


@then('the custom operation result should be "{expected}"')
def step_verify_custom_operation_result(context: Context, expected: str) -> None:
    assert context.tool_result == expected, (
        f"Expected '{expected}', got '{context.tool_result}'"
    )


# ---------------------------------------------------------------------------
# Scenario: Custom registered transform works
# ---------------------------------------------------------------------------


@given('a custom transform "{name}" is registered')
def step_register_custom_transform(context: Context, name: str) -> None:
    ReactiveStreamRouter.register_transform(name, lambda x: str(x) * 2)

    def cleanup() -> None:
        ReactiveStreamRouter._TRANSFORM_REGISTRY.pop(name, None)

    add_cleanup = getattr(context, "add_cleanup", None)
    if callable(add_cleanup):
        add_cleanup(cleanup)


@then("the custom transform should produce the doubled result")
def step_verify_custom_transform_doubled(context: Context) -> None:
    assert context.transform_result == "hello worldhello world", (
        f"Expected 'hello worldhello world', got '{context.transform_result}'"
    )