(AES) Encrypting Folder in Python

aesEncryption.py

from Crypto.Cipher import AES
from Crypto import Random
import os
import time
import sys
import argparse
from pathlib import Path
import colorama
from termcolor import colored
import shutil
import traceback

colorama.init()

def printMsg(msg, msgType = 'info'):
	if msgType == 'error':
		print(colored(msg, 'red'))
	elif msgType == 'success':
		print(colored(msg, 'green'))
	else:
		print(msg)

def getStrInput(prompt):
	try:
		strInput = str(input(prompt))
	except KeyboardInterrupt:
		printMsg('\n')
		printMsg('Don\'t want to continue? :(')
		sys.exit()

	return strInput

def aesObj(key, iv):
	return AES.new(key, AES.MODE_CFB, iv = iv)

def fileNameWithoutExtensions(fileObj):
	name = ''
	suffixes = ''.join(fileObj.suffixes)

	if fileObj.name.find(suffixes) > -1:
		name = fileObj.name[:fileObj.name.find(suffixes)]

	return name

def readFileContent(fileLocation, errText, fOpenMode = 'rb'):
	content = ''

	try:
		fileObj = open(fileLocation, fOpenMode)

		content = fileObj.read()
	except Exception as e:
		printMsg(errText, 'error')
	finally:
		fileObj.close()

	return content

def writeToFile(fileLocation, content, fOpenMode = 'wb'):
	fileObj = open(fileLocation, fOpenMode)

	fileObj.write(content)

	fileObj.close()

def proceedEncryption():
	folderLocation = getStrInput('\nEnter Folder Location To Lock: ')

	if len(folderLocation) > 0 and Path(folderLocation).is_dir():
		folderPointer = Path(folderLocation)
		parnetFolderLocation = folderPointer.parent

		# Create folder in [Folder Name]_locked format
		lockedFolderPath = parnetFolderLocation.joinpath(folderPointer.name + '_locked')

		if not(Path(lockedFolderPath).exists()):
			Path(lockedFolderPath).mkdir()

		namePrefix = str(lockedFolderPath.joinpath(folderPointer.name))

		# Archive the folder
		archiveLocation = shutil.make_archive(namePrefix, 'zip', folderLocation)

		# Uncomment Below to Delete the Source Folder
		# shutil.rmtree(folderLocation)

		try:
			archiveFile = open(archiveLocation, 'rb')
			archiveFileByteContent = archiveFile.read()
			archiveFile.close()

			# Delete the Archive file
			os.remove(archiveLocation)

			# Generate Key and IV(Initialization Vector)
			key = Random.get_random_bytes(32)
			iv = Random.get_random_bytes(16)

			encObj = aesObj(key, iv)

			# Encrypt the Folder Content
			encryptedText = encObj.encrypt(archiveFileByteContent)

			# Store Encrypted folder as Binary file
			writeToFile(namePrefix + '.encrypted.bin', encryptedText)

			# Store Key and IV as Binary files
			writeToFile(namePrefix + '.key.bin', key)
			writeToFile(namePrefix + '.iv.bin', iv)

			printMsg('\nThe Folder has been Locked at:\n' + str(lockedFolderPath) + '\nPlease keep the Key and IV files safe as these files are required to Unlock the Folder.', 'success')
		except FileNotFoundError as e:
			printMsg('Unable to access the Archived File', 'error')
			sys.exit()
	else:
		printMsg('Invalid Folder or Folder not Found!', 'error')
		sys.exit()

def proceedDecryption():
	encryptedFileLocation = getStrInput('\nEnter Locked File Location To Unlock: ')

	# Check Encrypted File exists and in correct file Format
	if not(len(encryptedFileLocation) > 0 and Path(encryptedFileLocation).is_file() and '.encrypted.bin' == ''.join(Path(encryptedFileLocation).suffixes)):
		printMsg('Invalid File or File not Found!', 'error')
		sys.exit()

	keyFileLocation = getStrInput('Enter Key file Location: ')

	# Check Key File exists and in correct file Format
	if not(len(keyFileLocation) > 0 and Path(keyFileLocation).is_file() and '.key.bin' == ''.join(Path(keyFileLocation).suffixes)):
		printMsg('Invalid File or File not Found!', 'error')
		sys.exit()

	ivFileLocation = getStrInput('Enter IV file Location: ')

	# Check IV File exists and in correct file Format
	if not(len(ivFileLocation) > 0 and Path(ivFileLocation).is_file() and '.iv.bin' == ''.join(Path(ivFileLocation).suffixes)):
		printMsg('Invalid File or File not Found!', 'error')
		sys.exit()

	# Get Encrypted File Content
	encryptedFileContent = readFileContent(encryptedFileLocation, 'Unable to Read Locked File Content!')

	# Get Key and IV Content
	keyFileContent = readFileContent(keyFileLocation, 'Unable to Read Locked File Content!')
	ivFileContent = readFileContent(ivFileLocation, 'Unable to Read Locked File Content!')

	try:
		encryptedFileObj = Path(encryptedFileLocation)
		namePrefix = fileNameWithoutExtensions(encryptedFileObj)

		if len(namePrefix) == 0:
			raise Exception('Unable to Get the Name prefix!')

		decObj = aesObj(keyFileContent, ivFileContent)

		# Decrypt the Encrypted Content
		originalFileContent = decObj.decrypt(encryptedFileContent)

		# Create folder in [Folder Name]_unlocked format
		unlockedFolderPath = encryptedFileObj.parent.joinpath(namePrefix + '_unlocked')

		if not(Path(unlockedFolderPath).exists()):
			Path(unlockedFolderPath).mkdir()

		namePrefix = str(unlockedFolderPath.joinpath(namePrefix));
		originalFileLocation = namePrefix + '.zip';

		# Restore the Archive file
		archiveFile = open(originalFileLocation, 'wb')
		archiveFile.write(originalFileContent)
		archiveFile.close()

		# Extract the Archive file as Original Folder
		shutil.unpack_archive(originalFileLocation, namePrefix)

		# Remove Archive file
		os.remove(originalFileLocation)

		# Remove Encryption files
		os.remove(encryptedFileLocation)
		os.remove(keyFileLocation)
		os.remove(ivFileLocation)

		printMsg('\nThe Folder has been Unlocked at:\n' + str(unlockedFolderPath), 'success')
	except Exception as e:
		printMsg('Error: ', 'error')
		traceback.print_exec()
		sys.exit()

action = getStrInput('Enter any Key to Continue (l/L - Lock, u/U - Unlock): ').lower()

if action == 'l':
	proceedEncryption()
elif action == 'u':
	proceedDecryption()
else:
	printMsg('Invalid Key!', 'error')