Update app.py

app.py

@@ -1,4 +1,3 @@
-
 import os
 import joblib
 import pefile
@@ -6,18 +5,57 @@ import numpy as np
 import pandas as pd
 import gradio as gr
 import hashlib
-import pickle
 import traceback
+from sklearn.ensemble import RandomForestClassifier
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score, recall_score
+
+# Chemin vers le modèle sauvegardé
+MODEL_PATH = 'random_forest_model.pkl'
+
+def train_and_save_model():
+    """Entraîner et sauvegarder le modèle si nécessaire."""
+    print("Aucun modèle trouvé. Entraînement en cours...")
+    # Chargement des données
+    data = pd.read_csv("DatasetmalwareExtrait.csv")
+
+    # Traitement des données
+    X = data.drop(['legitimate'], axis=1)
+    y = data['legitimate']
+
+    # Entraînement du modèle
+    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=42)
+    model = RandomForestClassifier(
+        n_estimators=196,
+        random_state=42,
+        criterion="gini",
+        max_depth=25,
+        min_samples_split=4,
+        min_samples_leaf=1
+    )
+    model.fit(X_train, y_train)
+
+    # Évaluation du modèle
+    y_pred = model.predict(X_test)
+    accuracy = accuracy_score(y_test, y_pred)
+    recall = recall_score(y_test, y_pred, average='weighted')
 
+    print(f"Précision du modèle supervisé : {accuracy:.3f}")
+    print(f"Rappel du modèle supervisé : {recall:.3f}")
 
-# Charger le modèle pré-entraîné
-model = None
-try:
-        model = joblib.load('random_forest_model.pkl')
-except Exception as e:
-    print(f"Erreur de chargement du modèle : {e}")
-    traceback.print_exc()  # Affiche la trace de l'erreur
+    # Sauvegarde du modèle
+    joblib.dump(model, MODEL_PATH)
+    print(f"Modèle sauvegardé sous : {MODEL_PATH}")
+    return model
 
+# Chargement ou entraînement du modèle
+if os.path.exists(MODEL_PATH):
+    print("Chargement du modèle existant...")
+    model = joblib.load(MODEL_PATH)
+else:
+    model = train_and_save_model()
+
+# Fonctions utilitaires
 def calculate_file_hash(file_path):
     """Calculer le hash SHA-256 du fichier."""
     sha256_hash = hashlib.sha256()
@@ -30,9 +68,7 @@ def extract_pe_attributes(file_path):
     """Extraction avancée des attributs du fichier PE."""
     try:
         pe = pefile.PE(file_path)
-
         attributes = {
-            # Attributs PE standard
             'AddressOfEntryPoint': pe.OPTIONAL_HEADER.AddressOfEntryPoint,
             'MajorLinkerVersion': pe.OPTIONAL_HEADER.MajorLinkerVersion,
             'MajorImageVersion': pe.OPTIONAL_HEADER.MajorImageVersion,
@@ -40,26 +76,12 @@ def extract_pe_attributes(file_path):
             'DllCharacteristics': pe.OPTIONAL_HEADER.DllCharacteristics,
             'SizeOfStackReserve': pe.OPTIONAL_HEADER.SizeOfStackReserve,
             'NumberOfSections': pe.FILE_HEADER.NumberOfSections,
-             'ResourceSize':pe.OPTIONAL_HEADER.DATA_DIRECTORY[2].Size
+            'ResourceSize': pe.OPTIONAL_HEADER.DATA_DIRECTORY[2].Size
         }
-        
-        """## Ressources
-        data_directory_entries = pe.OPTIONAL_HEADER.DATA_DIRECTORY
-        # Parcourir la liste pour trouver l'entrée du répertoire des ressources
-        for entry in data_directory_entries:
-            if entry.name == "IMAGE_DIRECTORY_ENTRY_RESOURCE":
-                resource_size = entry.Size
-                attributes['ResourceSize'] = resource_size
-                break
-        else:
-            attributes['ResourceSize'] = 0"""
-            
-
-
         return attributes
     except Exception as e:
         print(f"Erreur de traitement du fichier {file_path}: {str(e)}")
-        return f"Erreur de traitement du fichier {file_path}: {str(e)}"
+        return {"Erreur": str(e)}
 
 def predict_malware(file):
     """Prédiction de malware avec gestion d'erreurs."""
@@ -70,7 +92,7 @@ def predict_malware(file):
         # Extraire les attributs du fichier
         attributes = extract_pe_attributes(file.name)
         if "Erreur" in attributes:
-            return attributes
+            return attributes["Erreur"]
 
         # Convertir en DataFrame
         df = pd.DataFrame([attributes])
@@ -97,4 +119,4 @@ demo = gr.Interface(
 )
 
 if __name__ == "__main__":
-    demo.launch(share=True)  # Rend l'interface accessible publiquement
+    demo.launch(share=True)