feat: upgrade and improve secrets detection

guardrails_genie/guardrails/secrets_detection/__init__.py

@@ -1,17 +1,15 @@
 from guardrails_genie.guardrails.secrets_detection.secrets_detection import (
-    DEFAULT_SECRETS_PATTERNS,
     SecretsDetectionGuardrail,
     SecretsDetectionSimpleResponse,
     SecretsDetectionResponse,
     REDACTION,
-    redact,
+    redact_value,
 )
 
 __all__ = [
-    "DEFAULT_SECRETS_PATTERNS",
     "SecretsDetectionGuardrail",
     "SecretsDetectionSimpleResponse",
     "SecretsDetectionResponse",
     "REDACTION",
-    "redact",
+    "redact_value",
 ]

guardrails_genie/guardrails/secrets_detection/secrets_detection.py

@@ -1,41 +1,30 @@
 import hashlib
 import json
+import os
 import pathlib
+import tempfile
 from enum import Enum
-from typing import Union, Optional
+from typing import Optional, Any
 
 import weave
-from pydantic import BaseModel
+from pydantic import BaseModel, PrivateAttr
 
 from guardrails_genie.guardrails.base import Guardrail
-from guardrails_genie.regex_model import RegexModel
 
-
-def load_secrets_patterns() -> dict[str, list[str]]:
-    """
-    Load secret patterns from a JSONL file and return them as a dictionary.
-
-    Returns:
-        dict: A dictionary where keys are pattern names and values are lists of regex patterns.
-    """
-    default_patterns = {}
-    patterns = (
-        pathlib.Path(__file__).parent.absolute() / "secrets_patterns.jsonl"
-    ).read_text()
-
-    for pattern in patterns.splitlines():
-        pattern = json.loads(pattern)
-        default_patterns[pattern["name"]] = [rf"{pat}" for pat in pattern["patterns"]]
-    return default_patterns
-
-
-# Load default secret patterns from the JSONL file
-DEFAULT_SECRETS_PATTERNS = load_secrets_patterns()
+try:
+    from detect_secrets import SecretsCollection
+    from detect_secrets.settings import default_settings
+    import hyperscan
+except ImportError:
+    raise ImportError(
+        "The `detect-secrets` and the `hyperscan` packages are required for using the SecretsGuardrail. "
+        "Please install then by running `pip install detect-secrets hyperscan`."
+    )
 
 
 class REDACTION(str, Enum):
     """
-    Enum for different types of redaction methods.
+    Enum for different types of redaction modes.
     """
 
     REDACT_PARTIAL = "REDACT_PARTIAL"
@@ -44,31 +33,31 @@ class REDACTION(str, Enum):
     REDACT_NONE = "REDACT_NONE"
 
 
-def redact(text: str, matches: list[str], redaction_type: REDACTION) -> str:
+def redact_value(value: str, mode: str) -> str:
     """
-    Redact the given matches in the text based on the redaction type.
+    Redacts the given value based on the specified redaction mode.
 
     Args:
-        text (str): The input text to redact.
-        matches (list[str]): List of strings to be redacted.
-        redaction_type (REDACTION): The type of redaction to apply.
+        value (str): The string value to be redacted.
+        mode (str): The redaction mode to be applied. It can be one of the following:
+            - REDACTION.REDACT_PARTIAL: Partially redacts the value.
+            - REDACTION.REDACT_ALL: Fully redacts the value.
+            - REDACTION.REDACT_HASH: Redacts the value by hashing it.
+            - REDACTION.REDACT_NONE: No redaction is applied.
 
     Returns:
-        str: The redacted text.
-    """
-    for match in matches:
-        if redaction_type == REDACTION.REDACT_PARTIAL:
-            replacement = "[REDACTED:]" + match[:2] + ".." + match[-2:] + "[:REDACTED]"
-        elif redaction_type == REDACTION.REDACT_ALL:
-            replacement = "[REDACTED:]" + ("*" * len(match)) + "[:REDACTED]"
-        elif redaction_type == REDACTION.REDACT_HASH:
-            replacement = (
-                "[REDACTED:]" + hashlib.md5(match.encode()).hexdigest() + "[:REDACTED]"
-            )
-        else:
-            replacement = match
-        text = text.replace(match, replacement)
-    return text
+        str: The redacted value based on the specified mode.
+    """
+    replacement = value
+    if mode == REDACTION.REDACT_PARTIAL:
+        replacement = "[REDACTED:]" + value[:2] + ".." + value[-2:] + "[:REDACTED]"
+    elif mode == REDACTION.REDACT_ALL:
+        replacement = "[REDACTED:]" + ("*" * len(value)) + "[:REDACTED]"
+    elif mode == REDACTION.REDACT_HASH:
+        replacement = (
+            "[REDACTED:]" + hashlib.md5(value.encode()).hexdigest() + "[:REDACTED]"
+        )
+    return replacement
 
 
 class SecretsDetectionSimpleResponse(BaseModel):
@@ -79,11 +68,13 @@ class SecretsDetectionSimpleResponse(BaseModel):
         contains_secrets (bool): Indicates if secrets were detected.
         explanation (str): Explanation of the detection result.
         redacted_text (Optional[str]): The redacted text if secrets were found.
+        risk_score (float): The risk score of the detection result. (0.0, 0.5, 1.0)
     """
 
     contains_secrets: bool
     explanation: str
     redacted_text: Optional[str] = None
+    risk_score: float = 0.0
 
     @property
     def safe(self) -> bool:
@@ -104,54 +95,329 @@ class SecretsDetectionResponse(SecretsDetectionSimpleResponse):
         detected_secrets (dict[str, list[str]]): Dictionary of detected secrets.
     """
 
-    detected_secrets: dict[str, list[str]]
+    detected_secrets: dict[str, Any] | None = None
+
+
+class SecretsInfo(BaseModel):
+    """
+    Model representing information about a detected secret.
+
+    Attributes:
+        secret (str): The detected secret value.
+        line_number (int): The line number where the secret was found.
+    """
+
+    secret: str
+    line_number: int
+
+
+class ScanResult(BaseModel):
+    """
+    Model representing the result of a secrets scan.
+
+    Attributes:
+        detected_secrets (dict[str, Any] | None): Dictionary of detected secrets, or None if no secrets were found.
+        modified_prompt (str): The modified prompt with secrets redacted.
+        has_secret (bool): Indicates if any secrets were detected.
+        risk_score (float): The risk score of the detection result.
+    """
+
+    detected_secrets: dict[str, Any] | None = None
+    modified_prompt: str
+    has_secret: bool
+    risk_score: float
+
+
+class DetectSecretsModel(weave.Model):
+    """
+    Model for detecting secrets using the detect-secrets library.
+    """
+
+    @staticmethod
+    def scan(text: str) -> dict[str, list[SecretsInfo]]:
+        """
+        Scans the given text for secrets using the detect-secrets library.
+
+        Args:
+            text (str): The text to scan for secrets.
+
+        Returns:
+            dict[str, list[SecretsInfo]]: A dictionary where the keys are secret types and the values are lists of SecretsInfo objects.
+        """
+        secrets = SecretsCollection()
+        temp_file = tempfile.NamedTemporaryFile(delete=False)
+        temp_file.write(text.encode("utf-8"))
+        temp_file.close()
+
+        with default_settings():
+            secrets.scan_file(str(temp_file.name))
+
+        unique_secrets = {}
+        for file in secrets.files:
+            for found_secret in secrets[file]:
+                if found_secret.secret_value is None:
+                    continue
+
+                secret_type = found_secret.type
+                actual_secret = found_secret.secret_value
+                line_number = found_secret.line_number
+
+                if secret_type not in unique_secrets:
+                    unique_secrets[secret_type] = []
+
+                unique_secrets[secret_type].append(
+                    SecretsInfo(secret=actual_secret, line_number=line_number)
+                )
+
+        os.remove(temp_file.name)
+        return unique_secrets
+
+    @weave.op
+    def invoke(self, text: str) -> dict[str, list[SecretsInfo]]:
+        """
+        Invokes the scan method to detect secrets in the given text.
+
+        Args:
+            text (str): The text to scan for secrets.
+
+        Returns:
+            dict[str, list[SecretsInfo]]: A dictionary where the keys are secret types and the values are lists of SecretsInfo objects.
+        """
+        return self.scan(text)
+
+
+class HyperScanModel(weave.Model):
+    """
+    Model for detecting secrets using the Hyperscan library.
+    We use the Hyperscan library to scan for secrets using regex patterns.
+    The patterns are mined from https://github.com/mazen160/secrets-patterns-db
+    This model is used in conjunction with the DetectSecretsModel to improve the detection of secrets.
+    """
+
+    _db: Any = PrivateAttr()
+    _pattern_map: dict[str, str] = PrivateAttr()
+    only_high_confidence: bool = False
+    ids: list[str] = []
+
+    def _load_patterns(self) -> dict[str, str]:
+        """
+        Loads the patterns from a JSONL file.
+
+        Returns:
+            dict[str, str]: A dictionary where the keys are pattern names and the values are regex patterns.
+        """
+        patterns = (
+            pathlib.Path(__file__).parent.resolve() / "secrets_patterns.jsonl"
+        ).open()
+        patterns_list = [json.loads(line) for line in patterns]
+        if self.only_high_confidence:
+            patterns_list = [
+                pattern for pattern in patterns_list if pattern["confidence"] == "high"
+            ]
+        return {pattern["name"]: pattern["regex"] for pattern in patterns_list}
+
+    def __init__(self, **kwargs: Any):
+        """
+        Initializes the HyperScanModel instance.
+        """
+        super().__init__(**kwargs)
+
+    def model_post_init(self, __context: Any) -> None:
+        """
+        Post-initialization method to load patterns and compile the Hyperscan database.
+        """
+        self._pattern_map = self._load_patterns()
+        self.ids = list(self._pattern_map.keys())
+        expressions = [pattern.encode() for pattern in self._pattern_map.values()]
+        self._db = hyperscan.Database()
+        self._db.compile(expressions=expressions, ids=list(range(len(expressions))))
+
+    def scan(self, text: str) -> dict[str, list[SecretsInfo]]:
+        """
+        Scans the given text for secrets using the Hyperscan library.
+
+        Args:
+            text (str): The text to scan for secrets.
+
+        Returns:
+            dict[str, list[SecretsInfo]]: A dictionary where the keys are secret types and the values are lists of SecretsInfo objects.
+        """
+        unique_secrets = {}
+
+        def on_match(idx, start, end, flags, context):
+            """
+            Callback function for handling matches found by Hyperscan.
+
+            Args:
+                idx: The index of the matched pattern.
+                start: The start position of the match.
+                end: The end position of the match.
+                flags: The flags associated with the match.
+                context: The context provided to the scan method.
+            """
+            secret = context["text"][start:end]
+            line_number = context["line_number"]
+            current_match = unique_secrets.setdefault(self.ids[idx], [])
+
+            if not current_match or len(secret) > len(current_match[0].secret):
+                unique_secrets[self.ids[idx]] = [
+                    SecretsInfo(line_number=line_number, secret=secret)
+                ]
+
+        for line_no, line in enumerate(text.splitlines(), start=1):
+            self._db.scan(
+                line.encode(),
+                match_event_handler=on_match,
+                context={"text": line, "line_number": line_no},
+            )
+
+        return unique_secrets
+
+    @weave.op
+    def invoke(self, text: str) -> dict[str, list[SecretsInfo]]:
+        """
+        Invokes the scan method to detect secrets in the given text.
+
+        Args:
+            text (str): The text to scan for secrets.
+
+        Returns:
+            dict[str, list[SecretsInfo]]: A dictionary where the keys are secret types and the values are lists of SecretsInfo objects.
+        """
+        return self.scan(text)
 
 
 class SecretsDetectionGuardrail(Guardrail):
     """
-    A guardrail for detecting secrets in text using regex patterns.
-    reference: SecretBench: A Dataset of Software Secrets
-    https://arxiv.org/abs/2303.06729
+    Guardrail class for secrets detection using both detect-secrets and Hyperscan models.
 
     Attributes:
-        regex_model (RegexModel): The regex model used for detection.
-        patterns (Union[dict[str, str], dict[str, list[str]]]): The patterns used for detection.
-        redaction (REDACTION): The type of redaction to apply.
+        redaction (REDACTION): The redaction mode to be applied.
+        _detect_secrets_model (Any): Instance of the DetectSecretsModel.
+        _hyperscan_model (Any): Instance of the HyperScanModel.
     """
 
-    regex_model: RegexModel
-    patterns: Union[dict[str, str], dict[str, list[str]]] = {}
     redaction: REDACTION
+    _detect_secrets_model: Any = PrivateAttr()
+    _hyperscan_model: Any = PrivateAttr()
+
+    def model_post_init(self, __context: Any) -> None:
+        """
+        Post-initialization method to initialize the detect-secrets and Hyperscan models.
+        """
+        self._detect_secrets_model = DetectSecretsModel()
+        self._hyperscan_model = HyperScanModel()
 
     def __init__(
         self,
-        use_defaults: bool = True,
         redaction: REDACTION = REDACTION.REDACT_ALL,
         **kwargs,
     ):
         """
-        Initialize the SecretsDetectionGuardrail.
+        Initializes the SecretsDetectionGuardrail instance.
 
         Args:
-            use_defaults (bool): Whether to use default patterns.
-            redaction (REDACTION): The type of redaction to apply.
+            redaction (REDACTION): The redaction mode to be applied. Defaults to REDACTION.REDACT_ALL.
             **kwargs: Additional keyword arguments.
         """
-        patterns = {}
-        if use_defaults:
-            patterns = DEFAULT_SECRETS_PATTERNS.copy()
-        if kwargs.get("patterns"):
-            patterns.update(kwargs["patterns"])
-
-        regex_model = RegexModel(patterns=patterns)
-
         super().__init__(
-            regex_model=regex_model,
-            patterns=patterns,
             redaction=redaction,
         )
 
-    @weave.op()
+    def get_modified_value(
+        self, unique_secrets: dict[str, Any], lines: list[str]
+    ) -> str:
+        """
+        Redacts the detected secrets in the given lines of text.
+
+        Args:
+            unique_secrets (dict[str, Any]): Dictionary of detected secrets.
+            lines (list[str]): List of lines of text.
+
+        Returns:
+            str: The modified text with secrets redacted.
+        """
+        for _, secrets_list in unique_secrets.items():
+            for secret_info in secrets_list:
+                secret = secret_info.secret
+                line_number = secret_info.line_number
+                lines[line_number - 1] = lines[line_number - 1].replace(
+                    secret, redact_value(secret, self.redaction)
+                )
+
+        modified_value = "\n".join(lines)
+        return modified_value
+
+    def get_scan_result(
+        self, unique_secrets: dict[str, list[SecretsInfo]], lines: list[str]
+    ) -> ScanResult | None:
+        """
+        Generates a ScanResult based on the detected secrets.
+
+        Args:
+            unique_secrets (dict[str, list[SecretsInfo]]): Dictionary of detected secrets.
+            lines (list[str]): List of lines of text.
+
+        Returns:
+            ScanResult | None: The scan result if secrets are detected, otherwise None.
+        """
+        if unique_secrets:
+            modified_value = self.get_modified_value(unique_secrets, lines)
+            detected_secrets = {
+                k: [i.secret for i in v] for k, v in unique_secrets.items()
+            }
+
+            return ScanResult(
+                **{
+                    "detected_secrets": detected_secrets,
+                    "modified_prompt": modified_value,
+                    "has_secret": True,
+                    "risk_score": 1.0,
+                }
+            )
+        return None
+
+    def scan(self, prompt: str) -> ScanResult:
+        """
+        Scans the given prompt for secrets using both detect-secrets and Hyperscan models.
+
+        Args:
+            prompt (str): The text to scan for secrets.
+
+        Returns:
+            ScanResult: The scan result with detected secrets and redacted text.
+        """
+        if prompt.strip() == "":
+            return ScanResult(
+                **{
+                    "detected_secrets": None,
+                    "modified_prompt": prompt,
+                    "has_secret": False,
+                    "risk_score": 0.0,
+                }
+            )
+
+        unique_secrets = self._detect_secrets_model.invoke(text=prompt)
+        results = self.get_scan_result(unique_secrets, prompt.splitlines())
+        if results:
+            return results
+
+        unique_secrets = self._hyperscan_model.invoke(text=prompt)
+        results = self.get_scan_result(unique_secrets, prompt.splitlines())
+        if results:
+            results.risk_score = 0.5
+            return results
+
+        return ScanResult(
+            **{
+                "detected_secrets": None,
+                "modified_prompt": prompt,
+                "has_secret": False,
+                "risk_score": 0.0,
+            }
+        )
+
+    @weave.op
     def guard(
         self,
         prompt: str,
@@ -159,40 +425,38 @@ class SecretsDetectionGuardrail(Guardrail):
         **kwargs,
     ) -> SecretsDetectionResponse | SecretsDetectionResponse:
         """
-        Check if the input prompt contains any secrets based on the regex patterns.
+        Guards the given prompt by scanning for secrets and optionally returning detected secrets.
 
         Args:
-            prompt (str): Input text to check for secrets.
-            return_detected_secrets (bool): If True, returns detailed secrets type information.
+            prompt (str): The text to scan for secrets.
+            return_detected_secrets (bool): Whether to return detected secrets in the response. Defaults to True.
+            **kwargs: Additional keyword arguments.
 
         Returns:
-            SecretsDetectionResponse or SecretsDetectionResponse: Detection results.
+            SecretsDetectionResponse | SecretsDetectionSimpleResponse: The response with scan results and redacted text.
         """
-        result = self.regex_model.check(prompt)
+        results = self.scan(prompt)
 
         explanation_parts = []
-        if result.matched_patterns:
+        if results.has_secret:
             explanation_parts.append("Found the following secrets in the text:")
-            for secret_type, matches in result.matched_patterns.items():
+            for secret_type, matches in results.detected_secrets.items():
                 explanation_parts.append(f"- {secret_type}: {len(matches)} instance(s)")
         else:
             explanation_parts.append("No secrets detected in the text.")
 
-        redacted_text = prompt
-        if result.matched_patterns:
-            for secret_type, matches in result.matched_patterns.items():
-                redacted_text = redact(redacted_text, matches, self.redaction)
-
         if return_detected_secrets:
             return SecretsDetectionResponse(
-                contains_secrets=not result.passed,
-                detected_secrets=result.matched_patterns,
+                contains_secrets=results.has_secret,
+                detected_secrets=results.detected_secrets,
                 explanation="\n".join(explanation_parts),
-                redacted_text=redacted_text,
+                redacted_text=results.modified_prompt,
+                risk_score=results.risk_score,
             )
         else:
             return SecretsDetectionSimpleResponse(
-                contains_secrets=not result.passed,
+                contains_secrets=not results.has_secret,
                 explanation="\n".join(explanation_parts),
-                redacted_text=redacted_text,
+                redacted_text=results.modified_prompt,
+                risk_score=results.risk_score,
             )

pyproject.toml

@@ -25,6 +25,10 @@ dependencies = [
     "presidio-analyzer>=2.2.355",
     "presidio-anonymizer>=2.2.355",
     "instructor>=1.7.0",
+    "numpy<2.0.0",
+    "gibberish-detector>=0.1.1",
+    "detect-secrets>=1.5.0",
+    "hyperscan>=0.7.8"
 ]
 
 [project.optional-dependencies]

tests/guardrails_genie/guardrails/test_secrets_detection.py

@@ -2,15 +2,14 @@ import hashlib
 import re
 
 import pytest
-from hypothesis import given, strategies as st, settings
+from hypothesis import strategies as st, given, settings
 
+from guardrails_genie.guardrails import SecretsDetectionGuardrail
 from guardrails_genie.guardrails.secrets_detection import (
-    DEFAULT_SECRETS_PATTERNS,
     SecretsDetectionSimpleResponse,
     SecretsDetectionResponse,
-    SecretsDetectionGuardrail,
     REDACTION,
-    redact,
+    redact_value,
 )
 
 
@@ -18,7 +17,7 @@ from guardrails_genie.guardrails.secrets_detection import (
 def mock_secrets_guard(monkeypatch):
     def _mock_guard(*args, **kwargs):
         prompt = kwargs.get("prompt")
-        return_detected_types = kwargs.get("return_detected_types")
+        return_detected_types = kwargs.get("return_detected_secrets")
 
         if "safe text" in prompt:
             if return_detected_types:
@@ -27,12 +26,14 @@ def mock_secrets_guard(monkeypatch):
                     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:
@@ -41,12 +42,14 @@ def mock_secrets_guard(monkeypatch):
                     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(
@@ -56,32 +59,22 @@ def mock_secrets_guard(monkeypatch):
 
 
 def test_redact_partial():
-    text = "My secret key is ABCDEFGHIJKL"
-    matches = ["ABCDEFGHIJKL"]
-    redacted_text = redact(text, matches, REDACTION.REDACT_PARTIAL)
-    assert redacted_text == "My secret key is [REDACTED:]AB..KL[:REDACTED]"
+    text = "ABCDEFGHIJKL"
+    redacted_text = redact_value(text, REDACTION.REDACT_PARTIAL)
+    assert redacted_text == "[REDACTED:]AB..KL[:REDACTED]"
 
 
 def test_redact_all():
-    text = "My secret key is ABCDEFGHIJKL"
-    matches = ["ABCDEFGHIJKL"]
-    redacted_text = redact(text, matches, REDACTION.REDACT_ALL)
-    assert redacted_text == "My secret key is [REDACTED:]************[:REDACTED]"
+    text = "ABCDEFGHIJKL"
+    redacted_text = redact_value(text, REDACTION.REDACT_ALL)
+    assert redacted_text == "[REDACTED:]************[:REDACTED]"
 
 
 def test_redact_hash():
-    text = "My secret key is ABCDEFGHIJKL"
-    matches = ["ABCDEFGHIJKL"]
-    hashed_value = hashlib.md5("ABCDEFGHIJKL".encode()).hexdigest()
-    redacted_text = redact(text, matches, REDACTION.REDACT_HASH)
-    assert redacted_text == f"My secret key is [REDACTED:]{hashed_value}[:REDACTED]"
-
-
-def test_redact_no_match():
-    text = "My secret key is ABCDEFGHIJKL"
-    matches = ["XYZ"]
-    redacted_text = redact(text, matches, REDACTION.REDACT_ALL)
-    assert redacted_text == text
+    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):
@@ -134,16 +127,15 @@ def test_secrets_detection_guardrail_no_secrets(mock_secrets_guard):
     assert result.redacted_text == prompt
 
 
-# Create a strategy to generate strings that match the patterns
 def pattern_strategy(pattern):
     return st.from_regex(re.compile(pattern), fullmatch=True)
 
 
-@settings(deadline=1000)  # Set the deadline to 1000 milliseconds (1 second)
-@given(pattern_strategy(DEFAULT_SECRETS_PATTERNS["JwtToken"][0]))
+@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 "JwtToken" in result.detected_secrets
+    assert "AWS Access Key" in result.detected_secrets