encryption.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
File Encryption Module for Cloud Storage Bot

This module provides encryption and decryption functionality for the Cloud Storage Bot.
It uses Fernet symmetric encryption from the cryptography library to secure files.
"""

import os
import base64
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC


class FileEncryption:
    """Handles file encryption and decryption operations."""

    def __init__(self, password):
        """Initialize with a password for encryption/decryption.

        Args:
            password (str): The password used for encryption/decryption.
        """
        self.password = password.encode()
        self.salt = os.urandom(16)  # Generate a random salt
        self.key = self._generate_key()
        self.cipher = Fernet(self.key)

    def _generate_key(self):
        """Generate a key from the password using PBKDF2.

        Returns:
            bytes: The generated key for encryption/decryption.
        """
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=self.salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(self.password))
        return key

    def encrypt_file(self, file_path, output_path=None):
        """Encrypt a file using the generated key.

        Args:
            file_path (str): Path to the file to encrypt.
            output_path (str, optional): Path to save the encrypted file.
                If not provided, appends '.encrypted' to the original filename.

        Returns:
            str: Path to the encrypted file.
            bytes: The salt used for encryption (should be stored securely).
        """
        if not output_path:
            output_path = f"{file_path}.encrypted"

        # Read the file content
        with open(file_path, 'rb') as file:
            file_data = file.read()

        # Encrypt the data
        encrypted_data = self.cipher.encrypt(file_data)

        # Write the encrypted data to the output file
        with open(output_path, 'wb') as file:
            file.write(encrypted_data)

        return output_path, self.salt

    def decrypt_file(self, file_path, output_path=None, salt=None):
        """Decrypt a file using the generated key.

        Args:
            file_path (str): Path to the encrypted file.
            output_path (str, optional): Path to save the decrypted file.
                If not provided, removes '.encrypted' from the filename if present.
            salt (bytes, optional): The salt used during encryption.
                If provided, regenerates the key with this salt.

        Returns:
            str: Path to the decrypted file.
        """
        if salt:
            # Regenerate the key with the provided salt
            self.salt = salt
            self.key = self._generate_key()
            self.cipher = Fernet(self.key)

        if not output_path:
            if file_path.endswith('.encrypted'):
                output_path = file_path[:-10]  # Remove '.encrypted' suffix
            else:
                output_path = f"{file_path}.decrypted"

        # Read the encrypted file content
        with open(file_path, 'rb') as file:
            encrypted_data = file.read()

        # Decrypt the data
        decrypted_data = self.cipher.decrypt(encrypted_data)

        # Write the decrypted data to the output file
        with open(output_path, 'wb') as file:
            file.write(decrypted_data)

        return output_path


# Example usage
def encrypt_user_file(file_path, password, output_path=None):
    """Encrypt a user file with the provided password.

    Args:
        file_path (str): Path to the file to encrypt.
        password (str): Password for encryption.
        output_path (str, optional): Path to save the encrypted file.

    Returns:
        tuple: (encrypted_file_path, salt)
    """
    encryption = FileEncryption(password)
    return encryption.encrypt_file(file_path, output_path)


def decrypt_user_file(file_path, password, salt, output_path=None):
    """Decrypt a user file with the provided password and salt.

    Args:
        file_path (str): Path to the encrypted file.
        password (str): Password for decryption.
        salt (bytes): Salt used during encryption.
        output_path (str, optional): Path to save the decrypted file.

    Returns:
        str: Path to the decrypted file.
    """
    encryption = FileEncryption(password)
    return encryption.decrypt_file(file_path, output_path, salt)