backend.py

"""Backend functions used in the app."""

import os
import shutil
import gradio as gr
import numpy
import requests
import pickle
import pandas
from itertools import chain

from settings import (
    SERVER_URL,
    FHE_KEYS,
    CLIENT_FILES,
    SERVER_FILES,
    DEPLOYMENT_PATH,
    INITIAL_INPUT_SHAPE,
    INPUT_INDEXES,
    INPUT_SLICES,
    PRE_PROCESSOR_USER_PATH, 
    PRE_PROCESSOR_THIRD_PARTY_PATH,
    CLIENT_TYPES,
)

from utils.client_server_interface import MultiInputsFHEModelClient

# Load pre-processor instances
with PRE_PROCESSOR_USER_PATH.open('rb') as file:
    PRE_PROCESSOR_USER = pickle.load(file)

with PRE_PROCESSOR_THIRD_PARTY_PATH.open('rb') as file:
    PRE_PROCESSOR_THIRD_PARTY = pickle.load(file)


def shorten_bytes_object(bytes_object, limit=500):
    """Shorten the input bytes object to a given length.

    Encrypted data is too large for displaying it in the browser using Gradio. This function
    provides a shorten representation of it.

    Args:
        bytes_object (bytes): The input to shorten
        limit (int): The length to consider. Default to 500.

    Returns:
        str: Hexadecimal string shorten representation of the input byte object. 

    """
    # Define a shift for better display
    shift = 100
    return bytes_object[shift : limit + shift].hex()


def clean_temporary_files(n_keys=20):
    """Clean keys and encrypted inputs.

    A maximum of n_keys keys and associated temporary files are allowed to be stored. Once this 
    limit is reached, the oldest files are deleted.

    Args:
        n_keys (int): The maximum number of keys and associated files to be stored. Default to 20.

    """
    # Get the oldest key files in the key directory
    key_dirs = sorted(FHE_KEYS.iterdir(), key=os.path.getmtime)

    # If more than n_keys keys are found, remove the oldest
    user_ids = []
    if len(key_dirs) > n_keys:
        n_keys_to_delete = len(key_dirs) - n_keys
        for key_dir in key_dirs[:n_keys_to_delete]:
            user_ids.append(key_dir.name)
            shutil.rmtree(key_dir)

    # Get all the encrypted objects in the temporary folder
    client_files = CLIENT_FILES.iterdir()
    server_files = SERVER_FILES.iterdir()

    # Delete all files related to the ids whose keys were deleted
    for file in chain(client_files, server_files):
        for user_id in user_ids:
            if user_id in file.name:
                file.unlink()


def _get_client(client_id, client_type):
    """Get the client API.

    Args:
        client_id (int): The client ID to consider.
        client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party').

    Returns:
        FHEModelClient: The client API.
    """
    key_dir = FHE_KEYS / f"{client_type}_{client_id}"

    return MultiInputsFHEModelClient(DEPLOYMENT_PATH, key_dir=key_dir, nb_inputs=len(CLIENT_TYPES))


def _keygen(client_id, client_type):
    """Generate the private key associated to a client.

    Args:
        client_id (int): The client ID to consider.
        client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
    """
    # Clean temporary files
    clean_temporary_files()

    # Retrieve the client instance
    client = _get_client(client_id, client_type)

    # Generate a private key
    client.generate_private_and_evaluation_keys(force=True)

    # Retrieve the serialized evaluation key. In this case, as circuits are fully leveled, this
    # evaluation key is empty. However, for software reasons, it is still needed for proper FHE
    # execution
    evaluation_key = client.get_serialized_evaluation_keys()

    # Save evaluation_key as bytes in a file as it is too large to pass through regular Gradio
    # buttons (see https://github.com/gradio-app/gradio/issues/1877)
    evaluation_key_path = _get_client_file_path("evaluation_key", client_id, client_type)

    with evaluation_key_path.open("wb") as evaluation_key_file:
        evaluation_key_file.write(evaluation_key)


def _send_input(client_id, client_type):
    """Send the encrypted inputs as well as the evaluation key to the server.

    Args:
        client_id (int): The client ID to consider.
        client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
    """
    # Get the paths to the evaluation key and encrypted inputs
    evaluation_key_path = _get_client_file_path("evaluation_key", client_id, client_type)
    encrypted_input_path = _get_client_file_path("encrypted_inputs", client_id, client_type)

    # Define the data and files to post
    data = {
        "client_id": client_id,
        "client_type": client_type,
    }

    files = [
        ("files", open(encrypted_input_path, "rb")),
        ("files", open(evaluation_key_path, "rb")),
    ]

    # Send the encrypted inputs and evaluation key to the server
    url = SERVER_URL + "send_input"
    with requests.post(
        url=url,
        data=data,
        files=files,
    ) as response:
        return response.ok


def _get_client_file_path(name, client_id, client_type):
    """Get the correct temporary file path for the client.

    Args:
        name (str): The desired file name (either 'evaluation_key' or 'encrypted_inputs').
        client_id (int): The client ID to consider.
        client_type (str): The type of user to consider (either 'user', 'bank' or 'third_party').

    Returns:
        pathlib.Path: The file path.
    """
    return CLIENT_FILES / f"{name}_{client_type}_{client_id}"


def _keygen_encrypt_send(inputs, client_type):
    """Encrypt the given inputs for a specific client.

    Args:
        inputs (numpy.ndarray): The inputs to encrypt.
        client_type (str): The type of client to consider (either 'user', 'bank' or 'third_party').
    
    Returns:
        client_id, encrypted_inputs_short (int, bytes): Integer ID representing the current client 
            and a byte short representation of the encrypted input to send. 
    """
    # Create an ID for the current client to consider
    client_id = numpy.random.randint(0, 2**32)

    _keygen(client_id, client_type)

    # Retrieve the client instance
    client = _get_client(client_id, client_type)

    # TODO : pre-process the data first 

    # Quantize, encrypt and serialize the inputs
    encrypted_inputs = client.quantize_encrypt_serialize_multi_inputs(
        inputs, 
        input_index=INPUT_INDEXES[client_type], 
        initial_input_shape=INITIAL_INPUT_SHAPE, 
        input_slice=INPUT_SLICES[client_type],
    )

    # Save encrypted_inputs to bytes in a file, since too large to pass through regular Gradio
    # buttons, https://github.com/gradio-app/gradio/issues/1877
    encrypted_inputs_path = _get_client_file_path("encrypted_inputs", client_id, client_type)

    with encrypted_inputs_path.open("wb") as encrypted_inputs_file:
        encrypted_inputs_file.write(encrypted_inputs)

    # Create a truncated version of the encrypted inputs for display
    encrypted_inputs_short = shorten_bytes_object(encrypted_inputs)

    _send_input(client_id, client_type)

    # TODO: also return private key representation if possible
    return client_id, encrypted_inputs_short


def pre_process_keygen_encrypt_send_user(*inputs):
    """Pre-process the given inputs for a specific client.

    Args:
        *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
    
    Returns:
        (int, bytes): Integer ID representing the current client and a byte short representation of 
            the encrypted input to send. 
    """
    gender, bool_inputs, num_children, num_family, total_income, age, income_type, education_type, \
        family_status, occupation_type, housing_type = inputs
    
    # Encoding given in https://www.kaggle.com/code/samuelcortinhas/credit-cards-data-cleaning 
    # for "Gender" is M ('Male') -> 1 and F ('Female') -> 0
    gender = gender == "Male"

    # Retrieve boolean values
    own_car = "Car" in bool_inputs
    own_property = "Property" in bool_inputs
    work_phone = "Work phone" in bool_inputs
    phone = "Phone" in bool_inputs
    email = "Email" in bool_inputs

    user_inputs = pandas.DataFrame({
        "Gender": [gender],
        "Own_car": [own_car],
        "Own_property": [own_property],
        "Work_phone": [work_phone],
        "Phone": [phone],
        "Email": [email],
        "Num_children": num_children,
        "Num_family": num_family,
        "Total_income": total_income,
        "Age": age,
        "Income_type": income_type,
        "Education_type": education_type,
        "Family_status": family_status,
        "Occupation_type": occupation_type,
        "Housing_type": housing_type,
    })

    preprocessed_user_inputs = PRE_PROCESSOR_USER.transform(user_inputs)

    return _keygen_encrypt_send(preprocessed_user_inputs, "user")


def pre_process_keygen_encrypt_send_bank(*inputs):
    """Pre-process the given inputs for a specific client.

    Args:        
        *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.
    
    Returns:
        (int, bytes): Integer ID representing the current client and a byte short representation of 
            the encrypted input to send. 
    """
    account_length = inputs[0]
    
    return _keygen_encrypt_send(account_length, "bank")
    

def pre_process_keygen_encrypt_send_third_party(*inputs):
    """Pre-process the given inputs for a specific client.

    Args:
        *inputs (Tuple[numpy.ndarray]): The inputs to pre-process.

    Returns:
        (int, bytes): Integer ID representing the current client and a byte short representation of 
            the encrypted input to send. 
    """
    employed, years_employed = inputs
    
    # Original dataset contains an "unemployed" feature instead of "employed"
    unemployed = employed == "No"

    third_party_inputs = pandas.DataFrame({
        "Unemployed": [unemployed],
        "Years_employed": [years_employed],
    })

    preprocessed_third_party_inputs = PRE_PROCESSOR_THIRD_PARTY.transform(third_party_inputs)

    return _keygen_encrypt_send(preprocessed_third_party_inputs, "third_party")


def run_fhe(user_id, bank_id, third_party_id):
    """Run the model on the encrypted inputs previously sent using FHE.

    Args:
        user_id (int): The user ID to consider.
        bank_id (int): The bank ID to consider.
        third_party_id (int): The third party ID to consider.
    """

    # TODO : add a warning for users to send all client types' inputs 

    data = {
        "user_id": user_id,
        "bank_id": bank_id,
        "third_party_id": third_party_id,
    }

    # Trigger the FHE execution on the encrypted inputs previously sent
    url = SERVER_URL + "run_fhe"
    with requests.post(
        url=url,
        data=data,
    ) as response:
        if response.ok:
            return response.json()
        else:
            raise gr.Error("Please wait for the inputs to be sent to the server.")


def get_output(user_id, bank_id, third_party_id):
    """Retrieve the encrypted output.

    Args:
        user_id (int): The user ID to consider.
        bank_id (int): The bank ID to consider.
        third_party_id (int): The third party ID to consider.
    
    Returns:
        encrypted_output_short (bytes): A byte short representation of the encrypted output. 
    """
    data = {
        "user_id": user_id,
        "bank_id": bank_id,
        "third_party_id": third_party_id,
    }

    # Retrieve the encrypted output
    url = SERVER_URL + "get_output"
    with requests.post(
        url=url,
        data=data,
    ) as response:
        if response.ok:
            encrypted_output = response.content

            # Save the encrypted output to bytes in a file as it is too large to pass through regular
            # Gradio buttons (see https://github.com/gradio-app/gradio/issues/1877)
            encrypted_output_path = _get_client_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")

            with encrypted_output_path.open("wb") as encrypted_output_file:
                encrypted_output_file.write(encrypted_output)
            
            # Create a truncated version of the encrypted inputs for display
            encrypted_output_short = shorten_bytes_object(encrypted_output)

            return encrypted_output_short
        else:
            raise gr.Error("Please wait for the FHE execution to be completed.")


def decrypt_output(user_id, bank_id, third_party_id):
    """Decrypt the result.

    Args:
        user_id (int): The user ID to consider.
        bank_id (int): The bank ID to consider.
        third_party_id (int): The third party ID to consider.

    Returns:
        output(numpy.ndarray): The decrypted output

    """
    # Get the encrypted output path
    encrypted_output_path = _get_client_file_path("encrypted_output", user_id + bank_id + third_party_id, "output")

    if not encrypted_output_path.is_file():
        raise gr.Error("Please run the FHE execution first.")

    # Load the encrypted output as bytes
    with encrypted_output_path.open("rb") as encrypted_output_file:
        encrypted_output_proba = encrypted_output_file.read()

    # Retrieve the client API
    client = _get_client(user_id, "user")

    # Deserialize, decrypt and post-process the encrypted output
    output_proba = client.deserialize_decrypt_dequantize(encrypted_output_proba)

    # Determine the predicted class
    output = numpy.argmax(output_proba, axis=1)

    return output