File encryption and decryption utilities

Simple functions to encrypt and decrypt files.

It uses Fernet symmetric encryption scheme from the cryptography library.

A password is transformed into a Fernet-compatible key using SHA-256. Note: this is a simplified key-derivation method and is not as secure as PBKDF2HMAC or scrypt for real-world applications.

import base64
import hashlib
import os
import pathlib
import tempfile

from cryptography.fernet import Fernet

# -------------------------------------------------------------------------
# Key generation
# -------------------------------------------------------------------------


def generate_key(password: str) -> bytes:
    """
    Generate a Fernet-compatible key from a password using SHA-256.

    This function hashes the password twice with SHA-256 and encodes the
    result in URL-safe Base64, as required by Fernet.

    Parameters
    ----------
    password : str
        The password used to derive the encryption key.

    Returns
    -------
    bytes
        A 32-byte Base64-encoded key suitable for Fernet.

    Notes
    -----
    This is a simplified key derivation function. For stronger security,
    consider using PBKDF2HMAC with a salt and many iterations.
    """
    digest = hashlib.sha256(password.encode()).digest()
    return base64.urlsafe_b64encode(hashlib.sha256(digest).digest())


# -------------------------------------------------------------------------
# File encryption
# -------------------------------------------------------------------------

def encrypt_file(file_path: str, key: bytes):
    """
    Encrypt a file in place using the Fernet algorithm.

    The file is read entirely into memory, encrypted, and then overwritten
    with the encrypted content.

    Parameters
    ----------
    file_path : str
        Path to the file to encrypt.
    key : bytes
        A Fernet key generated by `generate_key()`.

    Returns
    -------
    Fernet
        The Fernet instance used for encryption.

    Notes
    -----
    - The original file content is overwritten and cannot be recovered
      without the correct key.
    - If you need to preserve the original file, copy it before encrypting.
    """
    fernet = Fernet(key)

    original = pathlib.Path(file_path).read_bytes()

    encrypted = fernet.encrypt(original)

    pathlib.Path(file_path).write_bytes(encrypted)

    return fernet


# -------------------------------------------------------------------------
# File decryption
# -------------------------------------------------------------------------

def decrypt_file(file_path: str, key: bytes):
    """
    Decrypt a file in place using the Fernet algorithm.

    The file is read entirely into memory, decrypted, and then overwritten
    with the decrypted content.

    Parameters
    ----------
    file_path : str
        Path to the file to decrypt.
    key : bytes
        A Fernet key generated by `generate_key()`.

    Returns
    -------
    None

    Notes
    -----
    - If the key is invalid, the function prints an error and leaves the
      encrypted file unchanged.
    - The decrypted content overwrites the encrypted file.
    """
    encrypted = pathlib.Path(file_path).read_bytes()

    fernet = Fernet(key)

    try:
        decrypted = fernet.decrypt(encrypted)
    except Exception:
        print("Error: invalid key. File was not decrypted.")
        return

    pathlib.Path(file_path).write_bytes(decrypted)

Example usage of the encryption functions

# 1. Create a temporary file with known content
original_content = b"Hello, this is a test message."
with tempfile.NamedTemporaryFile(delete=False) as tmp:
    file_path = tmp.name
    tmp.write(original_content)
print(f"Temporary file created at: {file_path}")

# 2. Generate a key from a password
password = "my_password"
key = generate_key(password)

# 3. Encrypt the file
encrypt_file(file_path, key)
print("File encrypted.")

# 4. Decrypt the file
decrypt_file(file_path, key)
print("File decrypted.")

# 5. Read back the content and verify equality
final_content = pathlib.Path(file_path).read_bytes()
if final_content == original_content:
    print("Success: decrypted content matches the original.")
else:
    print("Error: decrypted content does NOT match the original.")

# 6. Clean up
os.remove(file_path)

Temporary file created at: /tmp/tmp5289r_25
File encrypted.
File decrypted.
Success: decrypted content matches the original.

Estimated memory usage: 118 MB