Switched to the ChaCha20Poly1305 encryption algorithm for better security

src/encryption/accounts.py

@@ -1,17 +1,22 @@
-import base64
 import os
+from typing import Self
 
-from cryptography.fernet import Fernet
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import hashes
+from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
 from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
 
 from ..db.models import Account
 from ..decrypted_account import DecryptedAccount
 
 
-def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
-    """Generates key for fernet encryption"""
+class Cipher:
+    def __init__(self, key: bytes) -> None:
+        self._chacha = ChaCha20Poly1305(key)
+
+    @classmethod
+    def generate_cipher(cls, salt: bytes, password: bytes) -> Self:
+        """Generates cipher which uses key derived from a given password"""
         kdf = PBKDF2HMAC(
             algorithm=hashes.SHA256(),
             length=32,
@@ -19,8 +24,22 @@ def _generate_key(salt: bytes, master_pass: bytes) -> bytes:
             iterations=100000,
             backend=default_backend(),
         )
-    key = base64.urlsafe_b64encode(kdf.derive(master_pass))
-    return key
+        return cls(kdf.derive(password))
+
+    def encrypt(self, data: bytes) -> bytes:
+        nonce = os.urandom(12)
+        return nonce + self._chacha.encrypt(
+            nonce,
+            data,
+            associated_data=None,
+        )
+
+    def decrypt(self, data: bytes) -> bytes:
+        return self._chacha.decrypt(
+            nonce=data[:12],
+            data=data[12:],
+            associated_data=None,
+        )
 
 
 def encrypt(
@@ -29,15 +48,10 @@ def encrypt(
 ) -> Account:
     """Encrypts account using master password and returns Account object"""
     salt = os.urandom(64)
-    key = _generate_key(salt, master_pass.encode("utf-8"))
-    f = Fernet(key)
+    cipher = Cipher.generate_cipher(salt, master_pass.encode("utf-8"))
 
-    enc_login = base64.urlsafe_b64decode(
-        f.encrypt(account.login.encode("utf-8")),
-    )
-    enc_password = base64.urlsafe_b64decode(
-        f.encrypt(account.password.encode("utf-8")),
-    )
+    enc_login = cipher.encrypt(account.login.encode("utf-8"))
+    enc_password = cipher.encrypt(account.password.encode("utf-8"))
 
     return Account(
         user_id=account.user_id,
@@ -54,15 +68,10 @@ def decrypt(
 ) -> DecryptedAccount:
     """Decrypts account using master password and returns
     DecryptedAccount object"""
-    key = _generate_key(account.salt, master_pass.encode("utf-8"))
-    f = Fernet(key)
+    cipher = Cipher.generate_cipher(account.salt, master_pass.encode("utf-8"))
 
-    login = f.decrypt(
-        base64.urlsafe_b64encode(account.enc_login),
-    ).decode("utf-8")
-    password = f.decrypt(
-        base64.urlsafe_b64encode(account.enc_password),
-    ).decode("utf-8")
+    login = cipher.decrypt(account.enc_login).decode("utf-8")
+    password = cipher.decrypt(account.enc_password).decode("utf-8")
 
     return DecryptedAccount(
         user_id=account.user_id,