Spaces:
Sleeping
Sleeping
File size: 4,087 Bytes
4d9464e ba80d59 6e2bab3 4d9464e 6e2bab3 4d9464e 6e2bab3 4d9464e 6e2bab3 4d9464e 6e2bab3 4d9464e 6e2bab3 4d9464e 6e2bab3 4d9464e 6e2bab3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | import os
import joblib
import pefile
import numpy as np
import pandas as pd
import gradio as gr
import hashlib
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}")
# 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()
with open(file_path, "rb") as f:
for byte_block in iter(lambda: f.read(4096), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
def extract_pe_attributes(file_path):
"""Extraction avancée des attributs du fichier PE."""
try:
pe = pefile.PE(file_path)
attributes = {
'AddressOfEntryPoint': pe.OPTIONAL_HEADER.AddressOfEntryPoint,
'MajorLinkerVersion': pe.OPTIONAL_HEADER.MajorLinkerVersion,
'MajorImageVersion': pe.OPTIONAL_HEADER.MajorImageVersion,
'MajorOperatingSystemVersion': pe.OPTIONAL_HEADER.MajorOperatingSystemVersion,
'DllCharacteristics': pe.OPTIONAL_HEADER.DllCharacteristics,
'SizeOfStackReserve': pe.OPTIONAL_HEADER.SizeOfStackReserve,
'NumberOfSections': pe.FILE_HEADER.NumberOfSections,
'ResourceSize': pe.OPTIONAL_HEADER.DATA_DIRECTORY[2].Size
}
return attributes
except Exception as e:
print(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."""
if model is None:
return "Erreur : Modèle non chargé"
try:
# Extraire les attributs du fichier
attributes = extract_pe_attributes(file.name)
if "Erreur" in attributes:
return attributes["Erreur"]
# Convertir en DataFrame
df = pd.DataFrame([attributes])
# Prédiction
prediction = model.predict(df)
proba = model.predict_proba(df)[0]
# Résultat avec probabilité
if prediction[0] == 1:
return f"🚨 MALWARE (Probabilité: {proba[1] * 100:.2f}%)"
else:
return f"✅ Fichier Légitime (Probabilité: {proba[0] * 100:.2f}%)"
except Exception as e:
return f"Erreur d'analyse : {str(e)}"
# Interface Gradio
demo = gr.Interface(
fn=predict_malware,
inputs=gr.File(file_types=['.exe', '.dll', '.sys'], label="Télécharger un fichier exécutable"),
outputs="text",
title="🛡️ Détecteur de Malwares",
theme='huggingface' # Thème moderne
)
if __name__ == "__main__":
demo.launch(share=True)
|