Encrypt a file and get it decrypted using FAST-API

Tiempo de lectura: 2 minutos

Today we are going to learn how to store an encrypted file and how to decrypt it to return it in a response.

We will use the Fernet library.

pip install cryptography

First, let’s create an encryption key with Fernet.

def generate_key():
    key = Fernet.generate_key()
    return key.decode('utf-8')

With this function, we can obtain a random key for Fernet. It should be stored securely as it will be used to decrypt encrypted files.

Now, with these functions, we can encrypt and decrypt our files, which will be stored with a .enc extension.

from cryptography.fernet import Fernet
import os
from dotenv import load_dotenv

load_dotenv()
FERNET_KEY = bytes(os.getenv("FERNET_KEY"), "utf-8")

fernet = Fernet(FERNET_KEY)

def generate_key():
    key = Fernet.generate_key()
    return key.decode('utf-8')

def cipher_file(original_file: bytes, save_path: str):
    new_file_name = save_path + ".enc"
    encrypted_data = fernet.encrypt(original_file)

    with open(new_file_name, "wb") as encrypted_file:
        encrypted_file.write(encrypted_data)

def decipher_file(encrypted_file_path: str):
    # Read encrypted content from the file
    with open(encrypted_file_path, "rb") as encrypted_file:
        encrypted_data = encrypted_file.read()

    # Decrypt the content
    decrypted_data = fernet.decrypt(encrypted_data)

    # Return the decrypted data directly
    return decrypted_data

def decipher_file_in_path(encrypted_file_path: str):
    with open(encrypted_file_path, "rb") as encrypted_file:
        encrypted_data = encrypted_file.read()
        decrypted_data = fernet.decrypt(encrypted_data)

    return decrypted_data

With the cipher_file function, it will encrypt the file passed as bytes and store it in the specified path with .enc appended.

To convert files to bytes, we use this function:

@api.post("/add_file", status_code=status.HTTP_201_CREATED)
def add_file(pdf_file: UploadFile = File(...)):
    upload_folder = "pdf"

    # Get the current date
    now = datetime.now()
    dt_string = now.strftime("%d_%m_%Y_%H_%M_%S")

    file_name = dt_string + ".pdf"
    file_route = os.path.join(upload_folder, file_name)
    cipher_file(pdf_file.file.read(), file_route)
   
    return {"filename": file_name}

Files of type File(…) are uploaded and stored encrypted.

To retrieve the encrypted file, use the decipher_file method, which will return the decrypted file as bytes. To return this file in File format, do the following:

@api.get("/get_file/{path}", status_code=status.HTTP_200_OK)
def get_file(path: str):
   
    # Check if the file exists
    if not os.path.exists(path):
        return {"error": "File does not exist"}

    # Get decrypted_data from fernet.decrypt
    decrypted_data = decipher_file(file_path)
    
    # Return the decrypted file decrypted data as a pdf type (without a path)
    return StreamingResponse(io.BytesIO(decrypted_data), media_type="application/pdf", headers={"Content-Disposition": f"attachment; filename={"file"}"})

This way, the decrypted file is returned.

Finally, if you want to decrypt the file within the directory itself, use the decipher_file_in_path function, which directly overwrites the file within the directory.

Leave a Comment