app.py

# GRADIO APPLICATION FOR HUGGING FACE SPACES
# This version loads the master column list to ensure data alignment.

import os
import joblib
import numpy as np
import pandas as pd
import tensorflow as tf
import gradio as gr
from tensorflow.keras.models import load_model
from sklearn.preprocessing import LabelEncoder

# --- Model & Scaler Configuration ---
H5_MODEL_FILE = "intrusion_detector_model.h5"
SCALER_FILE_NAME = "scaler.pkl"
MASTER_COLS_FILE = "master_columns.pkl" # NEW: File containing all OHE column names
PREDICTION_THRESHOLD = 0.40 
# FEATURE_COUNT is now dynamic based on MASTER_COLS_FILE loading

# Raw features from the form (used to create the initial dictionary)
FEATURE_NAMES = [
    'duration', 'protocol_type', 'service', 'flag', 'src_bytes', 'dst_bytes', 'land',
    'wrong_fragment', 'urgent', 'hot', 'num_failed_logins', 'logged_in', 'num_compromised', 
    'root_shell', 'su_attempted', 'num_root', 'num_file_creations', 'num_shells', 'num_access_files', 
    'num_outbound_cmds', 'is_host_login', 'is_guest_login', 'count', 'srv_count', 'serror_rate', 
    'srv_serror_rate', 'rerror_rate', 'srv_rerror_rate', 'same_srv_rate', 'diff_srv_rate', 
    'srv_diff_host_rate', 'dst_host_count', 'dst_host_srv_count', 'dst_host_same_srv_rate', 
    'dst_host_diff_srv_rate', 'dst_host_same_src_port_rate', 'dst_host_srv_diff_host_rate', 
    'dst_host_serror_rate', 'dst_host_srv_serror_rate', 'dst_host_rerror_rate', 
    'dst_host_srv_rerror_rate'
]

# Lists for Gradio Dropdowns (must be comprehensive)
SERVICES = [
    'http', 'smtp', 'ftp_data', 'private', 'ecr_i', 'other', 'domain_u', 
    'finger', 'telnet', 'ftp', 'pop_3', 'courier', 'eco_i', 'imap4', 
    'domain_n', 'auth', 'time', 'shell', 'login', 'hostnames', 'ntp_service', 
    'echo', 'discard', 'systat', 'ctf', 'ssh', 'iso_tsap', 'whois', 'remote_job', 
    'sunrpc', 'rje', 'gopher', 'netbios_ssn', 'pm_srv', 'mtp', 'exec', 'klogin', 
    'kshell', 'daytime', 'message', 'icmp', 'netstat', 'Z39_50', 'bgp', 'nnsp', 
    'ctinrp', 'IRC', 'urp_i', 'pop_2', 'aol', 'rev_telnet', 'tftp_u'
]
FLAGS = ['SF', 'S0', 'REJ', 'RSTO', 'SH', 'S1', 'S2', 'RSTOS0', 'S3', 'OTH', 'RSTR']
PROTOCOLS = ['tcp', 'udp', 'icmp']
CATEGORICAL_COLS = ['protocol_type', 'service', 'flag']

# Global artifacts
model = None
scaler = None
master_columns = None
FEATURE_COUNT = 0
MAPPING = {'normal': 0, 'anomaly': 1}
label_encoder = LabelEncoder()
label_encoder.fit(list(MAPPING.keys()))


# --- Model Loading and Initialization ---

def load_artifacts():
    """Loads the trained model, scaler, and master column list globally."""
    global model, scaler, master_columns, FEATURE_COUNT
    
    print("--- Starting Artifact Loading ---")

    files_to_check = [SCALER_FILE_NAME, H5_MODEL_FILE, MASTER_COLS_FILE]
    missing = [f for f in files_to_check if not os.path.exists(f)]
    
    if missing:
        print(f"CRITICAL ERROR: The following required files are missing: {', '.join(missing)}")
        print("Please ensure these three files from the retrained OHE model are uploaded.")
        return False
    
    try:
        scaler = joblib.load(SCALER_FILE_NAME)
        master_columns = joblib.load(MASTER_COLS_FILE)
        model = load_model(H5_MODEL_FILE, compile=False)
        FEATURE_COUNT = len(master_columns)

        print(f"✓ Scaler, Master Columns, and Model loaded.")
        print(f"Model expects {FEATURE_COUNT} features.")
        
    except Exception as e:
        print(f"Error loading artifacts: {type(e).__name__}: {e}")
        return False

    print("--- Artifact Loading Complete ---")
    return True

# Load artifacts on startup
if not load_artifacts():
    pass


# --- Prediction Function ---

def predict_intrusion(*inputs):
    """
    Takes 41 raw network features, preprocesses them using OHE and scaling, 
    and makes a prediction, aligning the columns using master_columns.
    """
    if model is None or scaler is None or master_columns is None:
        return "<h2 style='color: red; text-align: center;'>FATAL ERROR: Artifacts Not Loaded. See Logs.</h2>", "N/A"

    try:
        # 1. Create a dictionary from the inputs
        raw_input_dict = {FEATURE_NAMES[i]: [inputs[i]] for i in range(len(FEATURE_NAMES))}
        df = pd.DataFrame(raw_input_dict)
        
        # 2. Apply One-Hot Encoding (OHE) for categorical features
        # This converts text columns into many 0/1 columns.
        df_ohe = pd.get_dummies(df, columns=CATEGORICAL_COLS, drop_first=False)
        
        # 3. Align columns to match training data (CRITICAL FIX APPLIED)
        # Convert the single row OHE DataFrame to a Series, reindex it against the 
        # master list, fill missing values with 0, and convert it back to a single row DataFrame.
        df_series = df_ohe.iloc[0].reindex(master_columns, fill_value=0)
        df_aligned = pd.DataFrame(df_series).T # Transpose back to a single row DataFrame
        
        # --- CRITICAL DEBUGGING PRINTS ---
        print(f"Debug: DataFrame aligned with {df_aligned.shape[1]} columns before scaling.")

        # 4. Scale and Reshape for CNN
        # Use .values to get a NumPy array
        data_scaled = scaler.transform(df_aligned.values)
        
        final_feature_count = data_scaled.shape[1]
        print(f"Debug: Scaler output size: {final_feature_count} features.")
        
        if final_feature_count != FEATURE_COUNT:
            error_msg = f"SCALER ERROR: Expected {FEATURE_COUNT} features, but scaled data has {final_feature_count}."
            print(f"CRITICAL: {error_msg}")
            return f"<h2 style='color: red; text-align: center;'>{error_msg}</h2>", "N/A"
            
        # Reshape for the 1D CNN: (1 sample, FEATURE_COUNT features, 1 channel)
        X_processed = data_scaled.reshape(1, FEATURE_COUNT, 1)

        # 5. Predict probability
        prediction_prob = model.predict(X_processed, verbose=0)[0][0]
        
        # 6. Apply optimized threshold
        prediction_int = 1 if prediction_prob >= PREDICTION_THRESHOLD else 0
        
        # 7. Inverse transform the prediction
        prediction_label = label_encoder.inverse_transform([prediction_int])[0].upper()
        
        
        # 8. Determine result display
        if prediction_label == 'ANOMALY':
            color = "red"
            confidence_value = prediction_prob
            message = f"🚨 ANOMALY DETECTED! (Confidence: {confidence_value:.4f})"
        else:
            color = "green"
            confidence_value = 1 - prediction_prob
            message = f"🟢 Connection is NORMAL. (Confidence: {confidence_value:.4f})"

        html_output = f"<h2 style='color: {color}; text-align: center;'>{message}</h2>"
        
        return html_output, f"{prediction_prob:.4f}"
        
    except Exception as e:
        error_msg = f"RUNTIME ERROR during prediction: {type(e).__name__}: {str(e)}"
        print(f"CRITICAL: {error_msg}")
        return f"<h2 style='color: red; text-align: center;'>{error_msg}</h2>", "N/A"


# --- Gradio Interface Definition ---

# Define input components corresponding to the 41 features
input_components = [
    gr.Number(label='duration (float, sec)', value=0.0),
    gr.Dropdown(label='protocol_type', choices=PROTOCOLS, value='tcp'),
    gr.Dropdown(label='service', choices=SERVICES, value='http'),
    gr.Dropdown(label='flag', choices=FLAGS, value='SF'),
    gr.Number(label='src_bytes (int)', value=491),
    gr.Number(label='dst_bytes (int)', value=0),
    gr.Dropdown(label='land (binary)', choices=[0, 1], value=0),
    gr.Number(label='wrong_fragment (int)', value=0),
    gr.Number(label='urgent (int)', value=0),
    gr.Number(label='hot (int)', value=0),
    gr.Number(label='num_failed_logins (int)', value=0),
    gr.Dropdown(label='logged_in (binary)', choices=[0, 1], value=0),
    gr.Number(label='num_compromised (int)', value=0),
    gr.Dropdown(label='root_shell (binary)', choices=[0, 1], value=0),
    gr.Dropdown(label='su_attempted (binary)', choices=[0, 1], value=0),
    gr.Number(label='num_root (int)', value=0),
    gr.Number(label='num_file_creations (int)', value=0),
    gr.Number(label='num_shells (int)', value=0),
    gr.Number(label='num_access_files (int)', value=0),
    gr.Number(label='num_outbound_cmds (int)', value=0),
    gr.Dropdown(label='is_host_login (binary)', choices=[0, 1], value=0),
    gr.Dropdown(label='is_guest_login (binary)', choices=[0, 1], value=0),
    gr.Number(label='count (float)', value=2.0),
    gr.Number(label='srv_count (float)', value=2.0),
    gr.Number(label='serror_rate (float)', value=0.0),
    gr.Number(label='srv_serror_rate (float)', value=0.0),
    gr.Number(label='rerror_rate (float)', value=0.0),
    gr.Number(label='srv_rerror_rate (float)', value=0.0),
    gr.Number(label='same_srv_rate (float)', value=1.0),
    gr.Number(label='diff_srv_rate (float)', value=0.0),
    gr.Number(label='srv_diff_host_rate (float)', value=0.0),
    gr.Number(label='dst_host_count (float)', value=150.0),
    gr.Number(label='dst_host_srv_count (float)', value=25.0),
    gr.Number(label='dst_host_same_srv_rate (float)', value=0.17),
    gr.Number(label='dst_host_diff_srv_rate (float)', value=0.03),
    gr.Number(label='dst_host_same_src_port_rate (float)', value=0.17),
    gr.Number(label='dst_host_srv_diff_host_rate (float)', value=0.0),
    gr.Number(label='dst_host_serror_rate (float)', value=0.0),
    gr.Number(label='dst_host_srv_serror_rate (float)', value=0.0),
    gr.Number(label='dst_host_rerror_rate (float)', value=0.05),
    gr.Number(label='dst_host_srv_rerror_rate (float)', value=0.0)
]

# Define output components
output_components = [
    gr.HTML(label="Prediction Result"),
    gr.Label(label="Attack Probability")
]


iface = gr.Interface(
    fn=predict_intrusion,
    inputs=input_components,
    outputs=output_components,
    title="CNN Network Intrusion Detector (KDDCup'99)",
    description=(
        "Enter the 41 features of a network connection record to determine if it is "
        "a **Normal** connection or an **Anomaly (Attack)**. This model is now trained using **One-Hot Encoding** "
        f"for robust feature alignment (total of {len(master_columns) if master_columns else '...'} features expected).<br>"
        "Default values are set for a NORMAL FTP data connection."
    ),
    live=False,
    allow_flagging='never'
)

iface.launch()