Upload 3 files

PE_main.py

@@ -76,9 +76,9 @@ def get_version_info(pe):
     return res
 
 #extract the info for a given file using pefile
-def extract_infos(file):
+def extract_infos(fpath):
     res = {}
-    pe = pefile.PE(file)
+    pe = pefile.PE(fpath)
     res['Machine'] = pe.FILE_HEADER.Machine
     res['SizeOfOptionalHeader'] = pe.FILE_HEADER.SizeOfOptionalHeader
     res['Characteristics'] = pe.FILE_HEADER.Characteristics
@@ -182,3 +182,20 @@ 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]))

fhe_utils.py

@@ -0,0 +1,223 @@
+import sys
+import os
+import pdb
+import numpy as np
+import random
+import json
+import shutil
+import time
+from scipy.stats import pearsonr
+from sklearn.model_selection import KFold
+import matplotlib.pyplot as plt
+import xgboost as xgb
+from tqdm import tqdm
+random.seed(42)
+
+
+import gzip
+import numpy as np
+import pandas as pd
+import requests
+from io import BytesIO
+from concrete.ml.deployment import FHEModelClient, FHEModelDev, FHEModelServer
+
+from concrete.ml.sklearn import DecisionTreeClassifier as DecisionTreeClassifierZAMA
+
+from concrete.ml.sklearn import LinearSVC as LinearSVCZAMA
+
+from sklearn.svm import LinearSVR as LinearSVR
+import time
+from shutil import copyfile
+from tempfile import TemporaryDirectory
+import pickle
+import os
+import time
+import numpy as np
+
+
+def convert_numpy(obj):
+    if isinstance(obj, np.integer):
+        return int(obj)
+    elif isinstance(obj, np.floating):
+        return float(obj)
+    elif isinstance(obj, np.ndarray):
+        return obj.tolist()
+    else:
+        return obj
+
+
+class OnDiskNetwork:
+    """Simulate a network on disk."""
+
+    def __init__(self):
+        # Create 3 temporary folder for server, client and dev with tempfile
+        self.server_dir = TemporaryDirectory()
+        self.client_dir = TemporaryDirectory()
+        self.dev_dir = TemporaryDirectory()
+
+    def client_send_evaluation_key_to_server(self, serialized_evaluation_keys):
+        """Send the public key to the server."""
+        with open(self.server_dir.name + "/serialized_evaluation_keys.ekl", "wb") as f:
+            f.write(serialized_evaluation_keys)
+
+    def client_send_input_to_server_for_prediction(self, encrypted_input):
+        """Send the input to the server and execute on the server in FHE."""
+        with open(self.server_dir.name + "/serialized_evaluation_keys.ekl", "rb") as f:
+            serialized_evaluation_keys = f.read()
+        time_begin = time.time()
+        encrypted_prediction = FHEModelServer(self.server_dir.name).run(
+            encrypted_input, serialized_evaluation_keys
+        )
+        time_end = time.time()
+        with open(self.server_dir.name + "/encrypted_prediction.enc", "wb") as f:
+            f.write(encrypted_prediction)
+        return time_end - time_begin
+
+    def dev_send_model_to_server(self):
+        """Send the model to the server."""
+        copyfile(
+            self.dev_dir.name + "/server.zip", self.server_dir.name + "/server.zip"
+        )
+
+    def server_send_encrypted_prediction_to_client(self):
+        """Send the encrypted prediction to the client."""
+        with open(self.server_dir.name + "/encrypted_prediction.enc", "rb") as f:
+            encrypted_prediction = f.read()
+        return encrypted_prediction
+
+    def dev_send_clientspecs_and_modelspecs_to_client(self):
+        """Send the clientspecs and evaluation key to the client."""
+        copyfile(
+            self.dev_dir.name + "/client.zip", self.client_dir.name + "/client.zip"
+        )
+
+    def cleanup(self):
+        """Clean up the temporary folders."""
+        self.server_dir.cleanup()
+        self.client_dir.cleanup()
+        self.dev_dir.cleanup()
+
+
+def generate_fingerprint(smiles, radius=2, bits=512):
+    mol = Chem.MolFromSmiles(smiles)
+    if mol is None:
+        return np.nan
+
+    fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=radius, nBits=bits)
+
+    return np.array(fp)
+
+
+def train_xgb_regressor(X_train, y_train, param_grid=None, verbose=10):
+    if param_grid is None:
+        param_grid = {
+            "max_depth": [3, 6],
+            "learning_rate": [0.01, 0.1, 0.2],
+            "n_estimators": [20],
+            "colsample_bytree": [0.3, 0.7],
+        }
+
+    xgb_regressor = xgb.XGBRegressor(objective="reg:squarederror")
+
+    kfold = KFold(n_splits=5, shuffle=True, random_state=42)
+    grid_search = GridSearchCV(
+        estimator=xgb_regressor,
+        param_grid=param_grid,
+        cv=kfold,
+        verbose=verbose,
+        n_jobs=-1,
+    )
+
+    grid_search.fit(X_train, y_train)
+    return (
+        grid_search.best_params_,
+        grid_search.best_score_,
+        grid_search.best_estimator_,
+    )
+
+
+def evaluate_model(model, X_test, y_test):
+    y_pred = model.predict(X_test)
+    pearsonr_score = pearsonr(y_test, y_pred).statistic
+    return pearsonr_score
+
+
+def setup_network(model_dev):
+    network = OnDiskNetwork()
+    fhemodel_dev = FHEModelDev(network.dev_dir.name, model_dev)
+    fhemodel_dev.save(via_mlir=True)
+    return network, fhemodel_dev
+
+
+def copy_directory(source, destination="deployment"):
+    try:
+        # Check if the source directory exists
+        if not os.path.exists(source):
+            return False, "Source directory does not exist."
+
+        # Check if the destination directory exists
+        if not os.path.exists(destination):
+            os.makedirs(destination)
+
+        # Copy each item in the source directory
+        for item in os.listdir(source):
+            s = os.path.join(source, item)
+            d = os.path.join(destination, item)
+            if os.path.isdir(s):
+                shutil.copytree(
+                    s, d, dirs_exist_ok=True
+                )  # dirs_exist_ok is available from Python 3.8
+            else:
+                shutil.copy2(s, d)
+
+        return True, None
+
+    except Exception as e:
+        return False, str(e)
+
+
+def client_server_interaction(network, fhemodel_client, X_client):
+    decrypted_predictions = []
+    execution_time = []
+    for i in tqdm(range(X_client.shape[0])):
+        clear_input = X_client[[i], :]
+        encrypted_input = fhemodel_client.quantize_encrypt_serialize(clear_input)
+        execution_time.append(
+            network.client_send_input_to_server_for_prediction(encrypted_input)
+        )
+        encrypted_prediction = network.server_send_encrypted_prediction_to_client()
+        decrypted_prediction = fhemodel_client.deserialize_decrypt_dequantize(
+            encrypted_prediction
+        )[0]
+        decrypted_predictions.append(decrypted_prediction)
+        #pdb.set_trace()
+    return decrypted_predictions, execution_time
+
+
+def train_zama(X_train, y_train):
+
+    model_dev = LinearSVCZAMA()
+    # LinearSVCZAMA()
+    # DecisionTreeClassifierZAMA()
+
+    print("Training Zama model...")
+    model_dev.fit(X_train, y_train)
+    print("compiling model...")
+    model_dev.compile(X_train)
+    print("done")
+
+    return model_dev
+
+
+def time_prediction(model, X_sample):
+    time_begin = time.time()
+    y_pred_fhe = model.predict(X_sample, fhe="execute")
+    time_end = time.time()
+    return time_end - time_begin
+
+
+def setup_client(network, key_dir):
+    fhemodel_client = FHEModelClient(network.client_dir.name, key_dir=key_dir)
+    fhemodel_client.generate_private_and_evaluation_keys()
+    serialized_evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()
+    return fhemodel_client, serialized_evaluation_keys

train_FHE.py

@@ -0,0 +1,171 @@
+import os
+import pandas as pd
+import numpy
+import pickle
+import pefile
+import sklearn.ensemble as ek
+from sklearn.feature_selection import SelectFromModel
+import joblib
+from sklearn.tree import DecisionTreeClassifier
+from sklearn.metrics import confusion_matrix
+from sklearn import svm
+import sklearn.metrics as metrics
+from sklearn.model_selection import train_test_split
+import pdb
+
+from tqdm import tqdm
+
+dataset = pd.read_csv("data.csv", sep="|")
+
+
+# Feature
+X = dataset.drop(
+    ["Name", "md5", "legitimate"], axis=1
+).values  # Droping this because classification model will not accept object type elements (float and int only)
+# Target variable
+
+ugly = [
+    "Machine",
+    "SizeOfOptionalHeader",
+    "Characteristics",
+    "MajorLinkerVersion",
+    "MinorLinkerVersion",
+    "SizeOfCode",
+    "SizeOfInitializedData",
+    "SizeOfUninitializedData",
+    "AddressOfEntryPoint",
+    "BaseOfCode",
+    "BaseOfData",
+    "ImageBase",
+    "SectionAlignment",
+    "FileAlignment",
+    "MajorOperatingSystemVersion",
+    "MinorOperatingSystemVersion",
+    "MajorImageVersion",
+    "MinorImageVersion",
+    "MajorSubsystemVersion",
+    "MinorSubsystemVersion",
+    "SizeOfImage",
+    "SizeOfHeaders",
+    "CheckSum",
+    "Subsystem",
+    "DllCharacteristics",
+    "SizeOfStackReserve",
+    "SizeOfStackCommit",
+    "SizeOfHeapReserve",
+    "SizeOfHeapCommit",
+    "LoaderFlags",
+    "NumberOfRvaAndSizes",
+    "SectionsNb",
+    "SectionsMeanEntropy",
+    "SectionsMinEntropy",
+    "SectionsMaxEntropy",
+    "SectionsMeanRawsize",
+    "SectionsMinRawsize",
+    #"SectionsMaxRawsize",
+    "SectionsMeanVirtualsize",
+    "SectionsMinVirtualsize",
+    "SectionMaxVirtualsize",
+    "ImportsNbDLL",
+    "ImportsNb",
+    "ImportsNbOrdinal",
+    "ExportNb",
+    "ResourcesNb",
+    "ResourcesMeanEntropy",
+    "ResourcesMinEntropy",
+    "ResourcesMaxEntropy",
+    "ResourcesMeanSize",
+    "ResourcesMinSize",
+    "ResourcesMaxSize",
+    "LoadConfigurationSize",
+    "VersionInformationSize",
+]
+
+X = dataset[ugly].values
+
+y = dataset["legitimate"].values
+
+
+extratrees = ek.ExtraTreesClassifier().fit(X[:1000], y[:1000])
+model = SelectFromModel(extratrees, prefit=True)
+X_new = model.transform(X)
+nbfeatures = X_new.shape[1]
+
+
+# splitting the data (70% - training and 30% - testing)
+
+
+X_train, X_test, y_train, y_test = train_test_split(
+    X_new, y, test_size=0.29, stratify=y
+)
+
+
+features = []
+index = numpy.argsort(extratrees.feature_importances_)[::-1][:nbfeatures]
+
+
+for f in range(nbfeatures):
+    print(
+        "%d. feature %s (%f)"
+        % (
+            f + 1,
+            dataset.columns[2 + index[f]],
+            extratrees.feature_importances_[index[f]],
+        )
+    )
+    features.append(dataset.columns[2 + f])
+
+
+model = {
+    "DecisionTree": DecisionTreeClassifier(max_depth=10),
+    "RandomForest": ek.RandomForestClassifier(n_estimators=50),
+}
+
+
+results = {}
+for algo in model:
+    clf = model[algo]
+    clf.fit(X_train, y_train)
+    score = clf.score(X_test, y_test)
+    print("%s : %s " % (algo, score))
+    results[algo] = score
+
+
+winner = max(results, key=results.get)  # Selecting the classifier with good result
+print("Using", winner, "for classification, with", len(features), "features.")
+
+
+joblib.dump(model[winner], "classifier.pkl")
+open("features.pkl", "wb").write(pickle.dumps(features))
+
+
+from fhe_utils import (
+    client_server_interaction, train_zama,
+    setup_network,
+    copy_directory,
+    setup_client,
+)
+
+model_dev_fhe = train_zama(X_train, y_train)
+#pdb.set_trace()
+network, _ = setup_network(model_dev_fhe)
+copied, error_message = copy_directory(network.dev_dir.name, destination="fhe_model")
+
+
+if not copied:
+    print(f"Error copying directory: {error_message}")
+
+
+network.dev_send_model_to_server()
+network.dev_send_clientspecs_and_modelspecs_to_client()
+
+fhemodel_client, serialized_evaluation_keys = setup_client(
+    network, network.client_dir.name
+)
+print(f"Evaluation keys size: {len(serialized_evaluation_keys)} B")
+
+network.client_send_evaluation_key_to_server(serialized_evaluation_keys)
+
+
+
+decrypted_predictions, execution_time  = client_server_interaction(network, fhemodel_client, X_test[:100])