File Encryption / Decryption Tool

 import os

import json

import hashlib

import base64

import secrets

import getpass

from datetime import datetime

from pathlib import Path


try:

    from cryptography.fernet import Fernet, InvalidToken

    from cryptography.hazmat.primitives import hashes

    from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC

    CRYPTO_OK = True

except ImportError:

    CRYPTO_OK = False


# ============================================================

# CONFIGURATION

# ============================================================


KEY_FILE      = "encryption.key"

VAULT_LOG     = "encryption_vault.json"

ENC_EXTENSION = ".encrypted"

SALT_SIZE     = 16       # bytes

ITERATIONS    = 480000   # PBKDF2 iterations (NIST recommended)


# ============================================================

# HELPERS

# ============================================================


def format_size(size_bytes):

    for unit in ["B", "KB", "MB", "GB"]:

        if size_bytes < 1024:

            return f"{size_bytes:.1f} {unit}"

        size_bytes /= 1024

    return f"{size_bytes:.1f} GB"



def file_checksum(filepath):

    """Return MD5 checksum of a file."""

    h = hashlib.md5()

    with open(filepath, "rb") as f:

        while chunk := f.read(8192):

            h.update(chunk)

    return h.hexdigest()



# ============================================================

# KEY MANAGEMENT

# ============================================================


def generate_key():

    """Generate a new random Fernet key."""

    return Fernet.generate_key()



def derive_key_from_password(password: str, salt: bytes = None):

    """

    Derive a Fernet-compatible key from a password using PBKDF2-HMAC-SHA256.

    Returns (key, salt).

    """

    if salt is None:

        salt = secrets.token_bytes(SALT_SIZE)


    kdf = PBKDF2HMAC(

        algorithm=hashes.SHA256(),

        length=32,

        salt=salt,

        iterations=ITERATIONS,

    )

    key = base64.urlsafe_b64encode(kdf.derive(password.encode()))

    return key, salt



def save_key(key: bytes, filepath=KEY_FILE):

    """Save a key to a file."""

    with open(filepath, "wb") as f:

        f.write(key)

    print(f"  Key saved: {filepath}")

    print("  IMPORTANT: Keep this key file safe. Without it, decryption is impossible!")



def load_key(filepath=KEY_FILE):

    """Load a key from file."""

    if not Path(filepath).exists():

        print(f"  Key file not found: {filepath}")

        return None

    with open(filepath, "rb") as f:

        return f.read().strip()



# ============================================================

# VAULT LOG

# ============================================================


def load_vault():

    if Path(VAULT_LOG).exists():

        try:

            with open(VAULT_LOG, "r") as f:

                return json.load(f)

        except:

            pass

    return {"files": []}



def save_vault(vault):

    with open(VAULT_LOG, "w") as f:

        json.dump(vault, f, indent=2)



def log_operation(vault, operation, filepath, output_path, method):

    vault["files"].append({

        "operation":   operation,

        "source":      str(filepath),

        "output":      str(output_path),

        "method":      method,

        "timestamp":   datetime.now().strftime("%d-%m-%Y %H:%M:%S"),

        "size":        format_size(Path(output_path).stat().st_size)

                       if Path(output_path).exists() else "?"

    })

    vault["files"] = vault["files"][-100:]

    save_vault(vault)



# ============================================================

# ENCRYPT FILE

# ============================================================


def encrypt_file(filepath, key: bytes, output_path=None,

                 salt: bytes = None):

    """

    Encrypt a file using Fernet.

    If salt is provided, it's prepended to the encrypted output

    (used for password-based encryption).

    Returns output path or None on failure.

    """

    filepath = Path(filepath)

    if not filepath.exists():

        print(f"  File not found: {filepath}")

        return None


    if output_path is None:

        output_path = filepath.with_suffix(filepath.suffix + ENC_EXTENSION)

    output_path = Path(output_path)


    try:

        fernet = Fernet(key)


        with open(filepath, "rb") as f:

            data = f.read()


        encrypted = fernet.encrypt(data)


        with open(output_path, "wb") as f:

            if salt:

                # Write salt length (2 bytes) + salt + encrypted data

                f.write(len(salt).to_bytes(2, "big"))

                f.write(salt)

            f.write(encrypted)


        orig_size = format_size(filepath.stat().st_size)

        enc_size  = format_size(output_path.stat().st_size)

        checksum  = file_checksum(filepath)


        print(f"\n  Encrypted successfully!")

        print(f"  Input  : {filepath.name}  ({orig_size})")

        print(f"  Output : {output_path.name}  ({enc_size})")

        print(f"  Checksum (original): {checksum}")

        return output_path


    except Exception as e:

        print(f"  Encryption failed: {e}")

        return None



# ============================================================

# DECRYPT FILE

# ============================================================


def decrypt_file(filepath, key: bytes, output_path=None,

                 password_based=False):

    """

    Decrypt a Fernet-encrypted file.

    If password_based=True, reads the salt from the file header.

    Returns output path or None on failure.

    """

    filepath = Path(filepath)

    if not filepath.exists():

        print(f"  File not found: {filepath}")

        return None


    # Determine output path

    if output_path is None:

        name = filepath.name

        if name.endswith(ENC_EXTENSION):

            out_name = name[:-len(ENC_EXTENSION)]

        else:

            out_name = name + ".decrypted"

        output_path = filepath.parent / out_name

    output_path = Path(output_path)


    try:

        with open(filepath, "rb") as f:

            raw = f.read()


        if password_based:

            # Read salt from header

            salt_len  = int.from_bytes(raw[:2], "big")

            salt      = raw[2:2 + salt_len]

            encrypted = raw[2 + salt_len:]


            # Re-derive key from password

            password = getpass.getpass("  Password: ")

            key, _   = derive_key_from_password(password, salt)

        else:

            encrypted = raw


        fernet    = Fernet(key)

        decrypted = fernet.decrypt(encrypted)


        with open(output_path, "wb") as f:

            f.write(decrypted)


        orig_size = format_size(filepath.stat().st_size)

        dec_size  = format_size(output_path.stat().st_size)

        checksum  = file_checksum(output_path)


        print(f"\n  Decrypted successfully!")

        print(f"  Input  : {filepath.name}  ({orig_size})")

        print(f"  Output : {output_path.name}  ({dec_size})")

        print(f"  Checksum (restored): {checksum}")

        return output_path


    except InvalidToken:

        print("  Decryption failed: Wrong key or password, or file is corrupted.")

        return None

    except Exception as e:

        print(f"  Decryption failed: {e}")

        return None



# ============================================================

# ENCRYPT / DECRYPT FOLDER

# ============================================================


def process_folder(folder_path, key, operation,

                   extensions=None, password_based=False, salt=None):

    folder = Path(folder_path)

    if not folder.is_dir():

        print("  Invalid folder.")

        return


    if operation == "encrypt":

        if extensions:

            files = [f for f in folder.rglob("*")

                     if f.is_file() and f.suffix.lower() in extensions

                     and not f.name.endswith(ENC_EXTENSION)]

        else:

            files = [f for f in folder.rglob("*")

                     if f.is_file() and not f.name.endswith(ENC_EXTENSION)]

    else:

        files = [f for f in folder.rglob("*")

                 if f.is_file() and f.name.endswith(ENC_EXTENSION)]


    if not files:

        print(f"  No files found to {operation}.")

        return


    print(f"\n  Found {len(files)} file(s) to {operation}.")

    confirm = input(f"  Proceed? (y/n): ").strip().lower()

    if confirm != "y":

        print("  Cancelled.")

        return


    success = 0

    failed  = 0

    vault   = load_vault()


    for f in files:

        print(f"\n  [{operation.upper()}] {f.name}")

        if operation == "encrypt":

            result = encrypt_file(f, key, salt=salt)

        else:

            result = decrypt_file(f, key, password_based=password_based)


        if result:

            success += 1

            log_operation(vault, operation, f, result,

                          "password" if password_based else "key-file")

        else:

            failed += 1


    save_vault(vault)

    print(f"\n  Done: {success} succeeded, {failed} failed.")



# ============================================================

# ENCRYPT TEXT STRING

# ============================================================


def encrypt_text(key: bytes):

    print("\n  Enter text to encrypt (type END on new line to finish):")

    lines = []

    while True:

        line = input()

        if line.strip().upper() == "END":

            break

        lines.append(line)

    text = "\n".join(lines)


    fernet    = Fernet(key)

    encrypted = fernet.encrypt(text.encode())

    encoded   = encrypted.decode()


    print(f"\n  Encrypted text (copy this):")

    print(f"  {encoded}")


    save = input("\n  Save to file? (y/n): ").strip().lower()

    if save == "y":

        fname = f"encrypted_text_{datetime.now().strftime('%Y%m%d_%H%M%S')}.txt"

        with open(fname, "w") as f:

            f.write(encoded)

        print(f"  Saved: {fname}")



def decrypt_text(key: bytes):

    print("\n  Paste encrypted text:")

    enc_text = input("  > ").strip()


    try:

        fernet    = Fernet(key)

        decrypted = fernet.decrypt(enc_text.encode()).decode()

        print(f"\n  Decrypted text:")

        print(f"  {decrypted}")

    except InvalidToken:

        print("  Decryption failed: Invalid token or wrong key.")

    except Exception as e:

        print(f"  Error: {e}")



# ============================================================

# DISPLAY VAULT LOG

# ============================================================


def display_vault():

    vault = load_vault()

    files = vault.get("files", [])


    if not files:

        print("\n  No operations logged yet.")

        return


    print("\n" + "="*65)

    print(f"  ENCRYPTION VAULT LOG  ({len(files)} entries)")

    print("="*65)

    print(f"  {'OP':<10} {'TIMESTAMP':<22} {'METHOD':<10} FILE")

    print("  " + "-"*60)


    for entry in reversed(files[-20:]):

        op     = entry["operation"].upper()

        ts     = entry["timestamp"]

        method = entry["method"]

        fname  = Path(entry["output"]).name

        print(f"  {op:<10} {ts:<22} {method:<10} {fname}")


    print("="*65)



# ============================================================

# SETUP: GET OR CREATE KEY

# ============================================================


def setup_key_menu():

    print("\n  KEY SETUP")

    print("  1. Generate new random key (save to file)")

    print("  2. Use existing key file")

    print("  3. Use password (derive key)")


    choice = input("\n  Choice (1/2/3): ").strip()


    if choice == "1":

        key = generate_key()

        fname = input(f"  Save key as [{KEY_FILE}]: ").strip() or KEY_FILE

        save_key(key, fname)

        return key, False, None


    elif choice == "2":

        fname = input(f"  Key file path [{KEY_FILE}]: ").strip() or KEY_FILE

        key = load_key(fname)

        if key:

            print(f"  Key loaded from: {fname}")

            return key, False, None

        return None, False, None


    elif choice == "3":

        password = getpass.getpass("  Enter password: ")

        confirm  = getpass.getpass("  Confirm password: ")

        if password != confirm:

            print("  Passwords do not match.")

            return None, False, None

        key, salt = derive_key_from_password(password)

        print("  Key derived from password successfully.")

        print(f"  Strength tip: Use 12+ chars with mixed letters, numbers, symbols.")

        return key, True, salt


    return None, False, None



# ============================================================

# MAIN MENU

# ============================================================


def print_menu(key_loaded, method):

    status = f"Key loaded ({method})" if key_loaded else "No key loaded"

    print("\n" + "-"*50)

    print(f"  FILE ENCRYPTION TOOL  [{status}]")

    print("-"*50)

    print("  1. Setup key (generate / load / password)")

    print("  2. Encrypt a file")

    print("  3. Decrypt a file")

    print("  4. Encrypt all files in folder")

    print("  5. Decrypt all .encrypted files in folder")

    print("  6. Encrypt text string")

    print("  7. Decrypt text string")

    print("  8. View vault log")

    print("  0. Exit")

    print("-"*50)



def main():

    print("\n" + "="*55)

    print("     FILE ENCRYPTION / DECRYPTION TOOL")

    print("="*55)


    if not CRYPTO_OK:

        print("\n  cryptography library not installed!")

        print("  Install it:  pip install cryptography")

        return


    print("\n  Uses Fernet symmetric encryption (AES-128-CBC + HMAC-SHA256)")

    print("  Supports: key file OR password-based encryption\n")


    current_key      = None

    password_based   = False

    current_salt     = None

    method_str       = ""


    # Auto-load key if exists

    if Path(KEY_FILE).exists():

        current_key = load_key(KEY_FILE)

        method_str  = "key-file"

        print(f"  Auto-loaded key from: {KEY_FILE}")


    vault = load_vault()


    while True:

        print_menu(current_key is not None, method_str)

        choice = input("  > ").strip()


        if choice == "1":

            current_key, password_based, current_salt = setup_key_menu()

            if current_key:

                method_str = "password" if password_based else "key-file"


        elif choice in ["2", "3", "4", "5", "6", "7"]:

            if not current_key:

                print("\n  No key loaded. Please set up a key first (Option 1).")

                continue


            if choice == "2":

                path = input("\n  File to encrypt: ").strip()

                out  = input("  Output path (Enter = auto): ").strip() or None

                result = encrypt_file(path, current_key, out,

                                      salt=current_salt)

                if result:

                    log_operation(vault, "encrypt", path, result, method_str)

                    save_vault(vault)


            elif choice == "3":

                path = input("\n  File to decrypt: ").strip()

                out  = input("  Output path (Enter = auto): ").strip() or None

                result = decrypt_file(path, current_key, out,

                                      password_based=password_based)

                if result:

                    log_operation(vault, "decrypt", path, result, method_str)

                    save_vault(vault)


            elif choice == "4":

                folder = input("\n  Folder path: ").strip()

                ext_in = input("  File types to encrypt (e.g. .txt .pdf, Enter=all): ").strip()

                exts   = [e if e.startswith(".") else f".{e}"

                          for e in ext_in.split()] if ext_in else None

                process_folder(folder, current_key, "encrypt",

                               extensions=exts,

                               salt=current_salt)


            elif choice == "5":

                folder = input("\n  Folder path: ").strip()

                process_folder(folder, current_key, "decrypt",

                               password_based=password_based)


            elif choice == "6":

                encrypt_text(current_key)


            elif choice == "7":

                decrypt_text(current_key)


        elif choice == "8":

            display_vault()


        elif choice == "0":

            print("\n  Goodbye!\n")

            break


        else:

            print("  Invalid choice.")



# ============================================================

# RUN

# ============================================================


if __name__ == "__main__":

    main()

No comments: