import os
import joblib
import pefile
import numpy as np
import pandas as pd
import gradio as gr
import hashlib


# Charger le modèle pré-entraîné
try:
    model = joblib.load('random_forest_model.pkl')
except Exception as e:
    print(f"Erreur de chargement du modèle : {e}")
    model = None

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 = {
            # Attributs PE standard
            '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
        }
        
        """## 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)}"

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

        # 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)  # Rend l'interface accessible publiquement