diff --git a/python/samples/concepts/README.md b/python/samples/concepts/README.md index bfde792aed39..e500b74829d0 100644 --- a/python/samples/concepts/README.md +++ b/python/samples/concepts/README.md @@ -150,6 +150,7 @@ ### Filtering - Creating and using Filters - [Auto Function Invoke Filters](./filtering/auto_function_invoke_filters.py) +- [Command Guard Filter](./filtering/command_guard_filter.py) - [Function Invocation Filters](./filtering/function_invocation_filters.py) - [Function Invocation Filters Stream](./filtering/function_invocation_filters_stream.py) - [Prompt Filters](./filtering/prompt_filters.py) diff --git a/python/samples/concepts/filtering/command_guard_filter.py b/python/samples/concepts/filtering/command_guard_filter.py new file mode 100644 index 000000000000..a9ab5c81f6e4 --- /dev/null +++ b/python/samples/concepts/filtering/command_guard_filter.py @@ -0,0 +1,125 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import inspect +import logging +from collections.abc import Awaitable, Callable, Coroutine, Iterable +from dataclasses import dataclass +from typing import Any, Protocol + +from semantic_kernel import Kernel +from semantic_kernel.filters import FilterTypes, FunctionInvocationContext +from semantic_kernel.functions import FunctionResult, kernel_function + +logger = logging.getLogger(__name__) + +# This sample shows how to put a guard in front of tools that run commands, such as a shell tool, +# using the existing function invocation filter. The guard is vendor neutral: it only needs an object +# with a `scan(command)` method that returns a verdict. That object can be a simple deny list, a policy +# engine, a call into an OS level sandbox, or a third party guardrail library. +# +# The filter runs before the function. If the scanner does not allow the command, the filter does not call +# `next`, so the function never runs, and it puts a short message in `context.result` instead. +# The same pattern works for the auto function invocation filter (FilterTypes.AUTO_FUNCTION_INVOCATION). +# In that case set `context.function_result` instead of `context.result`. +# +# Important: a check on the command text is a safety net, not a security boundary. For untrusted code, also +# run the tool inside real isolation (a container, a VM, or OS level sandboxing). + + +@dataclass +class ScanVerdict: + """The decision of a scanner for one command.""" + + allowed: bool + reason: str = "" + + +class CommandScanner(Protocol): + """Anything that can decide if a command is allowed to run. + + `scan` can be a normal method or a coroutine function. + """ + + def scan(self, command: str) -> ScanVerdict | Awaitable[ScanVerdict]: ... + + +def create_command_guard_filter( + scanner: CommandScanner, + argument_names: Iterable[str] = ("command",), +) -> Callable[[FunctionInvocationContext, Callable[[FunctionInvocationContext], Coroutine[Any, Any, None]]], Any]: + """Create a function invocation filter that checks command arguments with the scanner before the call. + + Args: + scanner: The scanner that decides if a command is allowed. + argument_names: The names of the function arguments that hold a command to check. + """ + names = tuple(argument_names) + + async def command_guard_filter( + context: FunctionInvocationContext, + next: Callable[[FunctionInvocationContext], Coroutine[Any, Any, None]], + ) -> None: + for name in names: + value = context.arguments.get(name) + if value is None: + continue + verdict = scanner.scan(str(value)) + if inspect.isawaitable(verdict): + verdict = await verdict + if not verdict.allowed: + # Only log the function name, the command itself may contain secrets. + logger.warning(f"Blocked call to {context.function.fully_qualified_name}: {verdict.reason}") + context.result = FunctionResult( + function=context.function.metadata, + value=f"Blocked by command guard: {verdict.reason or 'not allowed'}", + ) + # do not call next, so the function does not run + return + await next(context) + + return command_guard_filter + + +class DenyListScanner: + """A toy scanner for this sample: it blocks commands that contain one of the given fragments. + + Replace this with a real scanner. A deny list is easy to get around and is only used here to + keep the sample small and runnable without any model or service. + """ + + def __init__(self, blocked_fragments: Iterable[str]) -> None: + """Initialize the scanner with the fragments to block.""" + self.blocked_fragments = tuple(blocked_fragments) + + def scan(self, command: str) -> ScanVerdict: + """Return a blocked verdict if the command contains a blocked fragment.""" + for fragment in self.blocked_fragments: + if fragment in command: + return ScanVerdict(allowed=False, reason=f"command contains '{fragment}'") + return ScanVerdict(allowed=True) + + +class FakeShellPlugin: + """A stand in for a shell tool. It does not run anything, it only echoes the command.""" + + @kernel_function(name="run", description="Run a shell command.") + def run(self, command: str) -> str: + """Pretend to run the command.""" + return f"ran: {command}" + + +async def main() -> None: + kernel = Kernel() + kernel.add_plugin(FakeShellPlugin(), plugin_name="shell") + + scanner = DenyListScanner(blocked_fragments=["rm -rf /", "mkfs"]) + kernel.add_filter(FilterTypes.FUNCTION_INVOCATION, create_command_guard_filter(scanner)) + + for command in ["ls -la", "rm -rf /"]: + result = await kernel.invoke(plugin_name="shell", function_name="run", command=command) + print(f"{command!r} -> {result}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/tests/unit/kernel/test_command_guard_filter_sample.py b/python/tests/unit/kernel/test_command_guard_filter_sample.py new file mode 100644 index 000000000000..84e0a627674b --- /dev/null +++ b/python/tests/unit/kernel/test_command_guard_filter_sample.py @@ -0,0 +1,92 @@ +# Copyright (c) Microsoft. All rights reserved. + +import importlib.util +from pathlib import Path + +import pytest + +from semantic_kernel import Kernel +from semantic_kernel.filters import FilterTypes +from semantic_kernel.functions import kernel_function + +SAMPLE_PATH = ( + Path(__file__).parents[3] / "samples" / "concepts" / "filtering" / "command_guard_filter.py" +) # python/samples/concepts/filtering/command_guard_filter.py + + +@pytest.fixture(scope="module") +def sample(): + spec = importlib.util.spec_from_file_location("command_guard_filter_sample", SAMPLE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class RecordingShellPlugin: + def __init__(self) -> None: + self.commands: list[str] = [] + + @kernel_function(name="run") + def run(self, command: str) -> str: + self.commands.append(command) + return f"ran: {command}" + + @kernel_function(name="other") + def other(self, text: str) -> str: + return text + + +def _kernel_with_guard(sample, scanner, **filter_kwargs) -> tuple[Kernel, RecordingShellPlugin]: + kernel = Kernel() + plugin = RecordingShellPlugin() + kernel.add_plugin(plugin, plugin_name="shell") + kernel.add_filter(FilterTypes.FUNCTION_INVOCATION, sample.create_command_guard_filter(scanner, **filter_kwargs)) + return kernel, plugin + + +async def test_allowed_command_runs(sample): + kernel, plugin = _kernel_with_guard(sample, sample.DenyListScanner(["rm -rf /"])) + + result = await kernel.invoke(plugin_name="shell", function_name="run", command="ls -la") + + assert str(result) == "ran: ls -la" + assert plugin.commands == ["ls -la"] + + +async def test_blocked_command_does_not_run(sample): + kernel, plugin = _kernel_with_guard(sample, sample.DenyListScanner(["rm -rf /"])) + + result = await kernel.invoke(plugin_name="shell", function_name="run", command="rm -rf /") + + assert "Blocked by command guard" in str(result) + assert "rm -rf /" in str(result) # the reason from the scanner is passed on + assert plugin.commands == [] + + +async def test_async_scanner_is_supported(sample): + class AsyncScanner: + async def scan(self, command: str): + return sample.ScanVerdict(allowed=False, reason="async says no") + + kernel, plugin = _kernel_with_guard(sample, AsyncScanner()) + + result = await kernel.invoke(plugin_name="shell", function_name="run", command="anything") + + assert "async says no" in str(result) + assert plugin.commands == [] + + +async def test_function_without_the_command_argument_is_not_checked(sample): + kernel, _ = _kernel_with_guard(sample, sample.DenyListScanner(["hello"])) + + result = await kernel.invoke(plugin_name="shell", function_name="other", text="hello") + + assert str(result) == "hello" + + +async def test_custom_argument_names(sample): + kernel, _ = _kernel_with_guard(sample, sample.DenyListScanner(["hello"]), argument_names=("text",)) + + result = await kernel.invoke(plugin_name="shell", function_name="other", text="hello") + + assert "Blocked by command guard" in str(result)