Spaces:
Sleeping
Sleeping
File size: 4,481 Bytes
0ee8cc7 f6e3fba 0ee8cc7 f6e3fba 0ee8cc7 5878037 0ee8cc7 5878037 0ee8cc7 f6e3fba 0ee8cc7 f6e3fba e12b043 f6e3fba 0ee8cc7 f6e3fba 0ee8cc7 f6e3fba 0ee8cc7 f6e3fba 0ee8cc7 f6e3fba 0ee8cc7 42293a6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
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)
|