tests/guardrails_genie/guardrails/test_secrets_detection.py

import hashlib
import re

import pytest
from hypothesis import strategies as st, given, settings

from guardrails_genie.guardrails import SecretsDetectionGuardrail
from guardrails_genie.guardrails.secrets_detection import (
    SecretsDetectionSimpleResponse,
    SecretsDetectionResponse,
    REDACTION,
    redact_value,
)


@pytest.fixture
def mock_secrets_guard(monkeypatch):
    def _mock_guard(*args, **kwargs):
        prompt = kwargs.get("prompt")
        return_detected_types = kwargs.get("return_detected_secrets")

        if "safe text" in prompt:
            if return_detected_types:
                return SecretsDetectionResponse(
                    contains_secrets=False,
                    explanation="No secrets detected in the text.",
                    detected_secrets={},
                    redacted_text=prompt,
                    risk_score=0.0,
                )
            else:
                return SecretsDetectionSimpleResponse(
                    contains_secrets=False,
                    explanation="No secrets detected in the text.",
                    redacted_text=prompt,
                    risk_score=0.0,
                )
        else:
            if return_detected_types:
                return SecretsDetectionResponse(
                    contains_secrets=True,
                    explanation="The output contains secrets.",
                    detected_secrets={"secrets": ["API_KEY"]},
                    redacted_text="My secret key is [REDACTED:]************[:REDACTED]",
                    risk_score=1.0,
                )
            else:
                return SecretsDetectionSimpleResponse(
                    contains_secrets=True,
                    explanation="The output contains secrets.",
                    redacted_text="My secret key is [REDACTED:]************[:REDACTED]",
                    risk_score=1.0,
                )

    monkeypatch.setattr(
        "guardrails_genie.guardrails.secrets_detection.SecretsDetectionGuardrail.guard",
        _mock_guard,
    )


def test_redact_partial():
    text = "ABCDEFGHIJKL"
    redacted_text = redact_value(text, REDACTION.REDACT_PARTIAL)
    assert redacted_text == "[REDACTED:]AB..KL[:REDACTED]"


def test_redact_all():
    text = "ABCDEFGHIJKL"
    redacted_text = redact_value(text, REDACTION.REDACT_ALL)
    assert redacted_text == "[REDACTED:]************[:REDACTED]"


def test_redact_hash():
    text = "ABCDEFGHIJKL"
    hashed_value = hashlib.md5(text.encode()).hexdigest()
    redacted_text = redact_value(text, REDACTION.REDACT_HASH)
    assert redacted_text == f"[REDACTED:]{hashed_value}[:REDACTED]"


def test_secrets_detection_guardrail_detect_types(mock_secrets_guard):
    from guardrails_genie.guardrails.secrets_detection import (
        SecretsDetectionGuardrail,
        REDACTION,
    )

    guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
    prompt = "My secret key is ABCDEFGHIJKL"

    result = guardrail.guard(prompt=prompt, return_detected_secrets=True)

    assert result.contains_secrets is True
    assert result.explanation == "The output contains secrets."
    assert result.detected_secrets == {"secrets": ["API_KEY"]}
    assert result.redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"


def test_secrets_detection_guardrail_simple_response(mock_secrets_guard):
    from guardrails_genie.guardrails.secrets_detection import (
        SecretsDetectionGuardrail,
        REDACTION,
    )

    guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
    prompt = "My secret key is ABCDEFGHIJKL"

    result = guardrail.guard(prompt=prompt, return_detected_secrets=False)

    assert result.contains_secrets is True
    assert result.explanation == "The output contains secrets."
    assert result.redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"


def test_secrets_detection_guardrail_no_secrets(mock_secrets_guard):
    from guardrails_genie.guardrails.secrets_detection import (
        SecretsDetectionGuardrail,
        REDACTION,
    )

    guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
    prompt = "This is a safe text with no secrets."

    result = guardrail.guard(prompt=prompt, return_detected_secrets=True)

    assert result.contains_secrets is False
    assert result.explanation == "No secrets detected in the text."
    assert result.detected_secrets == {}
    assert result.redacted_text == prompt


def pattern_strategy(pattern):
    return st.from_regex(re.compile(pattern), fullmatch=True)


@settings(deadline=1000)
@given(pattern_strategy(r"AKIA[0-9A-Z]{16}"))
def test_specific_pattern_guardrail(text):
    guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
    result = guardrail.guard(prompt=text, return_detected_secrets=True)

    assert result.contains_secrets is True
    assert "AWS Access Key" in result.detected_secrets