Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import pandas as pd | |
| import pickle | |
| # Load the trained models | |
| with open("cyber_iot_8class.pkl", "rb") as model_file_8class: | |
| model_8class = pickle.load(model_file_8class) | |
| with open("cyber_iot_34class.pkl", "rb") as model_file_34class: | |
| model_34class = pickle.load(model_file_34class) | |
| # Use 15 features, drop other features as part of feature engineering | |
| dropped_features = [ | |
| "flow_duration", | |
| "urg_count", | |
| "Duration", | |
| "Number", | |
| "Weight", | |
| "Srate", | |
| "Rate", | |
| "ack_count", | |
| "fin_count", | |
| "TCP", | |
| "ack_flag_number", | |
| "syn_flag_number", | |
| "UDP", | |
| "ICMP", | |
| "fin_flag_number", | |
| "rst_flag_number", | |
| "psh_flag_number", | |
| "HTTP", | |
| "HTTPS", | |
| "IPv", | |
| "LLC", | |
| "SSH", | |
| "ece_flag_number", | |
| "SMTP", | |
| "ARP", | |
| "DHCP", | |
| "DNS", | |
| "cwr_flag_number", | |
| "Drate", | |
| "Telnet", | |
| "IRC" | |
| ] | |
| def make_prediction(file): | |
| if file: | |
| # Read the uploaded CSV file | |
| data = pd.read_csv(file.name) | |
| data = data.drop(columns=dropped_features, axis=1) | |
| # Make initial prediction using the 8-class model | |
| initial_prediction = model_8class.predict(data) | |
| # Convert the prediction results to a list of strings | |
| result_benign_malicious = [str(pred) for pred in initial_prediction] | |
| #result_benign_malicious = str(initial_prediction) | |
| # Initialize an empty list to store detailed predictions | |
| result_type_of_attack = [] | |
| # Check the predictions based on the given conditions | |
| if all(pred == "Benign" for pred in result_benign_malicious): | |
| # Condition 1: If all predictions are "Benign" | |
| result_type_of_attack.append("N/A") | |
| gr.Info("All connections are fine/Benign") | |
| elif "Benign" in result_benign_malicious and len(set(result_benign_malicious)) > 1: | |
| # Condition 2: If "Benign" and other attacks are present | |
| # Remove the Benign entries from the dataframe for 2nd level check | |
| for i in result_benign_malicious: | |
| if i == "Benign": | |
| row = result_benign_malicious.index(i) | |
| print("row = ", row) | |
| data.drop(data.index[row], inplace=True) | |
| detailed_prediction = model_34class.predict(data) | |
| result_type_of_attack.append(str(detailed_prediction)) | |
| gr.Warning("Malicious attack detected!") | |
| elif "Benign" not in result_benign_malicious: | |
| # Condition 3: If all are attack connections | |
| detailed_prediction = model_34class.predict(data) | |
| result_type_of_attack.append(str(detailed_prediction)) | |
| gr.Warning("Malicious attack detected!") | |
| # Convert the results to strings for display in the Gradio interface | |
| result_benign_malicious_str = ', '.join(result_benign_malicious) | |
| result_type_of_attack_str = ', '.join(result_type_of_attack) | |
| return result_benign_malicious_str, result_type_of_attack_str.strip("[]").replace("'", "").split() | |
| else: | |
| return "", "" | |
| # Define the Gradio interface | |
| with gr.Blocks() as cyberIoT: | |
| # Add title and description using Markdown | |
| gr.Markdown("# IoT Cybersecurity Connection Prediction") | |
| gr.Markdown("Upload a .csv file to check whether it is a Non-Attack(Benign) or Attack(Non-Benign) connection. If Malicious, the type of attack will be further classified.") | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_input = gr.File(label="Upload PCAP CSV File", file_types=['csv']) | |
| #gr.ClearButton(components=file_input) | |
| with gr.Column(): | |
| output_benign_malicious = gr.Textbox(label="Benign/Malicious(BruteForce,DDoS,DoS,Mirai,Recon,Spoofing,Web)", placeholder="Initial prediction will appear here...") | |
| output_type_of_attack = gr.Textbox(label="Type of Attack", placeholder="Further classification will appear here if applicable...") | |
| # When file is uploaded, make predictions and display the results | |
| file_input.change(make_prediction, inputs=file_input, outputs=[output_benign_malicious, output_type_of_attack]) | |
| gr.ClearButton(components=[file_input, output_benign_malicious,output_type_of_attack]) | |
| # Launch the Gradio app | |
| cyberIoT.launch(share=True) | |