Upload app.py

app.py

@@ -1,14 +1,198 @@
 import gradio as gr
+import time
+import os
+from pathlib import Path
+
+import subprocess
+from concrete.ml.deployment import FHEModelClient
+from requests import head
+import numpy
+import os
+from pathlib import Path
+import requests
+import json
+import base64
+import subprocess
+import shutil
+import time
+import pandas as pd
+import pickle
+import numpy as np
+
+# This repository's directory
+REPO_DIR = Path(__file__).parent
+subprocess.Popen(["uvicorn", "server:app"], cwd=REPO_DIR)
+
+# if not exists, create a directory for the FHE keys called .fhe_keys
+if not os.path.exists(".fhe_keys"):
+    os.mkdir(".fhe_keys")
+# if not exists, create a directory for the tmp files called tmp
+if not os.path.exists("tmp"):
+    os.mkdir("tmp")
+
+
+# Wait 4 sec for the server to start
+time.sleep(4)
+
+
+# Encrypted data limit for the browser to display
+# (encrypted data is too large to display in the browser)
+ENCRYPTED_DATA_BROWSER_LIMIT = 500
+N_USER_KEY_STORED = 20
+
+
+def clean_tmp_directory():
+    # Allow 20 user keys to be stored.
+    # Once that limitation is reached, deleted the oldest.
+    path_sub_directories = sorted(
+        [f for f in Path(".fhe_keys/").iterdir() if f.is_dir()], key=os.path.getmtime
+    )
+
+    user_ids = []
+    if len(path_sub_directories) > N_USER_KEY_STORED:
+        n_files_to_delete = len(path_sub_directories) - N_USER_KEY_STORED
+        for p in path_sub_directories[:n_files_to_delete]:
+            user_ids.append(p.name)
+            shutil.rmtree(p)
+
+    list_files_tmp = Path("tmp/").iterdir()
+    # Delete all files related to user_id
+    for file in list_files_tmp:
+        for user_id in user_ids:
+            if file.name.endswith(f"{user_id}.npy"):
+                file.unlink()
+
+
+def keygen():
+    # Clean tmp directory if needed
+    clean_tmp_directory()
+
+    print("Initializing FHEModelClient...")
+    # Let's create a user_id
+    user_id = numpy.random.randint(0, 2**32)
+    fhe_api = FHEModelClient(f"deployment/deployment_{task}", f".fhe_keys/{user_id}")
+    fhe_api.load()
+
+    # Generate a fresh key
+    fhe_api.generate_private_and_evaluation_keys(force=True)
+    evaluation_key = fhe_api.get_serialized_evaluation_keys()
+
+    numpy.save(f"tmp/tmp_evaluation_key_{user_id}.npy", evaluation_key)
+
+    return [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], user_id]
+
+
+def encode_quantize_encrypt(test_file, user_id):
+
+    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
+    fhe_api.load()
+    from PE_main import extract_infos
+
+    features = pickle.loads(open(os.path.join("features.pkl"), "rb").read())
+    encodings = extract_infos(test_file)
+    encodings = list(map(lambda x: encodings[x], features))
+
+    quantized_encodings = fhe_api.model.quantize_input(encodings).astype(numpy.uint8)
+    encrypted_quantized_encoding = fhe_api.quantize_encrypt_serialize(encodings)
+
+    # Save encrypted_quantized_encoding in a file, since too large to pass through regular Gradio
+    # buttons, https://github.com/gradio-app/gradio/issues/1877
+    numpy.save(
+        f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy",
+        encrypted_quantized_encoding,
+    )
+
+    # Compute size
+    encrypted_quantized_encoding_shorten = list(encrypted_quantized_encoding)[
+        :ENCRYPTED_DATA_BROWSER_LIMIT
+    ]
+    encrypted_quantized_encoding_shorten_hex = "".join(
+        f"{i:02x}" for i in encrypted_quantized_encoding_shorten
+    )
+    return (
+        encodings[0],
+        quantized_encodings[0],
+        encrypted_quantized_encoding_shorten_hex,
+    )
+
+
+def run_fhe(user_id):
+    encoded_data_path = Path(f"tmp/tmp_encrypted_quantized_encoding_{user_id}.npy")
+    encrypted_quantized_encoding = numpy.load(encoded_data_path)
+
+    # Read evaluation_key from the file
+    evaluation_key = numpy.load(f"tmp/tmp_evaluation_key_{user_id}.npy")
+
+    # Use base64 to encode the encodings and evaluation key
+    encrypted_quantized_encoding = base64.b64encode(
+        encrypted_quantized_encoding
+    ).decode()
+    encoded_evaluation_key = base64.b64encode(evaluation_key).decode()
+
+    query = {}
+    query["evaluation_key"] = encoded_evaluation_key
+    query["encrypted_encoding"] = encrypted_quantized_encoding
+    headers = {"Content-type": "application/json"}
+
+
+    response = requests.post(
+        "http://localhost:8000/predict",
+        data=json.dumps(query),
+        headers=headers,
+    )
+
+
+    encrypted_prediction = base64.b64decode(response.json()["encrypted_prediction"])
+
+    numpy.save(f"tmp/tmp_encrypted_prediction_{user_id}.npy", encrypted_prediction)
+    encrypted_prediction_shorten = list(encrypted_prediction)[
+        :ENCRYPTED_DATA_BROWSER_LIMIT
+    ]
+    encrypted_prediction_shorten_hex = "".join(
+        f"{i:02x}" for i in encrypted_prediction_shorten
+    )
+
+
+def decrypt_prediction(user_id):
+    encoded_data_path = Path(f"tmp/tmp_encrypted_prediction_{user_id}.npy")
+
+    # Read encrypted_prediction from the file
+
+    encrypted_prediction = numpy.load(encoded_data_path).tobytes()
+
+    fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
+    fhe_api.load()
+
+    # We need to retrieve the private key that matches the client specs (see issue #18)
+    fhe_api.generate_private_and_evaluation_keys(force=False)
+
+    predictions = fhe_api.deserialize_decrypt_dequantize(encrypted_prediction)
 
 
 def update(name):
     return f"Welcome to Gradio, {name}!"
 
 
-with gr.Blocks() as app:
-    gr.Markdown("### Start typing below and then click **Run** to see the output.")
-    with gr.Column():
-        gr.File(interactive=False)
-        gr.UploadButton(label="Upload an executable file")
-
-app.launch()
+if __name__ == "__main__":
+    app = gr.Interface(
+        [
+            keygen,
+            encode_quantize_encrypt,
+            run_fhe,
+            decrypt_prediction,
+        ],
+        [
+            gr.inputs.Textbox(label="Task", default="malware"),
+            gr.inputs.File(label="Test File"),
+            gr.inputs.Textbox(label="User ID"),
+        ],
+        [
+            gr.outputs.Textbox(label="Evaluation Key"),
+            gr.outputs.Textbox(label="Encodings"),
+            gr.outputs.Textbox(label="Encrypted Quantized Encoding"),
+            gr.outputs.Textbox(label="Encrypted Prediction"),
+        ],
+        title="FHE Model",
+        description="This is a FHE Model",
+    )
+    app.launch()