Update app.py

app.py

@@ -1,8 +1,4 @@
 import gradio as gr
-import time
-import os
-from pathlib import Path
-
 import subprocess
 from concrete.ml.deployment import FHEModelClient
 from requests import head
@@ -12,13 +8,13 @@ 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
 import pdb
+
 # This repository's directory
 REPO_DIR = Path(__file__).parent
 subprocess.Popen(["uvicorn", "server:app"], cwd=REPO_DIR)
@@ -72,7 +68,7 @@ def clean_tmp_directory():
                 file.unlink()
 
 
-def keygen(eval_key):
+def keygen(eval_key, user_id):
     # Clean tmp directory if needed
     clean_tmp_directory()
 
@@ -89,8 +85,8 @@ def keygen(eval_key):
 
     numpy.save(f"tmp/tmp_evaluation_key_{user_id}.npy", evaluation_key)
 
-    eval_key = [list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT], user_id]
-    return eval_key[0]
+    eval_key = list(evaluation_key)[:ENCRYPTED_DATA_BROWSER_LIMIT]
+    return eval_key, user_id
 
 def encode_quantize(test_file, eval_key, encodings):
     ugly = ['Machine', 'SizeOfOptionalHeader', 'Characteristics',
@@ -128,8 +124,7 @@ def encode_quantize(test_file, eval_key, encodings):
 
     return encodings
 
-def encrypt_encoded_quantize(encodings, eval_key):
-    global user_id
+def encrypt_encoded_quantize(encodings, user_id, eval_key):
     fhe_api = FHEModelClient(f"fhe_model", f".fhe_keys/{user_id}")
     fhe_api.load()
     
@@ -147,8 +142,7 @@ def encrypt_encoded_quantize(encodings, eval_key):
     encrypted_quantized_encoding_shorten_hex = "".join(f"{i:02x}" for i in encrypted_quantized_encoding_shorten)
     return (encodings,quantized_encodings,encrypted_quantized_encoding_shorten_hex)
 
-def run_fhe():
-    global user_id
+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)
 
@@ -180,8 +174,7 @@ def run_fhe():
     return encrypted_prediction_shorten_hex
 
 
-def decrypt_prediction():
-    global user_id
+def decrypt_prediction(user_id):
     encoded_data_path = Path(f"tmp/tmp_encrypted_prediction_{user_id}.npy")
 
     # Read encrypted_prediction from the file
@@ -351,11 +344,11 @@ if __name__ == "__main__":
             "The encrypted scan result is sent back to the client, who can finally decrypt it with their private key. Only the client is aware of the original file content and the scan result."
         )
         b_decrypt_result = gr.Button("🔍 Decrypt scan result")
+        user_id_input = gr.Number(visible=False)
         scan_result = gr.Textbox(label="Scan Result:")
-        
         eval_key_input = gr.Textbox(value=eval_key, visible=False)
         # Button for key generation
-        b_gen_key.click(fn=keygen, inputs=[eval_key_input], outputs=[evaluation_key])
+        b_gen_key.click(fn=keygen, inputs=[eval_key_input, user_id_input], outputs=[evaluation_key, user_id_input])
 
         encodings_input = gr.Textbox(value=encodings, visible=False)
         # Button to extract vector
@@ -368,15 +361,15 @@ if __name__ == "__main__":
         # Button to encrypt file
         b_encrypt_file.click(
             fn=encrypt_encoded_quantize,
-            inputs=[extracted_vector, eval_key_input],
+            inputs=[extracted_vector, user_id_input, eval_key_input],
             outputs=[encrypted_file],
         )
 
         # Button to run FHE-based malware scan
-        b_run_fhe_scan.click(fn=run_fhe, inputs=[], outputs=[encrypted_scan_result])
+        b_run_fhe_scan.click(fn=run_fhe, inputs=[user_id_input], outputs=[encrypted_scan_result])
 
         # Button to decrypt the scan result
-        b_decrypt_result.click(fn=decrypt_prediction, inputs=[], outputs=[scan_result])
+        b_decrypt_result.click(fn=decrypt_prediction, inputs=[user_id_input], outputs=[scan_result])
 
         gr.Markdown(
             "ClairVault is built using advanced Fully Homomorphic Encryption techniques to ensure your data remains private and secure throughout the entire malware scanning process."