|
|
|
|
|
|
|
|
import os |
|
|
import numpy as np |
|
|
import pandas as pd |
|
|
import tensorflow as tf |
|
|
from tensorflow.keras.models import load_model |
|
|
import pickle |
|
|
from flask import Flask, request, jsonify, render_template |
|
|
|
|
|
|
|
|
MODEL_PATH = 'cicids2017_mlp_model.keras' |
|
|
SCALER_PATH = 'cicids2017_scaler.pkl' |
|
|
EXPECTED_FEATURES = 49 |
|
|
app = Flask(__name__) |
|
|
|
|
|
|
|
|
loaded_model = None |
|
|
loaded_scaler = None |
|
|
|
|
|
def load_assets(): |
|
|
"""Load the model and scaler only once when the app starts.""" |
|
|
global loaded_model, loaded_scaler |
|
|
|
|
|
|
|
|
if not os.path.exists(MODEL_PATH) or not os.path.exists(SCALER_PATH): |
|
|
print("🚨 ERROR: Model or scaler files not found. Ensure they are in the same directory.") |
|
|
return False |
|
|
|
|
|
try: |
|
|
|
|
|
loaded_model = load_model(MODEL_PATH) |
|
|
|
|
|
|
|
|
with open(SCALER_PATH, 'rb') as file: |
|
|
loaded_scaler = pickle.load(file) |
|
|
|
|
|
print(f"✅ Assets loaded successfully. Model ready for {EXPECTED_FEATURES} features.") |
|
|
return True |
|
|
|
|
|
except Exception as e: |
|
|
print(f"🚨 FATAL ERROR loading assets: {e}") |
|
|
return False |
|
|
|
|
|
|
|
|
load_assets() |
|
|
|
|
|
|
|
|
|
|
|
HTML_TEMPLATE = """ |
|
|
<!doctype html> |
|
|
<title>NIDS Prediction</title> |
|
|
<style> |
|
|
body {{ font-family: sans-serif; max-width: 800px; margin: auto; padding: 20px; }} |
|
|
textarea {{ width: 100%; min-height: 150px; padding: 10px; box-sizing: border-box; }} |
|
|
h1 {{ color: #007bff; }} |
|
|
.result {{ padding: 15px; border-radius: 5px; margin-top: 20px; }} |
|
|
.attack {{ background-color: #f8d7da; border: 1px solid #f5c6cb; color: #721c24; }} |
|
|
.benign {{ background-color: #d4edda; border: 1px solid #c3e6cb; color: #155724; }} |
|
|
.error {{ background-color: #ffeeba; border: 1px solid #ffc720; color: #664d03; }} |
|
|
</style> |
|
|
<h1>Network Intrusion Detection System (NIDS)</h1> |
|
|
<p>Paste {num_features} comma-separated network flow features below for prediction.</p> |
|
|
|
|
|
<form method="POST" action="/"> |
|
|
<textarea name="features" placeholder="e.g., 0.1, 10.5, 0.003, ...">{example_data}</textarea><br><br> |
|
|
<input type="submit" value="Predict Traffic Type" style="padding: 10px 20px; background-color: #007bff; color: white; border: none; cursor: pointer;"> |
|
|
</form> |
|
|
|
|
|
{result_html} |
|
|
""" |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
@app.route('/', methods=['GET', 'POST']) |
|
|
def home(): |
|
|
"""Handles both the form display (GET) and prediction submission (POST).""" |
|
|
result_html = "" |
|
|
example_data = "" |
|
|
|
|
|
if request.method == 'POST': |
|
|
try: |
|
|
|
|
|
features_str = request.form.get('features') |
|
|
features_list = [float(x.strip()) for x in features_str.split(',') if x.strip()] |
|
|
example_data = features_str |
|
|
|
|
|
if len(features_list) != EXPECTED_FEATURES: |
|
|
raise ValueError(f"Input must contain exactly {EXPECTED_FEATURES} features, but received {len(features_list)}.") |
|
|
|
|
|
|
|
|
input_array = np.array(features_list).reshape(1, -1) |
|
|
input_scaled = loaded_scaler.transform(input_array) |
|
|
|
|
|
|
|
|
prediction_proba = loaded_model.predict(input_scaled, verbose=0)[0][0] |
|
|
prediction_class = 1 if prediction_proba > 0.5 else 0 |
|
|
|
|
|
|
|
|
class_label = "🚨 ATTACK" if prediction_class == 1 else "✅ BENIGN" |
|
|
css_class = "attack" if prediction_class == 1 else "benign" |
|
|
|
|
|
result_html = f""" |
|
|
<div class="result {css_class}"> |
|
|
<h3>Prediction: {class_label}</h3> |
|
|
<p>Probability of Attack (1): <b>{prediction_proba:.5f}</b></p> |
|
|
</div> |
|
|
""" |
|
|
|
|
|
except ValueError as e: |
|
|
result_html = f""" |
|
|
<div class="result error"> |
|
|
<h3>Input Error</h3> |
|
|
<p>{e}</p> |
|
|
</div> |
|
|
""" |
|
|
except Exception as e: |
|
|
result_html = f""" |
|
|
<div class="result error"> |
|
|
<h3>Prediction Error</h3> |
|
|
<p>An unexpected error occurred: {e}</p> |
|
|
</div> |
|
|
""" |
|
|
|
|
|
return render_template_string( |
|
|
HTML_TEMPLATE, |
|
|
num_features=EXPECTED_FEATURES, |
|
|
example_data=example_data, |
|
|
result_html=result_html |
|
|
) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
|
|
app.run(host='0.0.0.0', port=5000) |
