Sniffer_AI.py

import gradio as gr
import joblib
import requests
import os

from sklearn.ensemble import RandomForestClassifier

# Load the saved models
rf_model = joblib.load('rf_model.pkl')

# Define the feature names (excluding the target column 'type')
feature_names = [
    "date", "time", "door_state", "sphone_signal", "label"
]

class_labels = {
    0: "normal",
    1: "backdoor",
    2: "ddos",
    3: "injection",
    4: "password",
    5: "ransomware",
    6: "scanning",
    7: "xss",
}

# Placeholder model (replace with actual Random Forest model object)
rf_model = None  # Load the actual trained Random Forest model here

def detect_intrusion(file):
    # Read the uploaded log file as a CSV or structured data
    try:
        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)}"

    # 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

    # Predict the class (multi-class classification) for each row in the log file
    predictions = rf_model.predict(feature_values)

    # Return only the 'Prediction' and 'label' columns
    return log_data[['Prediction']].head().to_string()

# Create a Gradio interface
iface = gr.Interface(
    fn=detect_intrusion,
    inputs=[
        gr.File(label="Upload Log File (CSV format)")  # File input
    ],  
    outputs="text", 
    title="Intrusion Detection System",
    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
iface.launch()