changes

PE_main.py

@@ -78,7 +78,7 @@ def get_version_info(pe):
 #extract the info for a given file using pefile
 def extract_infos(file):
     res = {}
-    pe = pefile.PE(data=file)
+    pe = pefile.PE(file)
     res['Machine'] = pe.FILE_HEADER.Machine
     res['SizeOfOptionalHeader'] = pe.FILE_HEADER.SizeOfOptionalHeader
     res['Characteristics'] = pe.FILE_HEADER.Characteristics
@@ -182,20 +182,3 @@ def extract_infos(file):
         res['VersionInformationSize'] = 0
     return res
 
-
-if __name__ == '__main__':
-    
-    #Loading the classifier.pkl and features.pkl
-    clf = joblib.load('Classifier/classifier.pkl')
-    features = pickle.loads(open(os.path.join('Classifier/features.pkl'),'rb').read())
-    
-    #extracting features from the PE file mentioned in the argument 
-    data = extract_infos(sys.argv[1])
-    
-    #matching it with the features saved in features.pkl
-    pe_features = list(map(lambda x:data[x], features))
-    print("Features used for classification: ", pe_features)
-    
-    #prediciting if the PE is malicious or not based on the extracted features
-    res= clf.predict([pe_features])[0]
-    print ('The file %s is %s' % (os.path.basename(sys.argv[1]),['malicious', 'legitimate'][res]))

app.py

@@ -117,7 +117,7 @@ def encode_quantize(test_file, eval_key, encodings):
     # features = pickle.loads(open(os.path.join("features.pkl"), "rb").read())
 
     
-    encodings = extract_infos(test_file)
+    encodings = extract_infos(test_file.name)
 
     encodings = list(map(lambda x: encodings[x], ugly))
 
@@ -302,7 +302,7 @@ if __name__ == "__main__":
         gr.Markdown(
             "Upload a file you want to scan for malware. ClairVault will encrypt it locally before sending it to the cloud."
         )
-        file_input = gr.File(label="Upload a file:")
+        file_input = gr.File(label="Upload a file:", live=True, file_count = "single") 
 
         gr.Markdown("# 📥 Step 3: Extract executable file features")
 
@@ -348,7 +348,7 @@ if __name__ == "__main__":
         scan_result = gr.Textbox(label="Scan Result:")
 
         # Button for key generation
-        b_gen_key.click(keygen(eval_key), inputs=[], outputs=[])
+        b_gen_key.click(keygen(eval_key), inputs=[], outputs=[evaluation_key])
 
         # Button to extract vector
         b_extract.click(