encrypt_decrypt.py

import os
import sys
from cryptography.fernet import Fernet, InvalidToken


KEY_FILE = "secret.key"


def generate_key():
    if os.path.exists(KEY_FILE):
        return
    key = Fernet.generate_key()
    with open(KEY_FILE, "wb") as f:
        f.write(key)
    print(f"New key generated and saved to {KEY_FILE}")


def load_key():
    if not os.path.exists(KEY_FILE):
        print("No key found. Generating a new one...")
        generate_key()
    with open(KEY_FILE, "rb") as f:
        return f.read()


def encrypt_file(input_path):
    if not os.path.exists(input_path):
        print(f"File not found: {input_path}")
        return

    key = load_key()
    fernet = Fernet(key)

    with open(input_path, "rb") as f:
        data = f.read()

    encrypted_data = fernet.encrypt(data)

    output_path = input_path + ".enc"
    with open(output_path, "wb") as f:
        f.write(encrypted_data)

    print(f"File encrypted successfully -> {output_path}")


def decrypt_file(input_path):
    if not os.path.exists(input_path):
        print(f"File not found: {input_path}")
        return

    key = load_key()
    fernet = Fernet(key)

    with open(input_path, "rb") as f:
        encrypted_data = f.read()

    try:
        decrypted_data = fernet.decrypt(encrypted_data)
    except InvalidToken:
        print("Decryption failed. Invalid key or corrupted file.")
        return

    if input_path.endswith(".enc"):
        output_path = input_path[:-4] + ".dec"
    else:
        output_path = input_path + ".dec"

    with open(output_path, "wb") as f:
        f.write(decrypted_data)

    print(f"File decrypted successfully -> {output_path}")


def show_menu():
    print("\n===== File Encryption / Decryption Tool =====")
    print("1. Encrypt a file")
    print("2. Decrypt a file")
    print("3. Exit")


def main():
    generate_key()

    while True:
        show_menu()
        choice = input("Enter your choice (1-3): ").strip()

        if choice == "1":
            path = input("Enter the path of the file to encrypt: ").strip()
            encrypt_file(path)

        elif choice == "2":
            path = input("Enter the path of the file to decrypt: ").strip()
            decrypt_file(path)

        elif choice == "3":
            print("Exiting the program. Goodbye!")
            sys.exit(0)

        else:
            print("Invalid choice. Please select between 1 and 3.")


if __name__ == "__main__":
    main()