# app.py
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
# --- Configuration ---
MODEL_PATH = 'cicids2017_mlp_model.keras'
SCALER_PATH = 'cicids2017_scaler.pkl'
EXPECTED_FEATURES = 49 # Based on your training data shape
app = Flask(__name__)
# --- Global Assets ---
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
# Check if files exist
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:
# Load the Keras model
loaded_model = load_model(MODEL_PATH)
# Load the StandardScaler object
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 when the application starts
load_assets()
# --- HTML Template for Simple Interface ---
HTML_TEMPLATE = """
NIDS Prediction
Network Intrusion Detection System (NIDS)
Paste {num_features} comma-separated network flow features below for prediction.
{result_html}
"""
# --- Flask Routes ---
@app.route('/', methods=['GET', 'POST'])
def home():
"""Handles both the form display (GET) and prediction submission (POST)."""
result_html = ""
example_data = "" # You can load a real example here if you have a local test file
if request.method == 'POST':
try:
# 1. Get raw input data
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)}.")
# 2. Scale the input
input_array = np.array(features_list).reshape(1, -1)
input_scaled = loaded_scaler.transform(input_array)
# 3. Predict
prediction_proba = loaded_model.predict(input_scaled, verbose=0)[0][0]
prediction_class = 1 if prediction_proba > 0.5 else 0
# 4. Format Result
class_label = "🚨 ATTACK" if prediction_class == 1 else "✅ BENIGN"
css_class = "attack" if prediction_class == 1 else "benign"
result_html = f"""
Prediction: {class_label}
Probability of Attack (1): {prediction_proba:.5f}
"""
except ValueError as e:
result_html = f"""
"""
except Exception as e:
result_html = f"""
Prediction Error
An unexpected error occurred: {e}
"""
return render_template_string(
HTML_TEMPLATE,
num_features=EXPECTED_FEATURES,
example_data=example_data,
result_html=result_html
)
if __name__ == '__main__':
# Use 0.0.0.0 for compatibility with Docker/server environments
app.run(host='0.0.0.0', port=5000)