Update Sniffer_AI.py

Sniffer_AI.py

@@ -13,6 +13,19 @@ dt_model = joblib.load('decision_tree_model.pkl')
 bagging_model = joblib.load('model_bagging.pkl')
 ada_model = joblib.load('model_adaboost.pkl')
 
+# Define the feature names
+feature_names = [
+    "src_ip", "src_port", "dst_ip", "dst_port", "proto", "service", "duration", 
+    "src_bytes", "dst_bytes", "conn_state", "missed_bytes", "src_pkts", 
+    "src_ip_bytes", "dst_pkts", "dst_ip_bytes", "dns_query", "dns_qclass", 
+    "dns_qtype", "dns_rcode", "dns_AA", "dns_RD", "dns_RA", "dns_rejected", 
+    "ssl_version", "ssl_cipher", "ssl_resumed", "ssl_established", "ssl_subject", 
+    "ssl_issuer", "http_trans_depth", "http_method", "http_uri", "http_version", 
+    "http_request_body_len", "http_response_body_len", "http_status_code", 
+    "http_user_agent", "http_orig_mime_types", "http_resp_mime_types", 
+    "weird_name", "weird_addl", "weird_notice", "label"
+]
+
 class_labels = {
     0: "normal",
     1: "backdoor",
@@ -26,39 +39,53 @@ class_labels = {
     9: "mitm"
 }
 
-def detect_intrusion(features, model_choice="Random Forest"):
-    # Convert the input string (comma-separated values) into a list of floats
-    features = [list(map(float, features.split(",")))]
-    
+def detect_intrusion(feature_values, model_choice="Random Forest"):
+    # Ensure the length of feature_values matches feature_names
+    if len(feature_values) != len(feature_names):
+        return "Please fill in all the required feature values."
+
+    # Convert the input values to floats and match them with feature names
+    try:
+        feature_values = [float(value) for value in feature_values]
+    except ValueError:
+        return "Please enter valid numerical values for all fields."
+
     # Choose the model based on user selection
     if model_choice == "Random Forest":
         model = rf_model
     elif model_choice == "Decision Tree":
-        model = decision_tree_model
+        model = dt_model
     elif model_choice == "Bagging Classifier":
-        model = model_bagging
+        model = bagging_model
     elif model_choice == "AdaBoost Classifier":
-        model = model_adaboost
+        model = ada_model
     else:
         return "Invalid model choice!"
-    
+
     # Predict the class (multi-class classification)
-    prediction = model.predict(features)
-    predicted_class = prediction[0]  # Get the predicted class (an integer between 0-8)
-    
-    # Return the human-readable class description
+    prediction = model.predict([feature_values])
+    predicted_class = prediction[0]  # Get the predicted class (an integer between 0-9)
+
+    # Notify the user of the detected attack or normal traffic
     if predicted_class == 0:
         return "No Intrusion Detected"
     else:
         return f"Intrusion Detected: {class_labels.get(predicted_class, 'Unknown Attack')}"
 
-# Create a Gradio interface
-iface = gr.Interface(fn=detect_intrusion, 
-                     inputs=[gr.Textbox(label="Input Features (comma-separated)"), 
-                             gr.Dropdown(choices=["Random Forest", "Decision Tree", "Bagging Classifier", "AdaBoost Classifier"], label="Select Model")], 
-                     outputs="text", 
-                     title="Intrusion Detection System",
-                     description="Enter features in the format: feature1, feature2, feature3...")
+# Create Gradio input fields for each feature
+inputs = [gr.Textbox(label=feature_name) for feature_name in feature_names[:-1]]  # Exclude "label" field from inputs
+
+# Add model choice dropdown
+inputs.append(gr.Dropdown(choices=["Random Forest", "Decision Tree", "Bagging Classifier", "AdaBoost Classifier"], label="Select Model"))
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=detect_intrusion, 
+    inputs=inputs, 
+    outputs="text", 
+    title="Intrusion Detection System",
+    description="Fill in the blank fields for the network traffic features, and choose the model to detect intrusions."
+)
 
 # Launch the interface locally for testing
 iface.launch()