From c051c14f1f4da7b6ef676d70181051cef2fadb51 Mon Sep 17 00:00:00 2001
From: StNicolay <stnicolay2007@mail.ru>
Date: Tue, 10 Jan 2023 20:00:58 +0300
Subject: [PATCH] Switched to the ChaCha20Poly1305 encryption algorithm for
 better security

---
 src/encryption/accounts.py | 67 +++++++++++++++++++++-----------------
 1 file changed, 38 insertions(+), 29 deletions(-)

diff --git a/src/encryption/accounts.py b/src/encryption/accounts.py
index 6e841b6..35cbdb9 100644
--- a/src/encryption/accounts.py
+++ b/src/encryption/accounts.py
@@ -1,26 +1,45 @@
-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"""
-    kdf = PBKDF2HMAC(
-        algorithm=hashes.SHA256(),
-        length=32,
-        salt=salt,
-        iterations=100000,
-        backend=default_backend(),
-    )
-    key = base64.urlsafe_b64encode(kdf.derive(master_pass))
-    return key
+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,
+            salt=salt,
+            iterations=100000,
+            backend=default_backend(),
+        )
+        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,