guardrails_genie/guardrails/secrets_detection/secrets_detection.py

import hashlib
import json
import pathlib
from enum import Enum
from typing import Union, Optional

import weave
from pydantic import BaseModel

from guardrails_genie.guardrails.base import Guardrail
from guardrails_genie.regex_model import RegexModel


def load_secrets_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


DEFAULT_SECRETS_PATTERNS = load_secrets_patterns()


class REDACTION(str, Enum):
    REDACT_PARTIAL = "REDACT_PARTIAL"
    REDACT_ALL = "REDACT_ALL"
    REDACT_HASH = "REDACT_HASH"
    REDACT_NONE = "REDACT_NONE"


def redact(text: str, matches: list[str], redaction_type: REDACTION) -> str:
    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


class SecretsDetectionSimpleResponse(BaseModel):
    contains_secrets: bool
    explanation: str
    redacted_text: Optional[str] = None

    @property
    def safe(self) -> bool:
        return not self.contains_entities


class SecretsDetectionResponse(SecretsDetectionSimpleResponse):
    detected_secrets: dict[str, list[str]]


class SecretsDetectionGuardrail(Guardrail):
    regex_model: RegexModel
    patterns: Union[dict[str, str], dict[str, list[str]]] = {}
    redaction: REDACTION

    def __init__(
        self,
        use_defaults: bool = True,
        redaction: REDACTION = REDACTION.REDACT_ALL,
        **kwargs,
    ):
        patterns = {}
        if use_defaults:
            patterns = DEFAULT_SECRETS_PATTERNS.copy()
        if kwargs.get("patterns"):
            patterns.update(kwargs["patterns"])

        # Create the RegexModel instance
        regex_model = RegexModel(patterns=patterns)

        # Initialize the base class with both the regex_model and patterns
        super().__init__(
            regex_model=regex_model,
            patterns=patterns,
            redaction=redaction,
        )

    @weave.op()
    def guard(
        self,
        prompt: str,
        return_detected_types: bool = True,
        **kwargs,
    ) -> SecretsDetectionResponse | SecretsDetectionResponse:
        """
        Check if the input prompt contains any entities based on the regex patterns.

        Args:
            prompt: Input text to check for entities
            return_detected_types: If True, returns detailed entity type information

        Returns:
            SecretsDetectionResponse or SecretsDetectionResponse containing detection results
        """
        result = self.regex_model.check(prompt)

        # Create detailed explanation
        explanation_parts = []
        if result.matched_patterns:
            explanation_parts.append("Found the following secrets in the text:")
            for secret_type, matches in result.matched_patterns.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_types:
            return SecretsDetectionResponse(
                contains_secrets=not result.passed,
                detected_secrets=result.matched_patterns,
                explanation="\n".join(explanation_parts),
                redacted_text=redacted_text,
            )
        else:
            return SecretsDetectionSimpleResponse(
                contains_entities=not result.passed,
                explanation="\n".join(explanation_parts),
                redacted_text=redacted_text,
            )


def main():
    weave.init(project_name="parambharat/guardrails-genie")

    guardrail = SecretsDetectionGuardrail(redaction=REDACTION.REDACT_ALL)
    dataset = [
        {
            "input": 'I need to pass a key\naws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"',
        },
        {
            "input": "My github token is: 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",
        },
        {
            "input": "My JWT token is: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c",
        },
    ]

    for item in dataset:
        # Check text for entities
        result = guardrail.guard(prompt=item["input"])

        # Access results
        print(f"Contains entities: {result.contains_secrets}")
        print(f"Detected entities: {result.detected_secrets}")
        print(f"Explanation: {result.explanation}")
        print(f"Anonymized text: {result.redacted_text}")
    # import regex as re
    #
    # sample_input = "My JWT token is: eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJzdWIiOiIxMjM0NTY3ODkwIiwibmFtZSI6IkpvaG4gRG9lIiwiaWF0IjoxNTE2MjM5MDIyfQ.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c"
    # jwt_pattern = DEFAULT_SECRETS_PATTERNS["JwtToken"][0]
    # print(jwt_pattern)
    # pattern = re.compile(jwt_pattern)
    # print(pattern)
    # print(pattern.findall(sample_input))

    # import pandas as pd
    #
    # df = pd.read_json("secrets_patterns_bak.jsonl", lines=True)
    # df.loc[:, "patterns"] = df["patterns"].map(lambda x: [i[2:-1] for i in x])
    # df.to_json("secrets_patterns.jsonl", orient="records", lines=True)


if __name__ == "__main__":
    main()