SEC-File.py

#!/usr/bin/env python3
import os
import sys
import getpass
import base64
import logging
import shutil
import subprocess
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes

#Banner part
subprocess.run(["figlet", "SEC-File"])
print(">>---S-T-A-R-T-E-D---<<\n")

# Configure logging
logging.basicConfig(filename='encryption.log', level=logging.INFO, format='%(asctime)s - %(message)s')

def derive_key(password: str, salt: bytes) -> bytes:
    """Derive a key from the given password and salt using PBKDF2."""
    kdf = PBKDF2HMAC(
        algorithm=hashes.SHA256(),
        length=32,
        salt=salt,
        iterations=100000,
        backend=default_backend()
    )
    return base64.urlsafe_b64encode(kdf.derive(password.encode()))

def generate_salt() -> bytes:
    """Generate a random salt."""
    return os.urandom(16)

def secure_delete(file_path: str):
    """Securely delete a file by overwriting it and then removing it."""
    try:
        with open(file_path, "r+b") as file:
            length = os.path.getsize(file_path)
            file.write(os.urandom(length))  # Overwrite with random data
        os.remove(file_path)
        logging.info(f"Securely deleted {file_path}")
    except Exception as e:
        logging.error(f"Error securely deleting {file_path}: {e}")
        print(f"-> Error securely deleting {file_path}: {e}")

def encrypt_file(file_path: str, fernet: Fernet):
    """Encrypt a single file."""
    try:
        with open(file_path, "rb") as file:
            file_data = file.read()
        encrypted_data = fernet.encrypt(file_data)
        with open(file_path + ".enc", "wb") as file:
            file.write(encrypted_data)
        logging.info(f"Encrypted {file_path}")
        print(f"Encrypted {file_path}")
    except Exception as e:
        logging.error(f"Error encrypting {file_path}: {e}")
        print(f"Error encrypting {file_path}: {e}")

def decrypt_file(file_path: str, fernet: Fernet):
    """Decrypt a single file and delete the encrypted version."""
    try:
        with open(file_path, "rb") as file:
            encrypted_data = file.read()
        decrypted_data = fernet.decrypt(encrypted_data)
        with open(file_path[:-4], "wb") as file:  # Remove .enc extension
            file.write(decrypted_data)
        logging.info(f"Decrypted {file_path}")
        print(f"Decrypted {file_path}")
        
        # Securely delete the encrypted file
        secure_delete(file_path)  # Call to secure delete the encrypted file
    except Exception as e:
        logging.error(f"Error decrypting {file_path}: {e}")
        print(f"Error decrypting {file_path}: {e}")

def encrypt_folder(folder_path: str, password: str, backup: bool):
    """Encrypt all files in the specified folder."""
    if not os.path.exists(folder_path):
        print("-> Folder Does not exist. Try again.")
        return
    
    # Get unique file extensions in the folder
    file_types = set()
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            file_types.add(os.path.splitext(file)[1])  # Get file extension

    if not file_types:
        print("-> Folder is Empty, There is Nothing To Do.")
        return

    salt = generate_salt()
    key = derive_key(password, salt)
    fernet = Fernet(key)

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if os.path.splitext(file)[1] in file_types:
                file_path = os.path.join(root, file)
                if backup:
                    shutil.copy(file_path, file_path + ".bak")  # Backup original file
                encrypt_file(file_path, fernet)
                secure_delete(file_path)  # Securely delete the original file

    # Save the salt for decryption
    with open(os.path.join(folder_path, "salt.bin"), "wb") as salt_file:
        salt_file.write(salt)

def decrypt_folder(folder_path: str, password: str):
    """Decrypt all files in the specified folder."""
    if not os.path.exists(folder_path):
        print("-> Folder Does not exist. Check The Path again.")
        return
    
    try:
        with open(os.path.join(folder_path, "salt.bin"), "rb") as salt_file:
            salt = salt_file.read()
    except FileNotFoundError:
        print("-> Salt File Not Found. Cannot Decrypt.")
        return

    try:
        key = derive_key(password, salt)
        fernet = Fernet(key)
    except Exception as e:
        print(f"-> Error Deriving key or creating Fernet instance: {e}")
        return

    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.endswith('.enc'):  # Only decrypt files with .enc extension
                file_path = os.path.join(root, file)
                decrypt_file(file_path, fernet)  # This will now delete the encrypted file after decryption

def main():
    """Main function to handle user input and call encryption/decryption functions."""
    action = input(">> Enter (e)ncrypt or (d)ecrypt: ").strip().lower()
    folder_path = input(">> Enter The Folder Path (Full Path): ").strip()
    
    # Check if the folder exists
    if not os.path.exists(folder_path):
        print("The specified folder does not exist. Please check the path and try again.")
        return  # Exit the function if the folder does not exist
    
    if action == 'e':
        password = getpass.getpass("\n>> Enter A Password For Encryption: ")
        backup = input(">> Do You Want Original Files BackUP (y/n): ").strip().lower() == 'y'
        encrypt_folder(folder_path, password, backup)
    elif action == 'd':
        password = getpass.getpass("\n>> Enter The Password For Decryption: ")
        decrypt_folder(folder_path, password)
    else:
        print("-> Invalid action. Please enter 'e' for encrypt or 'd' for decrypt.")

if __name__ == "__main__":
    main()