# =============================================================================
# pyencryptdecrypt.py - AES File Encryption/Decryption Tool
# =============================================================================
# REQUIREMENTS:
#   pip install pycryptodome
#
# USAGE:
#   python pyencryptdecrypt.py <file_path> <password> <encrypt/decrypt>
#
# ENCRYPT EXAMPLE:
#   python pyencryptdecrypt.py C:\Users\you\Documents\TEST.XLSX MYPASSWORD123 encrypt
#   Result: TEST.XLSX.aes  (original TEST.XLSX is deleted after encryption)
#
# DECRYPT EXAMPLE:
#   python pyencryptdecrypt.py C:\Users\you\Documents\TEST.XLSX.aes MYPASSWORD123 decrypt
#   Result: TEST.XLSX  (the .aes file is deleted after decryption)
#
# PASSWORD RULES:
#   - Minimum 12 characters
#   - Must contain at least one number
#   - No special character required
#
# NOTES:
#   - Encryption uses AES-EAX mode (authenticated encryption)
#   - If decryption fails due to wrong password, output file is deleted
#     and "Incorrect password or file corrupted." is printed
#   - File is processed in 64KB chunks so large files are handled efficiently
#   - The original file is permanently deleted after successful encryption
#   - The .aes file is permanently deleted after successful decryption
# =============================================================================
import os
import sys
from Crypto.Cipher import AES
from Crypto.Protocol.KDF import scrypt
from Crypto.Random import get_random_bytes
CHUNK_SIZE = 64 * 1024  # 64KB
def derive_key(password, salt):
    return scrypt(password.encode(), salt, 32, N=16384, r=8, p=1)
def validate_password(password):
    if len(password) < 12:
        return "Password must be at least 12 characters long and include at least one number."
    if not any(char.isdigit() for char in password):
        return "Password must also include at least one number."
    return None
def encrypt_file(file_path, password):
    salt = get_random_bytes(16)
    key = derive_key(password, salt)
    cipher = AES.new(key, AES.MODE_EAX)
    output_path = file_path + ".aes"
    with open(file_path, 'rb') as f_in, open(output_path, 'wb') as f_out:
        # Write header
        f_out.write(salt)
        f_out.write(cipher.nonce)
        # Encrypt in chunks
        while True:
            chunk = f_in.read(CHUNK_SIZE)
            if not chunk:
                break
            f_out.write(cipher.encrypt(chunk))
        # Write authentication tag
        f_out.write(cipher.digest())
    print("Encryption complete:", output_path)
def decrypt_file(file_path, password):
    with open(file_path, 'rb') as f_in:
        salt = f_in.read(16)
        nonce = f_in.read(16)
        # Get file size to separate ciphertext and tag
        f_in.seek(0, os.SEEK_END)
        file_size = f_in.tell()
        tag_position = file_size - 16
        ciphertext_length = tag_position - 32  # subtract salt + nonce
        f_in.seek(32)  # move to start of ciphertext
        key = derive_key(password, salt)
        cipher = AES.new(key, AES.MODE_EAX, nonce=nonce)
        output_path = file_path.rsplit('.aes', 1)[0]
        if os.path.exists(output_path):
            os.remove(output_path)
        with open(output_path, 'wb') as f_out:
            bytes_remaining = ciphertext_length
            while bytes_remaining > 0:
                chunk_size = min(CHUNK_SIZE, bytes_remaining)
                chunk = f_in.read(chunk_size)
                f_out.write(cipher.decrypt(chunk))
                bytes_remaining -= chunk_size
            tag = f_in.read(16)
            try:
                cipher.verify(tag)
            except ValueError:
                f_out.close()
                os.remove(output_path)
                print("Incorrect password or file corrupted.")
                sys.exit(1)
    print("Decryption complete:", output_path)
if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("Usage: python pyencryptdecrypt.py <file_path> <password> <encrypt/decrypt>")
        sys.exit(1)
    file_path = sys.argv[1]
    password = sys.argv[2]
    action = sys.argv[3].lower()
    if action == "encrypt":
        validation_error = validate_password(password)
        if validation_error:
            print(validation_error)
            sys.exit(1)
    if not os.path.exists(file_path):
        print("File not found:", file_path)
        sys.exit(1)
    if action == 'encrypt':
        encrypt_file(file_path, password)
    elif action == 'decrypt':
        decrypt_file(file_path, password)
    else:
        print("Invalid action. Use 'encrypt' or 'decrypt'.")
        sys.exit(1)