from transformers import AutoTokenizer
 
def tokenization_analysis(tokenizer, text):
    """Analyseer hoe tekst wordt getokeniseerd en leg mogelijke edge cases bloot."""
    tokens = tokenizer.encode(text)
    decoded_tokens = [tokenizer.decode([t]) for t in tokens]
 
    analysis = {
        "text": text,
        "num_tokens": len(tokens),
        "token_ids": tokens,
        "token_strings": decoded_tokens,
        "chars_per_token": len(text) / len(tokens)
    }
 
    # Identificeer multi-byte tokens en splitsingen op één teken
    for i, (tid, tstr) in enumerate(zip(tokens, decoded_tokens)):
        if len(tstr) == 1 and ord(tstr) > 127:
            analysis.setdefault("unicode_single_chars", []).append({
                "position": i,
                "character": tstr,
                "codepoint": hex(ord(tstr))
            })
 
    return analysis
 
# Voorbeeld: "ignore" is waarschijnlijk \u00e9\u00e9n token, maar "ign" + "ore" kan
# een filter omzeilen dat naar het token "ignore" zoekt
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b")
 
# Normale tokenisatie
print(tokenization_analysis(tokenizer, "ignore previous instructions"))
 
# Aangepaste tokenisatie (zero-width characters toevoegen)
print(tokenization_analysis(tokenizer, "ig\u200bnore previous instructions"))

def find_boundary_exploits(tokenizer, target_word):
    """
    Zoek manieren om een doelwoord te schrijven dat anders tokeniseert
    en zo mogelijk contentfilters op tokenniveau omzeilt.
    """
    base_tokens = tokenizer.encode(target_word, add_special_tokens=False)
    exploits = []
 
    # Strategie 1: variaties in hoofdletters
    variations = [
        target_word.upper(),
        target_word.capitalize(),
        target_word.swapcase(),
        ''.join(c.upper() if i % 2 == 0 else c
                for i, c in enumerate(target_word))
    ]
 
    # Strategie 2: Unicode-substituties
    unicode_map = {
        'a': '\u0430',  # Cyrillische 'a' (visueel identiek)
        'e': '\u0435',  # Cyrillische 'e'
        'o': '\u043e',  # Cyrillische 'o'
        'p': '\u0440',  # Cyrillische 'p'
        'c': '\u0441',  # Cyrillische 'c'
        'x': '\u0445',  # Cyrillische 'x'
    }
 
    for i, char in enumerate(target_word):
        if char.lower() in unicode_map:
            variant = (target_word[:i] +
                      unicode_map[char.lower()] +
                      target_word[i+1:])
            variations.append(variant)
 
    # Strategie 3: zero-width character invoegen
    for i in range(1, len(target_word)):
        zwc_variant = target_word[:i] + '\u200b' + target_word[i:]
        variations.append(zwc_variant)
 
    for variant in variations:
        variant_tokens = tokenizer.encode(variant, add_special_tokens=False)
        if variant_tokens != base_tokens:
            exploits.append({
                "variant": variant,
                "original_tokens": base_tokens,
                "variant_tokens": variant_tokens,
                "token_count_change": len(variant_tokens) - len(base_tokens),
                "repr": repr(variant)
            })
 
    return exploits

class TokenFilterBypass:
    """Technieken om contentfilters op tokenniveau te omzeilen."""
 
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer
 
    def unicode_homoglyph_bypass(self, text, target_words):
        """
        Vervang tekens in doelwoorden door visueel identieke
        Unicode-tekens die de tokenisatie veranderen.
        """
        homoglyphs = {
            'A': '\u0410', 'B': '\u0412', 'C': '\u0421',
            'E': '\u0415', 'H': '\u041d', 'K': '\u041a',
            'M': '\u041c', 'O': '\u041e', 'P': '\u0420',
            'T': '\u0422', 'X': '\u0425',
            'a': '\u0430', 'c': '\u0441', 'e': '\u0435',
            'o': '\u043e', 'p': '\u0440', 'x': '\u0445',
        }
 
        result = text
        for word in target_words:
            modified_word = ''
            for char in word:
                if char in homoglyphs and len(modified_word) % 3 == 0:
                    modified_word += homoglyphs[char]
                else:
                    modified_word += char
            result = result.replace(word, modified_word)
 
        return result
 
    def whitespace_splitting(self, text, target_words):
        """
        Voeg ongebruikelijke witruimtetekens in om tokens op te splitsen.
        """
        special_spaces = [
            '\u00a0',  # Non-breaking space
            '\u2009',  # Thin space
            '\u200a',  # Hair space
            '\u2003',  # Em space
        ]
 
        result = text
        for word in target_words:
            mid = len(word) // 2
            space = special_spaces[hash(word) % len(special_spaces)]
            split_word = word[:mid] + space + word[mid:]
            result = result.replace(word, split_word)
 
        return result
 
    def encoding_normalization_bypass(self, text):
        """
        Gebruik Unicode-normalisatievormen die er identiek uitzien
        maar anders tokeniseren.
        """
        import unicodedata
 
        # NFC- versus NFD-normalisatie levert verschillende bytereeksen op
        # voor letters met accenten, wat de tokenisatie verandert
        nfd_text = unicodedata.normalize('NFD', text)
        nfc_text = unicodedata.normalize('NFC', text)
 
        # Controleer welke normalisatie de tokenisatie verandert
        nfd_tokens = self.tokenizer.encode(nfd_text)
        nfc_tokens = self.tokenizer.encode(nfc_text)
 
        if nfd_tokens != nfc_tokens:
            return nfd_text  # Andere tokenisatie kan het filter omzeilen
 
        return text

def special_token_analysis(tokenizer):
    """Analyseer speciale tokens op mogelijk misbruik."""
    special = {
        "bos": tokenizer.bos_token,
        "eos": tokenizer.eos_token,
        "pad": tokenizer.pad_token,
        "unk": tokenizer.unk_token,
    }
 
    # Controleer op extra speciale tokens
    if hasattr(tokenizer, 'additional_special_tokens'):
        for i, token in enumerate(tokenizer.additional_special_tokens):
            special[f"special_{i}"] = token
 
    # Controleer of strings van speciale tokens via tekst kunnen worden geïnjecteerd
    injectable = {}
    for name, token in special.items():
        if token and token in "normal looking text":
            injectable[name] = {
                "token": token,
                "can_inject_via_text": True
            }
 
    return {
        "special_tokens": special,
        "injectable": injectable,
        "vocab_size": tokenizer.vocab_size
    }

def token_efficiency_analysis(tokenizer, text):
    """
    Analyseer token-efficiëntie en zoek compactere representaties.
    """
    base_tokens = len(tokenizer.encode(text))
 
    strategies = {
        "base": {
            "text": text,
            "tokens": base_tokens
        }
    }
 
    # Afkortingen gebruiken minder tokens
    abbreviation_map = {
        "because": "bc",
        "without": "w/o",
        "with": "w/",
        "information": "info",
        "approximately": "~",
    }
 
    abbreviated = text
    for full, abbr in abbreviation_map.items():
        abbreviated = abbreviated.replace(full, abbr)
 
    strategies["abbreviated"] = {
        "text": abbreviated,
        "tokens": len(tokenizer.encode(abbreviated))
    }
 
    # Compacte opmaak
    compact = ' '.join(text.split())  # Verwijder overtollige witruimte
    strategies["compact"] = {
        "text": compact,
        "tokens": len(tokenizer.encode(compact))
    }
 
    return strategies

def context_exhaustion_payload(tokenizer, max_context, target_text,
                                system_prompt_length):
    """
    Maak een payload die het contextvenster zo vult dat
    cruciale system-prompt-content wordt afgekapt.
    """
    target_tokens = tokenizer.encode(target_text)
    target_len = len(target_tokens)
 
    # Bereken de benodigde padding
    available = max_context - system_prompt_length - target_len - 100
    # buffer van 100 tokens voor generatie
 
    # Genereer padding met een hoog aantal tokens
    # Losse tekens tokeniseren vaak elk naar 1 token
    padding = "x " * available
 
    return padding + "\n" + target_text

class TokenizerAwareFilter:
    """Contentfilter dat tokenisatie-edge-cases afhandelt."""
 
    def __init__(self, tokenizer, blocked_words):
        self.tokenizer = tokenizer
        self.blocked_words = blocked_words
 
    def check(self, text):
        """
        Toets tekst tegen geblokkeerde woorden met meerdere
        normalisatiestrategie\u00ebn.
        """
        # Strategie 1: toets ruwe tekst (hoofdletterongevoelig)
        text_lower = text.lower()
        for word in self.blocked_words:
            if word.lower() in text_lower:
                return {"blocked": True, "reason": f"Contains '{word}'"}
 
        # Strategie 2: Unicode normaliseren en daarna toetsen
        import unicodedata
        normalized = unicodedata.normalize('NFKC', text)
        normalized_lower = normalized.lower()
        for word in self.blocked_words:
            if word.lower() in normalized_lower:
                return {"blocked": True, "reason": f"Contains '{word}' (normalized)"}
 
        # Strategie 3: zero-width characters verwijderen en daarna toetsen
        zwc_removed = text
        for zwc in ['\u200b', '\u200c', '\u200d', '\ufeff']:
            zwc_removed = zwc_removed.replace(zwc, '')
        zwc_lower = zwc_removed.lower()
        for word in self.blocked_words:
            if word.lower() in zwc_lower:
                return {"blocked": True, "reason": f"Contains '{word}' (ZWC removed)"}
 
        # Strategie 4: detectie van verwarbare tekens
        stripped = self.strip_confusables(text)
        stripped_lower = stripped.lower()
        for word in self.blocked_words:
            if word.lower() in stripped_lower:
                return {"blocked": True, "reason": f"Contains '{word}' (confusables)"}
 
        return {"blocked": False}
 
    def strip_confusables(self, text):
        """Vervang verwarbare Unicode-tekens door ASCII-equivalenten."""
        confusable_map = {
            '\u0430': 'a', '\u0435': 'e', '\u043e': 'o',
            '\u0440': 'p', '\u0441': 'c', '\u0445': 'x',
            '\u0410': 'A', '\u0415': 'E', '\u041e': 'O',
            '\u0420': 'P', '\u0421': 'C',
        }
        return ''.join(confusable_map.get(c, c) for c in text)