# 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 "

FATAL ERROR: Artifacts Not Loaded. See Logs.

", "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"

{error_msg}

", "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"

{message}

" 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"

{error_msg}

", "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).
" "Default values are set for a NORMAL FTP data connection." ), live=False, allow_flagging='never' ) iface.launch()