A Python utility that allows users to securely encrypt or decrypt files using AES with a choice between 128‑bit and 256‑bit keys. It lists files from the current directory, prompts for password confirmation during encryption, and handles decryption seamlessly by deriving keys with PBKDF2 and using CBC mode.
This script relies on the PyCryptodome library to handle AES encryption and decryption operations. Be sure to install it via pip install pycryptodome before running the script.
Click to view script…
#!/usr/bin/env python3
import os
import sys
import getpass
from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Util.Padding import pad, unpad
HEADER_MAGIC = b"ENCR" # 4 bytes header to identify our encrypted files
SALT_SIZE = 16 # Salt size in bytes
IV_SIZE = 16 # AES block size
PBKDF2_ITERATIONS = 100000
def list_files():
files = [f for f in os.listdir('.') if os.path.isfile(f)]
if not files:
print("No files found in the current directory.")
sys.exit(1)
print("Files in the current directory:")
for i, f in enumerate(files, start=1):
print(f" {i}. {f}")
return files
def choose_file(files):
try:
choice = int(input("Select a file number: "))
if choice < 1 or choice > len(files):
raise ValueError
except ValueError:
print("Invalid selection.")
sys.exit(1)
return files[choice - 1]
def get_password(confirm=False):
pwd = getpass.getpass("Enter password: ")
if confirm:
pwd_confirm = getpass.getpass("Confirm password: ")
if pwd != pwd_confirm:
print("Passwords do not match.")
sys.exit(1)
return pwd.encode() # work with bytes
def encrypt_file(filepath):
print("\n--- Encryption ---")
# Choose key size
key_size_choice = input("Choose key size (enter 128 or 256): ").strip()
if key_size_choice not in ("128", "256"):
print("Invalid key size selection.")
sys.exit(1)
key_len = 16 if key_size_choice == "128" else 32
password = get_password(confirm=True)
# Generate a random salt and IV
salt = get_random_bytes(SALT_SIZE)
iv = get_random_bytes(IV_SIZE)
# Derive the AES key from the password and salt
key = PBKDF2(password, salt, dkLen=key_len, count=PBKDF2_ITERATIONS)
# Read file data
try:
with open(filepath, "rb") as f:
data = f.read()
except IOError as e:
print(f"Error reading file: {e}")
sys.exit(1)
# Pad and encrypt the data using CBC mode
cipher = AES.new(key, AES.MODE_CBC, iv)
ciphertext = cipher.encrypt(pad(data, AES.block_size))
# Create an output file format:
# [4 bytes header][1 byte key length indicator][16 bytes salt][16 bytes IV][ciphertext]
# We store key length as one byte: 16 (for 128-bit) or 32 (for 256-bit)
header = HEADER_MAGIC + bytes([key_len]) + salt + iv
out_data = header + ciphertext
# Determine output file name
out_filepath = filepath + ".enc"
try:
with open(out_filepath, "wb") as f:
f.write(out_data)
except IOError as e:
print(f"Error writing encrypted file: {e}")
sys.exit(1)
print(f"Encryption successful. Encrypted file: {out_filepath}")
def decrypt_file(filepath):
print("\n--- Decryption ---")
# Read the encrypted file and parse header information
try:
with open(filepath, "rb") as f:
file_data = f.read()
except IOError as e:
print(f"Error reading file: {e}")
sys.exit(1)
# Verify header length
header_length = len(HEADER_MAGIC) + 1 + SALT_SIZE + IV_SIZE
if len(file_data) < header_length:
print("File is too short to be a valid encrypted file.")
sys.exit(1)
header = file_data[:len(HEADER_MAGIC)]
if header != HEADER_MAGIC:
print("Invalid file format (missing header).")
sys.exit(1)
key_len = file_data[len(HEADER_MAGIC)]
if key_len not in (16, 32):
print("Invalid key length in file header.")
sys.exit(1)
salt = file_data[len(HEADER_MAGIC)+1 : len(HEADER_MAGIC)+1+SALT_SIZE]
iv_start = len(HEADER_MAGIC)+1+SALT_SIZE
iv = file_data[iv_start : iv_start+IV_SIZE]
ciphertext = file_data[iv_start+IV_SIZE:]
password = get_password(confirm=False)
key = PBKDF2(password, salt, dkLen=key_len, count=PBKDF2_ITERATIONS)
# Decrypt and unpad the data
cipher = AES.new(key, AES.MODE_CBC, iv)
try:
plaintext_padded = cipher.decrypt(ciphertext)
plaintext = unpad(plaintext_padded, AES.block_size)
except (ValueError, KeyError):
print("Decryption failed. Incorrect password or corrupted file.")
sys.exit(1)
# Determine output file name
if filepath.endswith(".enc"):
out_filepath = filepath[:-4] # remove the .enc extension
else:
out_filepath = filepath + ".dec"
try:
with open(out_filepath, "wb") as f:
f.write(plaintext)
except IOError as e:
print(f"Error writing decrypted file: {e}")
sys.exit(1)
print(f"Decryption successful. Decrypted file: {out_filepath}")
def main():
print("AES Encryption/Decryption Utility")
mode = input("Do you want to (E)ncrypt or (D)ecrypt? ").strip().lower()
if mode not in ("e", "d"):
print("Invalid mode selection.")
sys.exit(1)
files = list_files()
selected_file = choose_file(files)
if mode == "e":
encrypt_file(selected_file)
else:
decrypt_file(selected_file)
if __name__ == "__main__":
main()