Improve encryption

Author: VerdantFox

app/services/general/encryption_handler.py

@@ -1,112 +1,37 @@
 """encryption_helpers: functions for encrypting and decrypting data."""
 
-import base64
 import hashlib
-
-from cryptography.hazmat.backends import default_backend
-from cryptography.hazmat.primitives import padding
-from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+import hmac
 
 from app.settings import settings
 
 
-def derive_iv(plaintext: str) -> bytes:
-    """Derive a consistent IV from the plaintext using a hash function.
-
-    Args:
-    ----
-        plaintext (str): The plaintext from which to derive the IV.
-
-    Returns:
-    -------
-        bytes: A 16-byte IV derived from the plaintext.
-
-    """
-    return hashlib.sha256(plaintext.encode("utf-8")).digest()[:16]
+def hash_token(token: str) -> str:
+    """Return an HMAC-SHA256 hex digest of `token` keyed with the encryption key.
 
+    Using HMAC rather than a bare SHA-256 hash prevents offline rainbow-table
+    attacks: an attacker who obtains the DB rows cannot pre-compute hashes
+    without also knowing the secret key.
 
-def encrypt(data: str) -> str:
-    """Encrypt the given data using AES encryption and encodes it in Base64.
+    The token is a UUID (128 bits of entropy), so HMAC-SHA256 is sufficient —
+    a slow KDF such as bcrypt is not required.
 
     Args:
     ----
-        data (str): The data to be encrypted.
+        token: The plaintext token (UUID string) to hash.
 
     Returns:
     -------
-        str: The Base64 encoded encrypted data including the IV.
+        A 64-character lowercase hex string (256-bit HMAC digest).
 
     """
     key_bytes = bytes.fromhex(settings.encryption_key)
-    data_bytes = data.encode("utf-8")
-    iv = derive_iv(data)
-    cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend())
-    encryptor = cipher.encryptor()
-
-    # Pad data to be a multiple of block size
-    padder = padding.PKCS7(algorithms.AES.block_size).padder()
-    padded_data = padder.update(data_bytes) + padder.finalize()
-
-    encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
-    encrypted_data_with_iv = iv + encrypted_data
-
-    # Encode the encrypted data with IV in Base64
-    return base64.b64encode(encrypted_data_with_iv).decode("utf-8")
-
-
-def decrypt(encoded_data: str) -> str:
-    """Decrypt the given Base64 encoded encrypted data using AES decryption.
-
-    Args:
-    ----
-        encoded_data (str): The Base64 encoded encrypted data including the IV.
-
-    Returns:
-    -------
-        str: The decrypted data as a string.
-
-    """
-    key_bytes = bytes.fromhex(settings.encryption_key)
-    encrypted_data_with_iv = base64.b64decode(encoded_data)
-
-    iv = encrypted_data_with_iv[:16]
-    encrypted_data = encrypted_data_with_iv[16:]
-
-    cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend())
-    decryptor = cipher.decryptor()
-
-    padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
-
-    # Unpad data
-    unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
-    data_bytes = unpadder.update(padded_data) + unpadder.finalize()
-
-    return data_bytes.decode("utf-8")
+    return hmac.new(key_bytes, token.encode("utf-8"), hashlib.sha256).hexdigest()
 
 
 if __name__ == "__main__":  # pragma: no cover
-    # Encryption key as a 64-character hex string (256 bits / 32 bytes)
-    # key = os.urandom(32).hex()  # noqa: ERA001 (commented-out code)
-    print(f"Key (hex): {settings.encryption_key}")  # noqa: T201 (print used for example)
-
     import uuid
 
-    # Original plaintext data
     data = str(uuid.uuid4())
-    print("Original Data:", data)  # noqa: T201 (print used for example)
-
-    # Encrypt the data to get the encrypted value
-    encrypted_data = encrypt(data)
-    print("Encrypted Data:", encrypted_data)  # noqa: T201 (print used for example)
-
-    # Decrypt the data to get the original value back
-    decrypted_data = decrypt(encrypted_data)
-    print("Decrypted Data:", decrypted_data)  # noqa: T201 (print used for example)
-
-    # To search for the encrypted data if you know the decrypted data
-    known_decrypted_data = data
-    encrypted_search_value = encrypt(known_decrypted_data)
-    print("Encrypted Search Value:", encrypted_search_value)  # noqa: T201 (print used for example)
-
-    # Check if the encrypted data matches the search value
-    print("Match found:", encrypted_data == encrypted_search_value)  # noqa: T201 (print used for example)
+    print("Token:", data)  # noqa: T201 (print used for example)
+    print("Hashed:", hash_token(data))  # noqa: T201 (print used for example)

app/services/users/user_handler.py

@@ -178,7 +178,7 @@ async def create_pw_reset_token(
     query = str(uuid4())
     pw_reset_token = db_models.PasswordResetToken(
         user_id=user.id,
-        encrypted_query=encryption_handler.encrypt(query),
+        encrypted_query=encryption_handler.hash_token(query),
         created_timestamp=datetime.now(UTC),
         expires_timestamp=datetime.now(UTC) + timedelta(minutes=PW_RESET_TOKEN_EXPIRATION_MINUTES),
     )
@@ -190,7 +190,7 @@ async def create_pw_reset_token(
 
 async def assert_token_is_valid(db: AsyncSession, query: str) -> db_models.PasswordResetToken:
     """Get a password reset token by query."""
-    encrypted_query = encryption_handler.encrypt(query)
+    encrypted_query = encryption_handler.hash_token(query)
 
     stmt = select(db_models.PasswordResetToken).where(
         db_models.PasswordResetToken.encrypted_query == encrypted_query