app.py

import gradio as gr
from transformers import pipeline

# --- 1. Load Model AI ---
model_path = "elftsdmr/malware-url-detect"
try:
    classifier = pipeline("text-classification", model=model_path, top_k=None)
except Exception as e:
    classifier = None
    print(f"Error loading model: {e}")

# --- 2. Definisi Rule-Based (Daftar Hitam Kata Kunci) ---
# Ini adalah jaring pengaman jika AI gagal mendeteksi
SUSPICIOUS_KEYWORDS = [
    "bitcoin", "mining", "crypto", "wallet", "bonus", "claim", 
    "free-gift", "login-verify", "secure-account", "update-payment",
    "paypal", "bank", "prize", "giveaway"
]

def check_keywords(url):
    """Cek apakah URL mengandung kata kunci berbahaya"""
    found_keywords = []
    for keyword in SUSPICIOUS_KEYWORDS:
        if keyword in url.lower():
            found_keywords.append(keyword)
    return found_keywords

# --- 3. Fungsi Logika Utama (Hybrid) ---
def analyze_url(target_url):
    if classifier is None:
        return "Model gagal dimuat.", {}
    if not target_url:
        return "Mohon masukkan URL.", {}
    
    try:
        # A. Analisis AI
        results = classifier(target_url)[0]
        
        # Format data untuk grafik
        chart_data = {}
        ai_verdict = "AMAN"
        top_score = 0
        
        for res in results:
            label = res['label']
            score = res['score']
            
            if label == "LABEL_1" or "MAL" in label.upper():
                readable_label = "AI: BERBAHAYA"
                if score > 0.5: ai_verdict = "BAHAYA"
            else:
                readable_label = "AI: AMAN"
            
            chart_data[readable_label] = score
            if score > top_score: top_score = score

        # B. Analisis Keyword (Rule-Based)
        found_keywords = check_keywords(target_url)
        
        # C. KEPUTUSAN AKHIR (Logika Gabungan)
        
        # Skenario 1: AI bilang Bahaya
        if ai_verdict == "BAHAYA":
            final_status = "⚠️ BERBAHAYA (Terdeteksi AI)"
            detail_msg = f"AI sangat yakin ({top_score:.1%}) ini berbahaya."
            
        # Skenario 2: AI bilang Aman, TAPI ada Keyword mencurigakan (Kasus Bitcoin tadi)
        elif ai_verdict == "AMAN" and found_keywords:
            final_status = "⚠️ MENCURIGAKAN (Terdeteksi Keyword)"
            detail_msg = (f"Meskipun AI memberi skor aman, sistem mendeteksi kata kunci sensitif: {found_keywords}.\n"
                          f"Sistem melakukan 'Override' demi keamanan.")
            # Kita paksa grafik tampil seimbang agar user waspada
            chart_data["Rule-Based Detection"] = 0.99 
            
        # Skenario 3: AI bilang Aman dan tidak ada Keyword
        else:
            final_status = "✅ AMAN"
            detail_msg = f"Tidak ditemukan indikasi berbahaya oleh AI maupun sistem keyword."

        # Format Teks Output
        text_output = (f"STATUS: {final_status}\n"
                       f"--------------------------------\n"
                       f"{detail_msg}\n\n"
                       f"Confidence AI: {top_score:.1%}")
            
        return text_output, chart_data
        
    except Exception as e:
        return f"Error: {e}", {}

# --- 4. UI Gradio ---
with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🛡️ Hybrid AI Network Security
        ### Kombinasi Deep Learning + Rule Based System
        """
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="URL Target", 
                placeholder="Masukkan URL...",
                lines=3
            )
            submit_btn = gr.Button("🚀 Scan Hybrid", variant="primary")
            gr.Markdown("ℹ️ **Info:** Sistem ini menggunakan 2 lapis keamanan (AI + Daftar Kata Kunci).")

        with gr.Column(scale=1):
            output_text = gr.Textbox(label="Hasil Analisis", lines=8)
            output_chart = gr.Label(label="Indikator Deteksi")
            
    submit_btn.click(
        fn=analyze_url, 
        inputs=input_text, 
        outputs=[output_text, output_chart]
    )
    
    gr.Markdown("---")
    gr.Markdown("*Powered by Hugging Face Spaces*")

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)