cleveragents-core/robot/helper_security_eval.py

"""Helper utilities for security hardening Robot integration tests.

Covers SEC1: eval()/exec() removal from stream_router.py.
- Named operation registry (SimpleToolAgent._SAFE_OPERATIONS)
- Named transform registry (ReactiveStreamRouter._TRANSFORM_REGISTRY)
- Rejection of code blocks and unregistered transforms
"""

from __future__ import annotations

import sys

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


def _named_operations_work() -> None:
    """Integration test: registered named operations execute correctly."""
    # Operations are configured via tool dicts, not message metadata
    agent_upper = SimpleToolAgent(tools=[{"operation": "uppercase"}])
    result = agent_upper.process("hello world", {}, {})
    assert result == "HELLO WORLD", f"Expected 'HELLO WORLD', got {result}"

    agent_strip = SimpleToolAgent(tools=[{"operation": "strip"}])
    result2 = agent_strip.process("  HELLO  ", {}, {})
    assert result2 == "HELLO", f"Expected 'HELLO', got {result2}"

    agent_lower = SimpleToolAgent(tools=[{"operation": "lowercase"}])
    result3 = agent_lower.process("HELLO", {}, {})
    assert result3 == "hello", f"Expected 'hello', got {result3}"

    print("named-operations-ok")


def _code_blocks_rejected() -> None:
    """Integration test: code blocks are rejected (no eval/exec)."""
    # Code blocks are in the tool config, not metadata
    agent = SimpleToolAgent(tools=[{"code": "import os; os.system('echo pwned')"}])
    try:
        agent.process("test", {}, {})
        print("FAIL: Expected StreamRoutingError for code block")
        return
    except StreamRoutingError:
        pass

    # Attempt with import injection
    agent2 = SimpleToolAgent(tools=[{"code": "__import__('subprocess').call(['ls'])"}])
    try:
        agent2.process("test", {}, {})
        print("FAIL: Expected StreamRoutingError for import code block")
        return
    except StreamRoutingError:
        pass

    print("code-blocks-rejected-ok")


def _custom_operation_registration() -> None:
    """Integration test: custom operations can be registered and used."""
    # Register a custom operation (takes content, meta, ctx)
    SimpleToolAgent.register_operation(
        "reverse", lambda content, meta, ctx: content[::-1]
    )

    agent = SimpleToolAgent(tools=[{"operation": "reverse"}])
    result = agent.process("hello", {}, {})
    assert result == "olleh", f"Expected 'olleh', got {result}"

    print("custom-operation-ok")


def _unregistered_transform_rejected() -> None:
    """Integration test: unregistered transform expressions are rejected."""
    router = ReactiveStreamRouter()

    # Create a stream with an unregistered fn (lambda expression)
    config = StreamConfig(
        name="test_stream",
        type=StreamType.COLD,
        operators=[
            {
                "type": "transform",
                "params": {"fn": "lambda x: x.upper()"},
            }
        ],
    )

    # This should raise StreamRoutingError during stream creation
    try:
        router.create_stream(config)
        print("FAIL: Expected StreamRoutingError for unregistered transform")
        return
    except StreamRoutingError:
        pass
    except Exception:
        # Any error rejecting unregistered transforms is acceptable
        pass

    print("unregistered-transform-rejected-ok")


def _registered_transform_works() -> None:
    """Integration test: registered transforms work correctly."""
    # Register a named transform
    ReactiveStreamRouter.register_transform("shout", lambda x: x.upper() + "!")

    # Verify it's in the registry
    assert "shout" in ReactiveStreamRouter._TRANSFORM_REGISTRY
    assert ReactiveStreamRouter._TRANSFORM_REGISTRY["shout"]("hello") == "HELLO!"

    print("registered-transform-ok")


def _identity_operation_default() -> None:
    """Integration test: identity is a valid named operation."""
    agent = SimpleToolAgent(tools=[{"operation": "identity"}])
    result = agent.process("unchanged", {}, {})
    assert result == "unchanged", f"Expected 'unchanged', got {result}"

    # No tools means passthrough
    agent_empty = SimpleToolAgent(tools=[])
    result2 = agent_empty.process("unchanged", {}, {})
    assert result2 == "unchanged", f"Expected 'unchanged', got {result2}"

    print("identity-default-ok")


def main() -> None:
    if len(sys.argv) < 2:
        raise SystemExit("Expected command argument")
    command = sys.argv[1]
    commands = {
        "named-operations": _named_operations_work,
        "code-blocks-rejected": _code_blocks_rejected,
        "custom-operation": _custom_operation_registration,
        "unregistered-transform": _unregistered_transform_rejected,
        "registered-transform": _registered_transform_works,
        "identity-default": _identity_operation_default,
    }
    if command not in commands:
        raise SystemExit(f"Unknown command: {command}")
    commands[command]()


if __name__ == "__main__":
    main()