# requires openai/anthropic to run; didnt implement those. the script is just for guardrail review so that isnt nessecary

from pydantic import BaseModel, StringConstraints
from typing import Annotated

class LLMQuery(BaseModel, validate_assignment=True):
    prompt: Annotated[str, StringConstraints(pattern=r"^[0-9.\+\-\*/\(\)]+$")]
    response: Annotated[str, StringConstraints(pattern=r"^[0-9\.]+$")] = None

SYSTEM_PROMPT = '''You are a calculator. Please compute the result of the following mathematical expression.
Only respond with the result, no other text.

'''

import json

PRODUCTS = [
    'SnaccPayload',
    'PowerChex',
    'CrunchPacket',
    'ByteChew',
    'CryptoChunks'
]

# https://gist.githubusercontent.com/mspublic/f942cb80a83a7880cff689def8334c7b/raw/30527efde9e1d2b11ede9d9bf205d98c2ba9d550/Instruction-Bypass.yara
PROMPT_INJECTION_REGEX = r"^.*(Ignore|Disregard|Skip|Forget|Neglect|Overlook|Omit|Bypass)\s*(prior|previous|preceding|above|foregoing|earlier|initial)?\s*(content|text|instructions|instruction|directives|directive|commands|command).*$"

# https://gist.githubusercontent.com/mspublic/f942cb80a83a7880cff689def8334c7b/raw/30527efde9e1d2b11ede9d9bf205d98c2ba9d550/jailbreaks.json
JAILBREAKS = json.load(open('./jailbreaks.json')).get('jailbreak')

# check if any item in the passed list is contained in the passed string
def list_check(list, string):
    return any(item in string for item in list)

# detect jailbreak based on string similarity
from difflib import SequenceMatcher
def jailbreak_detect(prompt, jailbreak_threshold=0.5):
    return any(SequenceMatcher(None, jailbreak, prompt).ratio() > jailbreak_threshold for jailbreak in JAILBREAKS)

class GuardrailPromptException(Exception):
    pass

class GuardrailResponseException(Exception):
    pass


from pydantic import BaseModel, field_validator
import re

class LLMQuery(BaseModel, validate_assignment=True):
    prompt: str
    response: str = None

    @field_validator("prompt")
    @classmethod
    def validate_prompt(cls, prompt: str) -> str:
        prompt = prompt.strip()

        # Check for PacketSnacc product
        if not list_check(PRODUCTS, prompt):
            raise GuardrailPromptException("No PacketSnacc product mentioned in user prompt.")

        # Check for prompt injection
        if re.search(PROMPT_INJECTION_REGEX, prompt, re.IGNORECASE):
            raise GuardrailPromptException("Prompt injection attempt detected.")

        # Check for jailbreaking
        if jailbreak_detect(prompt):
            raise GuardrailPromptException("Jailbreak attempt detected.")

        return prompt

    @field_validator("response")
    @classmethod
    def validate_response(cls, response: str) -> str:
        # Check for competitors
        if list_check(COMPETITORS, response):
            raise GuardrailResponseException("Output does not satisfy company policy.")

        # Check for profanity
        if list_check(PROFANITY, response):
            raise GuardrailResponseException("Profane language detected.")

        # Check for credit card information
        if re.search(r"^.*[0-9]{13,19}.*$", response):
            raise GuardrailResponseException("Information leakage detected.")

        # Remove HTML-Tags
        response = re.sub(r"<.*?>", "", response)

        return response

COMPETITORS = [
    "SnackOverflow",
    "NullBite",
    "CyberChow"
]

# https://raw.githubusercontent.com/zacanger/profane-words/refs/heads/master/words.json
PROFANITY = json.load(open('./words.json'))

def query_llm(system_prompt:str, prompt: str) -> str:
    # openai/anthorpic etc; not relevent
    return response

def protected_query_llm(prompt: str) -> LLMQuery:
    query = LLMQuery(prompt=prompt)
    query.response = query_llm(SYSTEM_PROMPT, query.prompt)
    return query


while 1:
    try:
        prompt = input("> ")
        query_obj = protected_query_llm(prompt)
        print(query_obj.response)
    except Exception as e:
        print(f'Error: {e}')