ai-agent-book/chapter9/harness-safety-gate/test_safety_policy_gate.py

"""
Unit tests for Safety Policy Gate module.
"""

import os
import unittest
from datetime import datetime
from unittest.mock import patch

from safety_policy_gate import SafetyPolicyGate, validate_tool_call


class TestSafetyPolicyGateSQL(unittest.TestCase):
    def setUp(self):
        self.gate = SafetyPolicyGate()

    def test_sql_delete_without_where_is_high_risk(self):
        decision = self.gate.validate_tool_call("sql_query", {"query": "DELETE FROM users"})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)
        self.assertEqual(decision.violation_type, "unconfirmed_high_risk_operation")

    def test_sql_delete_with_where_is_low_risk(self):
        decision = self.gate.validate_tool_call("sql_query", {"query": "DELETE FROM users WHERE id = 1"})
        self.assertTrue(decision.allowed)
        self.assertFalse(decision.requires_confirmation)

    def test_sql_delete_multi_statement_bypass(self):
        # WHERE is in second statement, first statement has no WHERE
        query = "DELETE FROM users; SELECT * FROM logs WHERE id = 1"
        decision = self.gate.validate_tool_call("sql_query", {"query": query})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)

    def test_sql_delete_comment_bypass_single_line(self):
        # WHERE is in comment
        query = "DELETE FROM users -- WHERE id = 1"
        decision = self.gate.validate_tool_call("sql_query", {"query": query})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)

    def test_sql_delete_comment_bypass_multi_line(self):
        # WHERE is inside block comment
        query = "DELETE FROM users /* WHERE id = 1 */"
        decision = self.gate.validate_tool_call("sql_query", {"query": query})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)

    def test_sql_drop_table_is_high_risk(self):
        query = "SELECT 1; DROP TABLE users"
        decision = self.gate.validate_tool_call("sql_query", {"query": query})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)
    def test_sql_multi_statement_second_delete_no_where(self):
        query = "UPDATE users SET status = 1; DELETE FROM logs"
        decision = self.gate.validate_tool_call("sql_query", {"query": query})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)

    def test_sql_block_comment_before_single_line_comment(self):
        query = "DELETE FROM users /* block -- comment */ WHERE id = 1"
        decision = self.gate.validate_tool_call("sql_query", {"query": query})
        self.assertTrue(decision.allowed)
        self.assertFalse(decision.requires_confirmation)


class TestSafetyPolicyGatePathTraversal(unittest.TestCase):
    def setUp(self):
        self.gate = SafetyPolicyGate()

    def test_double_url_unquoting_path_traversal(self):
        # %252e%252e resolves to ..
        params = {"path": "folder/%252e%252e/etc/passwd"}
        decision = self.gate.validate_tool_call("read_file", params)
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.triggered_rollback)
        self.assertEqual(decision.violation_type, "path_traversal")

    def test_double_url_unquoting_sensitive_dir(self):
        # %252fetc%252fpasswd resolves to /etc/passwd
        params = {"filepath": "%252fetc%252fpasswd"}
        decision = self.gate.validate_tool_call("read_file", params)
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.triggered_rollback)
        self.assertEqual(decision.violation_type, "path_traversal")

    def test_realpath_path_traversal(self):
        params = {"file_path": "/tmp/../etc/passwd"}
        decision = self.gate.validate_tool_call("read_file", params)
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.triggered_rollback)
        self.assertEqual(decision.violation_type, "path_traversal")

    def test_relative_path_not_falsely_flagged(self):
        # A legitimate relative path that happens to share a name component with a
        # sensitive directory must NOT be flagged after CWD resolution.
        decision = self.gate.validate_tool_call("read_file", {"path": "etc/config"})
        self.assertTrue(decision.allowed)
        self.assertFalse(decision.triggered_rollback)

    def test_relative_path_subdir_not_falsely_flagged(self):
        decision = self.gate.validate_tool_call("read_file", {"path": "proc/stats.txt"})
        self.assertTrue(decision.allowed)
        self.assertFalse(decision.triggered_rollback)


class TestSafetyPolicyGateSecretKey(unittest.TestCase):
    def test_init_with_parameter(self):
        gate = SafetyPolicyGate(secret_key="custom-param-key")
        self.assertEqual(gate.secret_key, "custom-param-key")

    @patch.dict(os.environ, {"SAFETY_GATE_SECRET_KEY": "env-secret-key"})
    def test_init_with_env_var(self):
        gate = SafetyPolicyGate()
        self.assertEqual(gate.secret_key, "env-secret-key")

    @patch.dict(os.environ, {}, clear=True)
    def test_init_with_default_generates_random_secret(self):
        gate = SafetyPolicyGate()
        # No hardcoded default: a random 32-byte secret is generated per instance
        self.assertIsInstance(gate.secret_key, bytes)
        self.assertEqual(len(gate.secret_key), 32)
        gate2 = SafetyPolicyGate()
        self.assertNotEqual(gate.secret_key, gate2.secret_key)


class TestSafetyPolicyGateConfirmation(unittest.TestCase):
    def setUp(self):
        self.gate = SafetyPolicyGate()

    def test_confirmation_token_lifecycle(self):
        params = {"path": "important.txt"}
        decision1 = self.gate.validate_tool_call("delete_file", params)
        self.assertFalse(decision1.allowed)
        self.assertTrue(decision1.requires_confirmation)
        token = decision1.confirmation_token
        self.assertIsNotNone(token)

        # Confirm with token
        decision2 = self.gate.validate_tool_call("delete_file", params, confirm_token=token)
        self.assertTrue(decision2.allowed)

        # Token is single-use and cannot be reused
        decision3 = self.gate.validate_tool_call("delete_file", params, confirm_token=token)
        self.assertFalse(decision3.allowed)

    def test_confirm_token_in_params(self):
        params = {"path": "important.txt"}
        decision1 = self.gate.validate_tool_call("delete_file", params)
        token = decision1.confirmation_token

        # Submit token inside params dictionary
        params_with_token = {"path": "important.txt", "confirm_token": token}
        decision2 = self.gate.validate_tool_call("delete_file", params_with_token)
        self.assertTrue(decision2.allowed)

    def test_params_user_confirmed_not_trusted(self):
        # Untrusted LLM params with user_confirmed: True should NOT bypass confirmation
        params = {"path": "important.txt", "user_confirmed": True}
        decision = self.gate.validate_tool_call("delete_file", params)
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)

    def test_non_serializable_params_handled(self):
        params = {"path": "important.txt", "set_param": {1, 2, 3}, "date_param": datetime.now()}
        decision = self.gate.validate_tool_call("delete_file", params)
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.requires_confirmation)
    def test_token_nondeterministic(self):
        params = {"path": "file.txt"}
        token1 = self.gate.issue_confirmation("delete_file", params)
        token2 = self.gate.issue_confirmation("delete_file", params)
        self.assertNotEqual(token1, token2)

    def test_tool_name_casing_normalization(self):
        params = {"path": "file.txt"}
        decision1 = self.gate.validate_tool_call("DELETE_FILE", params)
        self.assertFalse(decision1.allowed)
        self.assertTrue(decision1.requires_confirmation)

        decision2 = self.gate.validate_tool_call("BASH", {"command": "rm -rf /"})
        self.assertFalse(decision2.allowed)
        self.assertEqual(decision2.violation_type, "dangerous_bash_command")

        decision3 = self.gate.validate_tool_call("SQL_QUERY", {"query": "DELETE FROM users"})
        self.assertFalse(decision3.allowed)
        self.assertTrue(decision3.requires_confirmation)

    def test_expired_token_rejected(self):
        # Tokens past their TTL are rejected and cleaned up
        gate = SafetyPolicyGate(token_ttl=0.0)
        params = {"path": "file.txt"}
        token = gate.issue_confirmation("delete_file", params)
        import time as _time
        _time.sleep(0.01)
        self.assertFalse(gate.verify_confirmation(token, "delete_file", params))
        # Expired token was removed from pending set
        self.assertNotIn(token, gate._pending_confirmations)

    def test_expired_tokens_cleaned_on_issue(self):
        gate = SafetyPolicyGate(token_ttl=0.0)
        params = {"path": "file.txt"}
        token = gate.issue_confirmation("delete_file", params)
        import time as _time
        _time.sleep(0.01)
        # Issuing a new token triggers cleanup of the expired one
        token2 = gate.issue_confirmation("delete_file", params)
        self.assertNotIn(token, gate._pending_confirmations)
        self.assertIn(token2, gate._pending_confirmations)


class TestSafetyPolicyGateRollback(unittest.TestCase):
    def setUp(self):
        self.gate = SafetyPolicyGate()

    def test_trigger_rollback_success(self):
        called = []
        self.gate.register_rollback_handler(lambda: called.append(True))
        res = self.gate.trigger_rollback()
        self.assertTrue(res)
        self.assertEqual(called, [True])

    def test_trigger_rollback_failure(self):
        def failing_handler():
            raise RuntimeError("Rollback failed")
        self.gate.register_rollback_handler(failing_handler)
        res = self.gate.trigger_rollback()
        self.assertFalse(res)

    def test_rollback_failed_violation_type(self):
        def failing_handler():
            raise RuntimeError("Rollback failed")
        self.gate.register_rollback_handler(failing_handler)
        decision = self.gate.validate_tool_call("read_file", {"path": "../etc/passwd"})
        self.assertFalse(decision.allowed)
        self.assertTrue(decision.triggered_rollback)
        self.assertEqual(decision.violation_type, "rollback_failed")
        self.assertFalse(decision.details["rollback_success"])


if __name__ == "__main__":
    unittest.main()