3
\$\begingroup\$

I've recently started looking into encryption in python and I came across with pycryptodome library. I made this simple program that encrypt/decrypt a file using AES-GCM. Could someone take a look if I implemented the encryption and authentication correctly? For sake of simplicity, I'm posting only the encryption and makeKey functions.

import os
import struct
import secrets
import hahslib
from Cryptodome.Cipher import AES


def makeKey(password, salt):    
    key = hashlib.pbkdf2_hmac('sha512',password.encode('ascii'), salt, 1000000)
    return key[:32]

def encrypt(password, inputFileName):
    #creates the needed variables for encryption
    salt = secrets.token_bytes(16)
    key = makeKey(password, salt)
    nonce = secrets.token_bytes(AES.block_size)
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    #put the file size, salt and nonce in outputfile
    fileSize = os.path.getsize(inputFileName)
    with open("output", "wb") as outputFile:
        outputFile.write(struct.pack('<Q', fileSize))
        outputFile.write(salt)
        outputFile.write(nonce)
        #beginning of encryption itself
        chunkSize = 64 * 1024
        with open(inputFileName, "rb") as inputFile:
            while True:
                chunk = inputFile.read(chunkSize)
                if len(chunk) == 0:
                    outputFile.write(cipher.digest())
                    break
                outputFile.write(cipher.encrypt(chunk))
\$\endgroup\$
2
\$\begingroup\$

Handling files - the Python way

Without being able to test run your code at the momemt, there is one thing I can already recommend to you, which is to make use of Python's powerful with statement. Using with will make sure that files will be closed no matter what happens (e.g. unexpected exceptions). with also works with more than one file as described in this SO post.

Also consider adding the output filename as parameter to your function. You can even give your momentary value of "output" as default value.


Minor notes on style

I'm not sure if you're aware of the official Style Guide for Python, often just called PEP8. The official recommendation is to use snake_case for variable and function names. This would create a more unified code appearance when working with other Python libraries which often stick to that convention. As always, exceptions apply.


Below is your code including these changes.

def make_key(password, salt):
    ...

def encrypt(password, input_filename, output_filename="output"):
    #creates the needed variables for encryption
    salt = secrets.token_bytes(16)
    key = make_key(password, salt)
    nonce = secrets.token_bytes(AES.block_size)
    cipher = AES.new(key, AES.MODE_GCM, nonce=nonce)
    # put the file size, salt and nonce in outputfile
    filesize = os.path.getsize(input_filename)
    with open(output_filename, "wb") as output_file,\
         open(input_filename, "rb") as input_file:
        output_file.write(struct.pack('<Q', filesize))
        output_file.write(salt)
        output_file.write(nonce)
        #beginning of encryption itself
        chunkSize = 64 * 1024
        while True:
            chunk = input_file.read(chunkSize)
            if len(chunk) == 0:
                output_file.write(cipher.digest())
                break
            output_file.write(cipher.encrypt(chunk))
        print("File encrypted successfully!")
\$\endgroup\$

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.