api/llms/key_utils.py

import hashlib
import base64

def generate_key(password: str) -> bytes:
    return hashlib.sha256(password.encode()).digest()

def xor_bytes(data: bytes, key: bytes) -> bytes:
    return bytes([b ^ key[i % len(key)] for i, b in enumerate(data)])

def encrypt_string(plain_text: str, password: str) -> str:
    key = generate_key(password)
    plain_bytes = plain_text.encode()
    encrypted_bytes = xor_bytes(plain_bytes, key)
    return base64.b64encode(encrypted_bytes).decode()

def decrypt_string(encrypted_text: str, password: str) -> str:
    key = generate_key(password)
    encrypted_bytes = base64.b64decode(encrypted_text)
    decrypted_bytes = xor_bytes(encrypted_bytes, key)
    return decrypted_bytes.decode()