Update Sniffer_AI.py

Sniffer_AI.py

@@ -3,88 +3,66 @@ import joblib
 import requests
 import os
 
-from sklearn.ensemble import RandomForestClassifier, BaggingClassifier, AdaBoostClassifier
-from sklearn.tree import DecisionTreeClassifier
-
+from sklearn.ensemble import RandomForestClassifier
 
 # Load the saved models
 rf_model = joblib.load('rf_model.pkl')
-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
+# Define the feature names (excluding the target column 'type')
 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"
+    "date", "time", "door_state", "sphone_signal", "label"
 ]
 
 class_labels = {
     0: "normal",
     1: "backdoor",
     2: "ddos",
-    3: "dos",
-    4: "injection",
-    5: "password",
-    6: "ransomware",
-    7: "scanning",
-    8: "xss",
-    9: "mitm"
+    3: "injection",
+    4: "password",
+    5: "ransomware",
+    6: "scanning",
+    7: "xss",
 }
 
-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."
+# Placeholder model (replace with actual Random Forest model object)
+rf_model = None  # Load the actual trained Random Forest model here
 
-    # Convert the input values to floats and match them with feature names
+def detect_intrusion(file):
+    # Read the uploaded log file as a CSV or structured data
     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 = dt_model
-    elif model_choice == "Bagging Classifier":
-        model = bagging_model
-    elif model_choice == "AdaBoost Classifier":
-        model = ada_model
-    else:
-        return "Invalid model choice!"
-
-    # Predict the class (multi-class classification)
-    prediction = model.predict([feature_values])
-    predicted_class = prediction[0]  # Get the predicted class (an integer between 0-9)
+        log_data = pd.read_csv(file.name)  # Use file.name to get the path for reading
+    except Exception as e:
+        return f"Error reading file: {str(e)}"
 
-    # 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')}"
+    # Check if all required feature columns are in the log file
+    missing_features = [feature for feature in feature_names if feature not in log_data.columns]
+    if missing_features:
+        return f"Missing features in file: {', '.join(missing_features)}"
+    
+    # Extract the feature values (excluding the 'type' column which is the target)
+    feature_values = log_data[feature_names].astype(float).values
 
-# 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
+    # Predict the class (multi-class classification) for each row in the log file
+    predictions = rf_model.predict(feature_values)
 
-# Add model choice dropdown
-inputs.append(gr.Dropdown(choices=["Random Forest", "Decision Tree", "Bagging Classifier", "AdaBoost Classifier"], label="Select Model"))
+    # Return only the 'Prediction' and 'label' columns
+    return log_data[['Prediction']].head().to_string()
 
-# Create the Gradio interface
+# Create a Gradio interface
 iface = gr.Interface(
-    fn=detect_intrusion, 
-    inputs=inputs, 
+    fn=detect_intrusion,
+    inputs=[
+        gr.File(label="Upload Log File (CSV format)")  # File input
+    ],  
     outputs="text", 
     title="Intrusion Detection System",
-    description="Fill in the blank fields for the network traffic features, and choose the model to detect intrusions."
+    description=("""
+        Upload a CSV log file containing the following features: 
+        date, time, door_state, sphone_signal, label (without the 'type' column).
+        Example file structure:
+        date,time,door_state,sphone_signal,label
+        2025-03-12,10:45:00,1,-85,normal
+    """)
 )
 
 # Launch the interface locally for testing