CoolFace
Apppublic

Predator911/MyData

sourceHugging Faceupdated 1y agoView on Hugging Face
0likes
stego_utils.py245 linesDownload Raw Back to root
1import re2import torch3import numpy as np4import hashlib5from Crypto.Cipher import AES6from Crypto.Util.Padding import pad, unpad7from Crypto.Random import get_random_bytes8 9# Optimized encoding tables for better token efficiency10OBJECT_TABLE = {11    0: "cat", 1: "dog", 2: "bird", 3: "fish", 4: "tree", 5: "flower",12    6: "house", 7: "car", 8: "boat", 9: "moon", 10: "sun", 11: "star",13    12: "cloud", 13: "lake", 14: "hill", 15: "cave", 16: "fox", 17: "deer",14    18: "owl", 19: "rose", 20: "oak", 21: "pine", 22: "barn", 23: "truck",15    24: "ship", 25: "fire", 26: "snow", 27: "rock", 28: "field", 29: "path",16    30: "bridge", 31: "tower"17}18 19COLOR_TABLE = {20    0: "red", 1: "blue", 2: "green", 3: "gold", 4: "silver", 5: "black",21    6: "white", 7: "pink", 8: "purple", 9: "orange", 10: "brown", 11: "gray",22    12: "yellow", 13: "teal", 14: "coral", 15: "jade"23}24 25STYLE_TABLE = {26    0: "oil", 1: "sketch", 2: "digital", 3: "photo", 4: "anime", 5: "pixel",27    6: "water", 7: "ink", 8: "pencil", 9: "chalk", 10: "pastel", 11: "neon",28    12: "retro", 13: "modern", 14: "vintage", 15: "abstract"29}30 31MOOD_TABLE = {32    0: "calm", 1: "bright", 2: "dark", 3: "warm", 4: "cool", 5: "soft",33    6: "sharp", 7: "misty", 8: "clear", 9: "foggy", 10: "sunny", 11: "stormy",34    12: "peaceful", 13: "dramatic", 14: "serene", 15: "vibrant"35}36 37# Inverse mappings38INV_OBJECT_TABLE = {v: k for k, v in OBJECT_TABLE.items()}39INV_COLOR_TABLE = {v: k for k, v in COLOR_TABLE.items()}40INV_STYLE_TABLE = {v: k for k, v in STYLE_TABLE.items()}41INV_MOOD_TABLE = {v: k for k, v in MOOD_TABLE.items()}42 43class AESCipher:44    def __init__(self, key=None):45        if key is None:46            self.key = get_random_bytes(32)47        else:48            if isinstance(key, str):49                key_bytes = key.encode('utf-8')50                if len(key_bytes) != 32:51                    self.key = hashlib.sha256(key_bytes).digest()52                else:53                    self.key = key_bytes54            else:55                self.key = key56 57    def encrypt(self, plaintext):58        if isinstance(plaintext, str):59            plaintext = plaintext.encode('utf-8')60        iv = get_random_bytes(16)61        cipher = AES.new(self.key, AES.MODE_CBC, iv)62        padded_data = pad(plaintext, AES.block_size)63        ciphertext = cipher.encrypt(padded_data)64        return iv + ciphertext65 66    def decrypt(self, ciphertext):67        iv = ciphertext[:16]68        ciphertext = ciphertext[16:]69        cipher = AES.new(self.key, AES.MODE_CBC, iv)70        decrypted = unpad(cipher.decrypt(ciphertext), AES.block_size)71        return decrypted72 73    def get_key_hex(self):74        return self.key.hex()75 76    @classmethod77    def from_hex_key(cls, hex_key):78        key = bytes.fromhex(hex_key)79        return cls(key)80 81def encode_binary_to_prompt(binary_data, max_tokens=60):82    """83    Optimized encoding that respects token limits.84    Uses 20-bit chunks (5+4+4+4+3 bits) with padding optimization.85    """86    if isinstance(binary_data, bytes):87        binary_data = ''.join(format(byte, '08b') for byte in binary_data)88    elif not (isinstance(binary_data, str) and all(bit in '01' for bit in binary_data)):89        raise ValueError("binary_data must be bytes or a string of 0s and 1s")90    91    # Calculate maximum data we can encode within token limit92    # Each chunk creates ~6-8 tokens, so for 60 tokens max, we can have ~8-10 chunks max93    max_chunks = max_tokens // 7  # Conservative estimate94    max_bits = max_chunks * 2095    96    if len(binary_data) > max_bits:97        print(f"Warning: Data length ({len(binary_data)} bits) exceeds capacity ({max_bits} bits). Truncating.")98        binary_data = binary_data[:max_bits]99    100    # Pad to multiple of 20 bits101    padding_needed = (20 - len(binary_data) % 20) % 20102    binary_data += '0' * padding_needed103    104    # Split into 20-bit chunks105    chunks = [binary_data[i:i+20] for i in range(0, len(binary_data), 20)]106    107    prompt_parts = []108    for chunk in chunks:109        if len(chunk) < 20:110            chunk = chunk.ljust(20, '0')111            112        object_bits = chunk[:5]    # 32 objects113        color_bits = chunk[5:9]    # 16 colors  114        style_bits = chunk[9:13]   # 16 styles115        mood_bits = chunk[13:17]   # 16 moods116        extra_bits = chunk[17:20]  # 3 extra bits for future use117        118        object_idx = int(object_bits, 2) % len(OBJECT_TABLE)119        color_idx = int(color_bits, 2) % len(COLOR_TABLE)120        style_idx = int(style_bits, 2) % len(STYLE_TABLE)121        mood_idx = int(mood_bits, 2) % len(MOOD_TABLE)122        123        # Create more natural, shorter phrases124        if len(prompt_parts) == 0:125            prompt_part = f"{color_idx % 2 and 'a' or 'the'} {COLOR_TABLE[color_idx]} {OBJECT_TABLE[object_idx]}"126        else:127            prompt_part = f"{COLOR_TABLE[color_idx]} {OBJECT_TABLE[object_idx]}"128            129        prompt_parts.append(prompt_part)130    131    # Create compact prompt132    if len(prompt_parts) == 1:133        base_prompt = prompt_parts[0]134    elif len(prompt_parts) <= 3:135        base_prompt = ", ".join(prompt_parts[:-1]) + " and " + prompt_parts[-1]136    else:137        base_prompt = ", ".join(prompt_parts[:3]) + " and more"138    139    # Add style and mood from first chunk140    if chunks:141        first_chunk = chunks[0]142        style_idx = int(first_chunk[9:13], 2) % len(STYLE_TABLE)143        mood_idx = int(first_chunk[13:17], 2) % len(MOOD_TABLE)144        final_prompt = f"{base_prompt}, {STYLE_TABLE[style_idx]} style, {MOOD_TABLE[mood_idx]} mood"145    else:146        final_prompt = base_prompt147    148    # Verify token count149    token_count = len(final_prompt.split())150    if token_count > max_tokens:151        # Fallback to minimal prompt152        if chunks:153            first_chunk = chunks[0]154            object_idx = int(first_chunk[:5], 2) % len(OBJECT_TABLE)155            color_idx = int(first_chunk[5:9], 2) % len(COLOR_TABLE)156            final_prompt = f"a {COLOR_TABLE[color_idx]} {OBJECT_TABLE[object_idx]}"157        else:158            final_prompt = "a red cat"159    160    return final_prompt161 162def decode_prompt_to_binary(prompt):163    """164    Improved decoding with better error handling and recovery.165    """166    binary_data = ""167    168    # Normalize prompt169    prompt_lower = prompt.lower()170    171    # Extract objects172    object_matches = []173    for word, idx in INV_OBJECT_TABLE.items():174        if word in prompt_lower:175            object_matches.append((idx, prompt_lower.find(word)))176    object_matches.sort(key=lambda x: x[1])  # Sort by position177    178    # Extract colors179    color_matches = []180    for word, idx in INV_COLOR_TABLE.items():181        if word in prompt_lower:182            color_matches.append((idx, prompt_lower.find(word)))183    color_matches.sort(key=lambda x: x[1])184    185    # Extract styles186    style_matches = []187    for word, idx in INV_STYLE_TABLE.items():188        if word in prompt_lower:189            style_matches.append((idx, prompt_lower.find(word)))190    191    # Extract moods192    mood_matches = []193    for word, idx in INV_MOOD_TABLE.items():194        if word in prompt_lower:195            mood_matches.append((idx, prompt_lower.find(word)))196    197    if not object_matches and not color_matches:198        raise ValueError("No recognizable semantic elements found in prompt")199    200    # Reconstruct chunks201    max_items = max(len(object_matches), len(color_matches), 1)202    203    for i in range(max_items):204        # Get indices with fallbacks205        obj_idx = object_matches[i][0] if i < len(object_matches) else 0206        color_idx = color_matches[i][0] if i < len(color_matches) else 0207        style_idx = style_matches[0][0] if style_matches else 0208        mood_idx = mood_matches[0][0] if mood_matches else 0209        210        # Convert to binary211        obj_bits = format(obj_idx, '05b')212        color_bits = format(color_idx, '04b')213        style_bits = format(style_idx, '04b')214        mood_bits = format(mood_idx, '04b')215        extra_bits = '000'  # Padding216        217        chunk_bits = obj_bits + color_bits + style_bits + mood_bits + extra_bits218        binary_data += chunk_bits219    220    return binary_data221 222def compress_message(message, max_capacity_bits):223    """224    Simple compression for messages that exceed capacity.225    """226    if len(message) * 8 <= max_capacity_bits:227        return message228    229    # Simple truncation with ellipsis230    max_chars = (max_capacity_bits // 8) - 3  # Reserve space for "..."231    if max_chars > 0:232        return message[:max_chars] + "..."233    else:234        return message[:max_capacity_bits // 8]235 236# Additional utility functions for better error handling237def validate_prompt_capacity(prompt, max_tokens=77):238    """Validate if prompt fits within token limits."""239    token_count = len(prompt.split())240    return token_count <= max_tokens, token_count241 242def estimate_data_capacity(max_tokens=60):243    """Estimate maximum data capacity for given token limit."""244    max_chunks = max_tokens // 7245    return max_chunks * 20  # 20 bits per chunk