8
\$\begingroup\$

I'm looking for any feedback in order to improve code styling, performance etc. Also I want to make sure that the way I've implemented AES here is actually secure

import os
import hashlib
import sys
import base64

from cryptography.fernet import Fernet

class Cryptor:
    def __init__(self, password):
        self.__password = password

    def operate(self, file_):
        if file_.endswith("enc"):
            self.decrypt(file_)
        else:
            self.encrypt(file_)

    def decrypt(self, file_):
        file_name = "".join(file_.split(".")[:-1])

        with open(file_, "rb") as file_:
            save1 = file_.readline()
            save2 = file_.readline()
            save3 = file_.readline()
            salt = save3
            key = base64.urlsafe_b64encode(hashlib.pbkdf2_hmac("sha256", bytes(self.__password, encoding="utf-8"), salt, 100000, dklen=32))
            f = Fernet(key)
            decrypted = f.decrypt(save1)
            file_extention = f.decrypt(save2)
            save = decrypted

        with open(file_name + file_extention.decode(encoding="utf-8"), "wb") as file_:
            file_.write(save)

        os.remove(file_name + ".enc")

    def encrypt(self, file_):
        salt = os.urandom(64)
        key = base64.urlsafe_b64encode(hashlib.pbkdf2_hmac("sha256", bytes(self.__password, encoding="utf-8"), salt, 100000, dklen=32))
        f = Fernet(key)

        file_name = "".join(file_.split(".")[:-1])
        file_extention = "." + file_.split(".")[-1]

        with open(file_, "rb") as file_:
            file__ = file_.read()
            encrypted = f.encrypt(bytes(file__))
            save1 = encrypted
            save2 = f.encrypt(bytes(file_extention, encoding="utf-8"))
            save3 = salt

        with open(file_name + ".enc", "wb") as file_:
            file_.write(save1)
            file_.write(b"\n")
            file_.write(save2)
            file_.write(b"\n")
            file_.write(save3)

        os.remove(file_name + file_extention)

def main():
    password = sys.argv[1]
    files = sys.argv[2:]
    crypt = Cryptor(password)

    for file_ in files:
        crypt.operate(file_)

if __name__ == "__main__":
    main()
\$\endgroup\$
6
\$\begingroup\$

Ways of improving:

Splitting filename

Both methods encrypt and decrypt try to split a filename into "root" part and extension.
encrypt method makes that even worse and verbose with 2 statements:

    file_name = "".join(file_.split(".")[:-1])
    file_extention = "." + file_.split(".")[-1]

Instead, that's easily achievable with os.path.splittext(path) function (split the pathname path into a pair (root, ext)):

file_name, file_ext = os.path.splitext(file_)

Generating Fernet instance

Both methods encrypt and decrypt have the same 2 statements for generating Fernet instance (with only difference in salt value):

key = base64.urlsafe_b64encode(hashlib.pbkdf2_hmac("sha256", bytes(self.__password, encoding="utf-8"), salt, 100000, dklen=32))
f = Fernet(key)

It's good to extract that common behavior into a separate internal function:

def _get_fernet_instance(self, salt):
    key = base64.urlsafe_b64encode(hashlib.pbkdf2_hmac("sha256", bytes(self.__password, encoding="utf-8"), 
                                                       salt, 100000, dklen=32))
    return Fernet(key)

  • decrypt function.
    The variables save1, save2, save3 and save reassignment are confusing things rather than clarify.
    The variables save1, save2, save3 can be renamed to line1, line2 and salt.
    Aliasing save = decrypted gives no benefit - just refer decrypted directly.
    Including the above mentioned optimizations the decrypt method would look as:

    def decrypt(self, file_):
        file_name, _ = os.path.splitext(file_)
    
        with open(file_, "rb") as f:
            line1 = f.readline()
            line2 = f.readline()
            salt = f.readline()
    
            fernet_ = self._get_fernet_instance(salt)
            decrypted = fernet_.decrypt(line1)
            file_ext = fernet_.decrypt(line2)
    
        with open(file_name + file_ext.decode(encoding="utf-8"), "wb") as f:
            f.write(decrypted)
    
        os.remove(file_)
    
  • encrypt function.
    The variable file__ in file__ = file_.read() does not give an explicit meaning for the content being read. Let the content be content or data.
    The same issues around variables save1, save2, save3 and redundant aliases.

    In case if there would be a need to refer the input argument file_ (in both methods) as its original value in different places in methods body - you should not reassign it in context managers with open(file_, "rb") as file_:.

    os.remove(file_name + file_extention) is essentially the same as os.remove(file_)

    The final optimized encrypt method would look as:

    def encrypt(self, file_):
        salt = os.urandom(64)
        fernet_ = self._get_fernet_instance(salt)
        file_name, file_ext = os.path.splitext(file_)
    
        with open(file_, "rb") as f:
            content = f.read()
            encrypted_content = fernet_.encrypt(bytes(content))
            encrypted_ext = fernet_.encrypt(bytes(file_ext, encoding="utf-8"))
    
        with open(f'{file_name}.enc', 'wb') as f:
            f.write(encrypted_content)
            f.write(b"\n")
            f.write(encrypted_ext)
            f.write(b"\n")
            f.write(salt)
    
        os.remove(file_)
    
\$\endgroup\$
  • \$\begingroup\$ Cheers for taking the time and giving a thorough answer! \$\endgroup\$ – Krishi Nov 16 at 23:17
  • \$\begingroup\$ @Krishi, welcome, see my last edit for adjusting variable name for Fernet instance \$\endgroup\$ – RomanPerekhrest Nov 17 at 6:30

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.