diff --git a/utils/crypto.py b/utils/crypto.py index 642d7aa..9a2f9c1 100644 --- a/utils/crypto.py +++ b/utils/crypto.py @@ -1,3 +1,5 @@ +import zlib + from Crypto.Cipher import AES from Crypto.Random import get_random_bytes from base64 import urlsafe_b64encode, urlsafe_b64decode @@ -9,10 +11,13 @@ from db.schemas import UserData def encrypt_user_data(user_data: UserData) -> str: data = user_data.model_dump_json( exclude_none=True, exclude_defaults=True, exclude_unset=True, round_trip=True - ).encode("utf-8") + ) + compressed_data = zlib.compress(data.encode("utf-8")) iv = get_random_bytes(16) cipher = AES.new(settings.secret_key.encode("utf-8"), AES.MODE_CBC, iv) - encrypted_data = cipher.encrypt(data + b"\0" * (16 - len(data) % 16)) + encrypted_data = cipher.encrypt( + compressed_data + b"\0" * (16 - len(compressed_data) % 16) + ) encrypted_str = urlsafe_b64encode(iv + encrypted_data).decode("utf-8") return encrypted_str @@ -26,7 +31,8 @@ def decrypt_user_data(secret_str: str | None = None) -> UserData: cipher = AES.new(settings.secret_key.encode("utf-8"), AES.MODE_CBC, iv) decrypted_data = cipher.decrypt(encrypted_data[16:]) decrypted_data = decrypted_data.rstrip(b"\0") - user_data = UserData.model_validate_json(decrypted_data.decode("utf-8")) + decompressed_data = zlib.decompress(decrypted_data) + user_data = UserData.model_validate_json(decompressed_data.decode("utf-8")) except Exception: user_data = UserData() return user_data