AES File Encryption & Decryption Utility Python 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()