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,