610 lines
21 KiB
Python
610 lines
21 KiB
Python
# -*- coding: utf-8 -*-
|
|
# pylint: disable=redefined-outer-name,protected-access
|
|
"""Tests for pineagents.security.tool_guard.engine."""
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from pineagents.security.tool_guard.engine import (
|
|
ToolGuardEngine,
|
|
_guard_enabled,
|
|
get_guard_engine,
|
|
)
|
|
from pineagents.security.tool_guard.models import (
|
|
GuardFinding,
|
|
GuardSeverity,
|
|
GuardThreatCategory,
|
|
ToolGuardResult,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_guardian(
|
|
name: str = "test_guardian",
|
|
always_run: bool = False,
|
|
findings: list[GuardFinding] | None = None,
|
|
side_effect=None,
|
|
) -> MagicMock:
|
|
"""Build a lightweight MagicMock that quacks like a BaseToolGuardian."""
|
|
g = MagicMock()
|
|
g.name = name
|
|
g.always_run = always_run
|
|
if side_effect is not None:
|
|
g.guard.side_effect = side_effect
|
|
else:
|
|
g.guard.return_value = findings or []
|
|
return g
|
|
|
|
|
|
def _make_finding(
|
|
severity: GuardSeverity = GuardSeverity.HIGH,
|
|
rule_id: str = "R001",
|
|
) -> GuardFinding:
|
|
return GuardFinding(
|
|
id="f-1",
|
|
rule_id=rule_id,
|
|
category=GuardThreatCategory.PATH_TRAVERSAL,
|
|
severity=severity,
|
|
title="Test finding",
|
|
description="desc",
|
|
tool_name="some_tool",
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_singleton():
|
|
"""Reset the module-level singleton before and after each test."""
|
|
import pineagents.security.tool_guard.engine as eng_mod
|
|
|
|
eng_mod._engine_instance = None
|
|
yield
|
|
eng_mod._engine_instance = None
|
|
|
|
|
|
@pytest.fixture()
|
|
def engine_with_defaults():
|
|
"""Return an engine with mocked defaults (no filesystem/config access)."""
|
|
g1 = _make_guardian("alpha", always_run=False)
|
|
g2 = _make_guardian("beta", always_run=True)
|
|
with patch.object(
|
|
ToolGuardEngine,
|
|
"_default_guardians",
|
|
return_value=[g1, g2],
|
|
), patch.object(ToolGuardEngine, "_reload_tool_sets"):
|
|
eng = ToolGuardEngine(enabled=True)
|
|
eng._guardians = [g1, g2]
|
|
eng._guarded_tools = None
|
|
eng._denied_tools = set()
|
|
eng._auto_denied_rules = set()
|
|
return eng
|
|
|
|
|
|
# ===================================================================
|
|
# TestGuardEnabled
|
|
# ===================================================================
|
|
|
|
|
|
class TestGuardEnabled:
|
|
"""Tests for the _guard_enabled() free function."""
|
|
|
|
@patch("pineagents.security.tool_guard.engine.EnvVarLoader")
|
|
def test_enabled_when_env_var_is_true_string(self, mock_loader):
|
|
mock_loader.get_str.return_value = "true"
|
|
assert _guard_enabled() is True
|
|
|
|
@patch("pineagents.security.tool_guard.engine.EnvVarLoader")
|
|
def test_enabled_when_env_var_is_false_string(self, mock_loader):
|
|
mock_loader.get_str.return_value = "false"
|
|
assert _guard_enabled() is False
|
|
|
|
@patch("pineagents.security.tool_guard.engine.EnvVarLoader")
|
|
def test_enabled_defaults_to_true_when_no_env_no_config(self, mock_loader):
|
|
mock_loader.get_str.return_value = ""
|
|
with patch(
|
|
"pineagents.config.load_config",
|
|
side_effect=Exception("no config"),
|
|
):
|
|
assert _guard_enabled() is True
|
|
|
|
@patch("pineagents.security.tool_guard.engine.EnvVarLoader")
|
|
def test_enabled_reads_from_config_when_no_env(self, mock_loader):
|
|
mock_loader.get_str.return_value = ""
|
|
mock_cfg = MagicMock()
|
|
mock_cfg.security.tool_guard.enabled = False
|
|
with patch(
|
|
"pineagents.config.load_config",
|
|
return_value=mock_cfg,
|
|
):
|
|
assert _guard_enabled() is False
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineInit
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineInit:
|
|
@patch.object(ToolGuardEngine, "_reload_tool_sets")
|
|
@patch.object(ToolGuardEngine, "_default_guardians")
|
|
def test_init_with_no_guardians_creates_defaults(
|
|
self,
|
|
mock_defaults,
|
|
mock_reload, # pylint: disable=unused-argument
|
|
):
|
|
default_g = _make_guardian("default_g")
|
|
mock_defaults.return_value = [default_g]
|
|
# Patch the guardian imports inside _default_guardians to be safe
|
|
eng = ToolGuardEngine(enabled=True)
|
|
mock_defaults.assert_called_once()
|
|
assert eng._guardians == [default_g]
|
|
|
|
@patch.object(ToolGuardEngine, "_reload_tool_sets")
|
|
def test_init_with_custom_guardians(
|
|
self,
|
|
mock_reload, # pylint: disable=unused-argument
|
|
):
|
|
g1 = _make_guardian("custom1")
|
|
g2 = _make_guardian("custom2")
|
|
eng = ToolGuardEngine(guardians=[g1, g2], enabled=True)
|
|
assert eng._guardians == [g1, g2]
|
|
|
|
@patch.object(ToolGuardEngine, "_reload_tool_sets")
|
|
def test_init_with_enabled_false(
|
|
self,
|
|
mock_reload, # pylint: disable=unused-argument
|
|
):
|
|
eng = ToolGuardEngine(enabled=False)
|
|
assert eng.enabled is False
|
|
|
|
@patch.object(ToolGuardEngine, "_reload_tool_sets")
|
|
@patch.object(ToolGuardEngine, "_default_guardians")
|
|
def test_init_calls_reload_tool_sets(self, mock_defaults, mock_reload):
|
|
mock_defaults.return_value = []
|
|
ToolGuardEngine(enabled=True)
|
|
mock_reload.assert_called_once()
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineProperties
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineProperties:
|
|
def test_guardian_names_returns_list_of_names(self, engine_with_defaults):
|
|
assert engine_with_defaults.guardian_names == ["alpha", "beta"]
|
|
|
|
def test_enabled_getter_setter(self, engine_with_defaults):
|
|
assert engine_with_defaults.enabled is True
|
|
engine_with_defaults.enabled = False
|
|
assert engine_with_defaults.enabled is False
|
|
engine_with_defaults.enabled = True
|
|
assert engine_with_defaults.enabled is True
|
|
|
|
def test_guarded_tools_returns_set_or_none(self, engine_with_defaults):
|
|
# Default: None (guard all)
|
|
assert engine_with_defaults.guarded_tools is None
|
|
|
|
engine_with_defaults._guarded_tools = {"read_file", "write_file"}
|
|
assert engine_with_defaults.guarded_tools == {
|
|
"read_file",
|
|
"write_file",
|
|
}
|
|
|
|
def test_denied_tools_returns_set(self, engine_with_defaults):
|
|
assert engine_with_defaults.denied_tools == set()
|
|
engine_with_defaults._denied_tools = {"dangerous_tool"}
|
|
assert engine_with_defaults.denied_tools == {"dangerous_tool"}
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineRegister
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineRegister:
|
|
def test_register_guardian_adds_guardian(self, engine_with_defaults):
|
|
g = _make_guardian("gamma")
|
|
engine_with_defaults.register_guardian(g)
|
|
assert "gamma" in engine_with_defaults.guardian_names
|
|
assert g in engine_with_defaults._guardians
|
|
|
|
def test_unregister_guardian_removes_by_name(self, engine_with_defaults):
|
|
assert "alpha" in engine_with_defaults.guardian_names
|
|
result = engine_with_defaults.unregister_guardian("alpha")
|
|
assert result is True
|
|
assert "alpha" not in engine_with_defaults.guardian_names
|
|
|
|
def test_unregister_guardian_returns_false_when_not_found(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
result = engine_with_defaults.unregister_guardian("nonexistent")
|
|
assert result is False
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineIsDenied
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineIsDenied:
|
|
def test_is_denied_true_when_in_denied_set(self, engine_with_defaults):
|
|
engine_with_defaults._denied_tools = {"rm_rf", "drop_db"}
|
|
assert engine_with_defaults.is_denied("rm_rf") is True
|
|
|
|
def test_is_denied_false_when_not_in_denied_set(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
engine_with_defaults._denied_tools = {"rm_rf"}
|
|
assert engine_with_defaults.is_denied("read_file") is False
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineShouldAutoDenyResult
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineShouldAutoDenyResult:
|
|
"""Tests for ToolGuardEngine.should_auto_deny_result.
|
|
|
|
The method returns True only when at least one finding's rule_id is
|
|
listed in the engine's auto-deny rule set; otherwise False. None /
|
|
empty results / empty rule set all short-circuit to False.
|
|
"""
|
|
|
|
def _make_result(self, findings):
|
|
return ToolGuardResult(
|
|
tool_name="execute_shell_command",
|
|
params={"command": "x"},
|
|
findings=findings,
|
|
guard_duration_seconds=0.0,
|
|
)
|
|
|
|
def test_returns_false_when_result_is_none(self, engine_with_defaults):
|
|
engine_with_defaults._auto_denied_rules = {"R1"}
|
|
assert engine_with_defaults.should_auto_deny_result(None) is False
|
|
|
|
def test_returns_false_when_findings_empty(self, engine_with_defaults):
|
|
engine_with_defaults._auto_denied_rules = {"R1"}
|
|
result = self._make_result([])
|
|
assert engine_with_defaults.should_auto_deny_result(result) is False
|
|
|
|
def test_returns_false_when_auto_deny_rules_empty(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
"""No auto-deny rules → never auto-deny regardless of findings."""
|
|
engine_with_defaults._auto_denied_rules = set()
|
|
result = self._make_result([_make_finding(rule_id="R1")])
|
|
assert engine_with_defaults.should_auto_deny_result(result) is False
|
|
|
|
def test_returns_true_when_finding_matches_rule(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
engine_with_defaults._auto_denied_rules = {"R1"}
|
|
result = self._make_result([_make_finding(rule_id="R1")])
|
|
assert engine_with_defaults.should_auto_deny_result(result) is True
|
|
|
|
def test_returns_false_when_findings_dont_match(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
engine_with_defaults._auto_denied_rules = {"DANGEROUS_RULE"}
|
|
result = self._make_result(
|
|
[
|
|
_make_finding(rule_id="OTHER_1"),
|
|
_make_finding(rule_id="OTHER_2"),
|
|
],
|
|
)
|
|
assert engine_with_defaults.should_auto_deny_result(result) is False
|
|
|
|
def test_returns_true_when_any_finding_matches(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
"""Among multiple findings, a single matching rule_id triggers."""
|
|
engine_with_defaults._auto_denied_rules = {"DANGEROUS"}
|
|
result = self._make_result(
|
|
[
|
|
_make_finding(rule_id="HARMLESS"),
|
|
_make_finding(rule_id="DANGEROUS"),
|
|
_make_finding(rule_id="OTHER"),
|
|
],
|
|
)
|
|
assert engine_with_defaults.should_auto_deny_result(result) is True
|
|
|
|
def test_auto_denied_rules_property_reflects_state(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
"""Public ``auto_denied_rules`` property exposes the configured set."""
|
|
engine_with_defaults._auto_denied_rules = {"R1", "R2"}
|
|
assert engine_with_defaults.auto_denied_rules == {"R1", "R2"}
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineIsGuarded
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineIsGuarded:
|
|
def test_is_guarded_true_when_in_guarded_set(self, engine_with_defaults):
|
|
engine_with_defaults._guarded_tools = {"read_file", "write_file"}
|
|
assert engine_with_defaults.is_guarded("read_file") is True
|
|
|
|
def test_is_guarded_true_when_guarded_tools_is_none(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
engine_with_defaults._guarded_tools = None
|
|
# None means guard all tools
|
|
assert engine_with_defaults.is_guarded("any_tool") is True
|
|
|
|
def test_is_guarded_false_when_not_in_set(self, engine_with_defaults):
|
|
engine_with_defaults._guarded_tools = {"read_file"}
|
|
assert engine_with_defaults.is_guarded("execute_shell") is False
|
|
|
|
|
|
# ===================================================================
|
|
# TestToolGuardEngineGuard
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineGuard:
|
|
def test_guard_returns_none_when_disabled(self, engine_with_defaults):
|
|
engine_with_defaults.enabled = False
|
|
result = engine_with_defaults.guard("some_tool", {})
|
|
assert result is None
|
|
|
|
def test_guard_runs_all_guardians_when_tool_is_guarded(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
finding = _make_finding()
|
|
engine_with_defaults._guardians[0].guard.return_value = [finding]
|
|
engine_with_defaults._guardians[1].guard.return_value = []
|
|
|
|
result = engine_with_defaults.guard(
|
|
"some_tool",
|
|
{"path": "/etc/passwd"},
|
|
)
|
|
|
|
assert isinstance(result, ToolGuardResult)
|
|
assert "alpha" in result.guardians_used
|
|
assert "beta" in result.guardians_used
|
|
assert finding in result.findings
|
|
engine_with_defaults._guardians[0].guard.assert_called_once_with(
|
|
"some_tool",
|
|
{"path": "/etc/passwd"},
|
|
)
|
|
engine_with_defaults._guardians[1].guard.assert_called_once_with(
|
|
"some_tool",
|
|
{"path": "/etc/passwd"},
|
|
)
|
|
|
|
def test_guard_skips_non_always_run_guardians_when_only_always_run(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
# alpha has always_run=False, beta has always_run=True
|
|
engine_with_defaults.guard(
|
|
"some_tool",
|
|
{},
|
|
only_always_run=True,
|
|
)
|
|
|
|
# Only beta should be called
|
|
engine_with_defaults._guardians[0].guard.assert_not_called()
|
|
engine_with_defaults._guardians[1].guard.assert_called_once()
|
|
|
|
def test_guard_aggregates_findings_from_all_guardians(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
f1 = _make_finding(severity=GuardSeverity.HIGH, rule_id="R001")
|
|
f2 = _make_finding(severity=GuardSeverity.MEDIUM, rule_id="R002")
|
|
engine_with_defaults._guardians[0].guard.return_value = [f1]
|
|
engine_with_defaults._guardians[1].guard.return_value = [f2]
|
|
|
|
result = engine_with_defaults.guard("tool", {})
|
|
|
|
assert result.findings == [f1, f2]
|
|
assert result.findings_count == 2
|
|
|
|
def test_guard_returns_result_with_guardians_used(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
engine_with_defaults._guardians[0].guard.return_value = []
|
|
engine_with_defaults._guardians[1].guard.return_value = []
|
|
|
|
result = engine_with_defaults.guard("tool", {})
|
|
|
|
assert result.guardians_used == ["alpha", "beta"]
|
|
assert not result.guardians_failed
|
|
assert result.tool_name == "tool"
|
|
assert not result.params
|
|
|
|
def test_guard_handles_guardian_exception_gracefully(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
engine_with_defaults._guardians[0].guard.side_effect = RuntimeError(
|
|
"boom",
|
|
)
|
|
engine_with_defaults._guardians[1].guard.return_value = []
|
|
|
|
result = engine_with_defaults.guard("tool", {})
|
|
|
|
# Failing guardian should appear in
|
|
# guardians_failed, not guardians_used
|
|
assert "alpha" not in result.guardians_used
|
|
assert len(result.guardians_failed) == 1
|
|
assert result.guardians_failed[0]["name"] == "alpha"
|
|
assert "boom" in result.guardians_failed[0]["error"]
|
|
# Non-failing guardian still runs fine
|
|
assert "beta" in result.guardians_used
|
|
|
|
def test_guard_sets_duration(self, engine_with_defaults):
|
|
engine_with_defaults._guardians[0].guard.return_value = []
|
|
engine_with_defaults._guardians[1].guard.return_value = []
|
|
|
|
result = engine_with_defaults.guard("tool", {})
|
|
|
|
assert result.guard_duration_seconds >= 0.0
|
|
|
|
|
|
# ===================================================================
|
|
# TestGetGuardEngine
|
|
# ===================================================================
|
|
|
|
|
|
class TestGetGuardEngine:
|
|
@patch(
|
|
"pineagents.security.tool_guard.engine.ToolGuardEngine",
|
|
autospec=True,
|
|
)
|
|
def test_returns_singleton_instance(self, MockEngine):
|
|
MockEngine.return_value = MagicMock(spec=ToolGuardEngine)
|
|
result = get_guard_engine()
|
|
MockEngine.assert_called_once()
|
|
assert result is MockEngine.return_value
|
|
|
|
@patch(
|
|
"pineagents.security.tool_guard.engine.ToolGuardEngine",
|
|
autospec=True,
|
|
)
|
|
def test_returns_same_instance_on_multiple_calls(self, MockEngine):
|
|
instance = MagicMock(spec=ToolGuardEngine)
|
|
MockEngine.return_value = instance
|
|
r1 = get_guard_engine()
|
|
r2 = get_guard_engine()
|
|
# ToolGuardEngine() should only have been called once
|
|
MockEngine.assert_called_once()
|
|
assert r1 is r2
|
|
|
|
|
|
# ===================================================================
|
|
# TestReloadRules
|
|
# ===================================================================
|
|
|
|
|
|
class TestToolGuardEngineReloadRules:
|
|
def test_reload_rules_calls_reload_on_guardians_and_reloads_tool_sets(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
g_with_reload = _make_guardian("reloadable")
|
|
g_with_reload.reload = MagicMock()
|
|
engine_with_defaults._guardians = [g_with_reload]
|
|
|
|
with patch.object(
|
|
engine_with_defaults,
|
|
"_reload_tool_sets",
|
|
) as mock_reload_ts:
|
|
engine_with_defaults.reload_rules()
|
|
|
|
g_with_reload.reload.assert_called_once()
|
|
mock_reload_ts.assert_called_once()
|
|
|
|
def test_reload_rules_skips_guardians_without_reload(
|
|
self,
|
|
engine_with_defaults,
|
|
):
|
|
# Default MagicMock has 'reload' attribute, so delete it explicitly
|
|
g_no_reload = _make_guardian("no_reload")
|
|
del g_no_reload.reload
|
|
engine_with_defaults._guardians = [g_no_reload]
|
|
|
|
with patch.object(
|
|
engine_with_defaults,
|
|
"_reload_tool_sets",
|
|
):
|
|
# Should not raise
|
|
engine_with_defaults.reload_rules()
|
|
|
|
|
|
# ===================================================================
|
|
# Shared safety → auto-deny integration
|
|
# ===================================================================
|
|
|
|
|
|
class TestSharedSafetyAutoDenyIntegration:
|
|
"""finding → should_auto_deny_result for default catastrophic policy."""
|
|
|
|
def _engine_with_shared_safety(self) -> ToolGuardEngine:
|
|
from pineagents.security.tool_guard.guardians.rule_guardian import (
|
|
SharedSafetyToolGuardian,
|
|
)
|
|
|
|
with patch.object(ToolGuardEngine, "_reload_tool_sets"):
|
|
eng = ToolGuardEngine(
|
|
guardians=[SharedSafetyToolGuardian()],
|
|
enabled=True,
|
|
)
|
|
eng._auto_denied_rules = {"SAFETY_CHECKS_DESTRUCTIVE_COMMAND"}
|
|
eng._guarded_tools = set() # shell not in guarded set
|
|
return eng
|
|
|
|
def test_rm_root_auto_denied(self):
|
|
eng = self._engine_with_shared_safety()
|
|
result = eng.guard(
|
|
"execute_shell_command",
|
|
{"command": "rm -rf /"},
|
|
only_always_run=True,
|
|
)
|
|
assert result is not None
|
|
assert any(
|
|
f.rule_id == "SAFETY_CHECKS_DESTRUCTIVE_COMMAND"
|
|
for f in result.findings
|
|
)
|
|
assert eng.should_auto_deny_result(result) is True
|
|
|
|
def test_npm_run_reboot_not_auto_denied(self):
|
|
eng = self._engine_with_shared_safety()
|
|
result = eng.guard(
|
|
"execute_shell_command",
|
|
{"command": "npm run reboot"},
|
|
only_always_run=True,
|
|
)
|
|
assert result is not None
|
|
assert not result.findings
|
|
assert eng.should_auto_deny_result(result) is False
|
|
|
|
def test_reboot_finding_not_auto_denied(self):
|
|
"""System power is CRITICAL for approval, not default hard DENY."""
|
|
eng = self._engine_with_shared_safety()
|
|
result = eng.guard(
|
|
"execute_shell_command",
|
|
{"command": "reboot"},
|
|
only_always_run=True,
|
|
)
|
|
assert result is not None
|
|
assert any(
|
|
f.rule_id == "SAFETY_CHECKS_SYSTEM_POWER" for f in result.findings
|
|
)
|
|
assert eng.should_auto_deny_result(result) is False
|
|
|
|
def test_path_traversal_rm_auto_denied(self):
|
|
eng = self._engine_with_shared_safety()
|
|
result = eng.guard(
|
|
"execute_shell_command",
|
|
{"command": "rm -rf /tmp/../etc"},
|
|
only_always_run=True,
|
|
)
|
|
assert result is not None
|
|
assert eng.should_auto_deny_result(result) is True
|