# scan.py
import disnake
from disnake.ext import commands
import hashlib
import os
import json
import re

MALWARE_DB = "malware.json"
SUSPICIOUS_EXTENSIONS = [".exe", ".scr", ".bat", ".js", ".vbs", ".cmd"]
MALWARE_PATTERNS = [
    rb"powershell\s+-nop",
    rb"cmd\s+/c",
    rb"eval\(",
    rb"base64_decode\(",
    rb"CreateObject\(\"WScript\.Shell\"\)"
]

def load_malware_db():
    if not os.path.isfile(MALWARE_DB):
        with open(MALWARE_DB, "w") as f:
            json.dump([], f)
        return []
    with open(MALWARE_DB, "r") as f:
        return json.load(f)

def get_file_hash(file_path):
    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 scan_file(file_path):
    if not os.path.isfile(file_path):
        return {"error": "Fichier introuvable."}

    file_hash = get_file_hash(file_path)
    malware_db = load_malware_db()

    if file_hash in malware_db:
        return {"hash": file_hash, "malicious": True, "reason": "Hash connu dans la base malveillante."}

    _, ext = os.path.splitext(file_path)
    if ext.lower() in SUSPICIOUS_EXTENSIONS:
        return {"hash": file_hash, "malicious": True, "reason": f"Extension suspecte : {ext}"}

    with open(file_path, "rb") as f:
        content = f.read()
        for pattern in MALWARE_PATTERNS:
            if re.search(pattern, content):
                return {"hash": file_hash, "malicious": True, "reason": f"Pattern suspect détecté : {pattern.decode(errors='ignore')}"}

    return {"hash": file_hash, "malicious": False, "reason": "Fichier sain."}


class Scan(commands.Cog):
    def __init__(self, bot: commands.Bot):
        self.bot = bot

    @commands.slash_command(name="scan", description="Scanner un fichier pour détecter un malware.")
    async def scan(self, interaction: disnake.ApplicationCommandInteraction, file: disnake.Attachment):
        await interaction.response.defer()  # Permet de prendre un peu de temps pour le scan

        # Téléchargement temporaire du fichier
        tmp_path = f"tmp_{file.filename}"
        await file.save(tmp_path)

        result = scan_file(tmp_path)
        os.remove(tmp_path)  # Supprime le fichier temporaire

        if result.get("error"):
            await interaction.followup.send(f"❌ Erreur : {result['error']}")
        elif result.get("malicious"):
            await interaction.followup.send(f"⚠️ Fichier malveillant détecté !\nRaison : {result['reason']}\nHash : `{result['hash']}`")
        else:
            await interaction.followup.send(f"✅ Fichier sain.\nHash : `{result['hash']}`")


async def setup(bot: commands.Bot):
    await bot.add_cog(Scan(bot))