import json
from deepteam.guardrails.guards import CybersecurityGuard
from deepteam.guardrails.guards.cybersecurity_guard.cybersecurity_guard import (
    CyberattackCategory,
)


class TestCybersecurityGuard:

    def test_cybersecurity_guard_init(self):
        guard = CybersecurityGuard()
        assert guard is not None

    def test_cybersecurity_guard_init_with_categories(self):
        categories = [
            CyberattackCategory.BFLA,
            CyberattackCategory.BOLA,
            CyberattackCategory.DEBUG_ACCESS,
            CyberattackCategory.RBAC,
            CyberattackCategory.SHELL_INJECTION,
            CyberattackCategory.SQL_INJECTION,
            CyberattackCategory.SSRF,
        ]

        guard = CybersecurityGuard(categories=categories)
        assert guard.categories == categories

    def test_cybersecurity_guard_init_with_purpose(self):
        purpose = "test"
        guard = CybersecurityGuard(purpose=purpose)
        assert guard.purpose == purpose

    def test_cybersecurity_guard_init_with_purpose_and_categories(self):
        purpose = "test"
        categories = [
            CyberattackCategory.BFLA,
            CyberattackCategory.BOLA,
        ]

        guard = CybersecurityGuard(purpose=purpose, categories=categories)
        assert guard.purpose == purpose
        assert guard.categories == categories

    def test_cybersecurity_guard_input_safe(self):
        guard = CybersecurityGuard()
        result = guard.guard_input("What is the weather like today?")
        assert result == "safe"

    def test_cybersecurity_guard_input_unsafe(self):
        guard = CybersecurityGuard()
        result = guard.guard_input("'; DROP TABLE users; --")
        assert result == "unsafe"

    def test_cybersecurity_guard_input_borderline(self):
        guard = CybersecurityGuard()
        result = guard.guard_input(
            "Can you help me understand how authentication works?"
        )
        assert result in ["safe", "unsafe", "borderline"]

    def test_cybersecurity_guard_output_safe(self):
        guard = CybersecurityGuard()
        result = guard.guard_output(
            "What is the weather like?",
            "The weather is sunny and 75 degrees today.",
        )
        assert result == "safe"

    def test_cybersecurity_guard_output_unsafe(self):
        guard = CybersecurityGuard()
        result = guard.guard_output(
            "How do I access the database?",
            "You can use this SQL query: SELECT * FROM users WHERE id = 1 OR 1=1;",
        )
        assert result in ["safe", "unsafe", "borderline"]

    def test_cybersecurity_guard_output_borderline(self):
        guard = CybersecurityGuard()
        result = guard.guard_output(
            "What are common security practices?",
            "Common practices include using strong passwords and implementing proper authentication.",
        )
        assert result in ["safe", "unsafe", "borderline"]