662 lines
18 KiB
Python
662 lines
18 KiB
Python
from flask import Flask, jsonify, send_file, request
|
|
from flask_cors import CORS
|
|
import threading
|
|
import multiprocessing as mp
|
|
import gc
|
|
import traceback
|
|
import json
|
|
import time
|
|
import os
|
|
import platform
|
|
import psutil
|
|
import traceback
|
|
try:
|
|
import winreg
|
|
except ImportError:
|
|
winreg = None
|
|
from datetime import datetime
|
|
from reportlab.pdfgen import canvas
|
|
from reportlab.lib.pagesizes import A4
|
|
from reportlab.lib.units import cm
|
|
|
|
from isolation_forest_cicids import run_isolation_forest
|
|
from autoencoder import run_autoencoder
|
|
from isolation_forest_custom import run_isolation_forest_custom
|
|
from autoencoder_custom import run_autoencoder_custom
|
|
import signal
|
|
import sys
|
|
import shutil
|
|
from flask import send_from_directory
|
|
|
|
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
TEMP_DIR = os.path.join(BASE_DIR, "temp")
|
|
REPORT_DIR = os.path.join(BASE_DIR, "reports")
|
|
PROGRESS_PATH = os.path.join(BASE_DIR, "progress.json")
|
|
FRONTEND_DIR = os.path.join(BASE_DIR, "..", "Frontend")
|
|
HTML_DIR = os.path.join(FRONTEND_DIR, "HTML Files")
|
|
JS_DIR = os.path.join(FRONTEND_DIR, "JS")
|
|
IMG_DIR = os.path.join(FRONTEND_DIR, "img")
|
|
|
|
os.makedirs(TEMP_DIR, exist_ok=True)
|
|
os.makedirs(REPORT_DIR, exist_ok=True)
|
|
|
|
app = Flask(__name__)
|
|
CORS(app)
|
|
|
|
def force_memory_cleanup():
|
|
gc.collect()
|
|
try:
|
|
import ctypes
|
|
ctypes.CDLL("libc.so.6").malloc_trim(0)
|
|
except Exception:
|
|
pass
|
|
|
|
@app.route("/", methods=["GET"])
|
|
def serve_index():
|
|
return send_from_directory(HTML_DIR, "index.html")
|
|
|
|
@app.route("/upload", methods=["GET"])
|
|
def serve_upload():
|
|
return send_from_directory(HTML_DIR, "upload.html")
|
|
|
|
@app.route("/default", methods=["GET"])
|
|
def serve_default():
|
|
return send_from_directory(HTML_DIR, "default.html")
|
|
|
|
@app.route("/JS/<path:filename>", methods=["GET"])
|
|
def serve_js(filename):
|
|
return send_from_directory(JS_DIR, filename)
|
|
|
|
@app.route("/img/<path:filename>", methods=["GET"])
|
|
def serve_img(filename):
|
|
return send_from_directory(IMG_DIR, filename)
|
|
|
|
model_result_iforest = None
|
|
model_result_autoencoder = None
|
|
model_result_iforest_custom = None
|
|
model_result_autoencoder_custom = None
|
|
|
|
state_lock = threading.Lock()
|
|
is_running = False
|
|
current_job = None
|
|
worker_process = None
|
|
|
|
RESULT_IFOREST = os.path.join(TEMP_DIR, "result_iforest.json")
|
|
RESULT_AUTOENCODER = os.path.join(TEMP_DIR, "result_autoencoder.json")
|
|
RESULT_IFOREST_CUSTOM = os.path.join(TEMP_DIR, "result_iforest_custom.json")
|
|
RESULT_AUTOENCODER_CUSTOM = os.path.join(TEMP_DIR, "result_autoencoder_custom.json")
|
|
|
|
def handle_exit(signum, frame):
|
|
cleanup()
|
|
sys.exit(0)
|
|
|
|
signal.signal(signal.SIGINT, handle_exit)
|
|
signal.signal(signal.SIGTERM, handle_exit)
|
|
|
|
|
|
def clear_directory(path):
|
|
if not os.path.exists(path):
|
|
return
|
|
|
|
for item in os.listdir(path):
|
|
item_path = os.path.join(path, item)
|
|
try:
|
|
if os.path.isfile(item_path) or os.path.islink(item_path):
|
|
os.unlink(item_path)
|
|
elif os.path.isdir(item_path):
|
|
shutil.rmtree(item_path)
|
|
except Exception as e:
|
|
print(f"Failed to delete {item_path}: {e}")
|
|
|
|
def cleanup():
|
|
|
|
clear_directory(TEMP_DIR)
|
|
clear_directory(REPORT_DIR)
|
|
|
|
|
|
def update_progress(value):
|
|
with open(PROGRESS_PATH, "w") as f:
|
|
json.dump({"progress": value}, f)
|
|
|
|
def read_json_file(path, default=None):
|
|
try:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
return json.load(f)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def write_json_file(path, data):
|
|
with open(path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=2)
|
|
|
|
|
|
def get_result_path(job_name):
|
|
mapping = {
|
|
"iforest": RESULT_IFOREST,
|
|
"autoencoder": RESULT_AUTOENCODER,
|
|
"iforest_custom": RESULT_IFOREST_CUSTOM,
|
|
"autoencoder_custom": RESULT_AUTOENCODER_CUSTOM,
|
|
}
|
|
return mapping[job_name]
|
|
|
|
|
|
def get_job_runner(job_name):
|
|
if job_name == "iforest":
|
|
return lambda: run_isolation_forest(plot=False, table=False)
|
|
elif job_name == "autoencoder":
|
|
return lambda: run_autoencoder(plot=False, table=False)
|
|
elif job_name == "iforest_custom":
|
|
return run_isolation_forest_custom
|
|
elif job_name == "autoencoder_custom":
|
|
return run_autoencoder_custom
|
|
else:
|
|
raise ValueError(f"Unknown job name: {job_name}")
|
|
|
|
|
|
def worker_entry(job_name):
|
|
result_path = get_result_path(job_name)
|
|
start = time.time()
|
|
|
|
try:
|
|
update_progress(1)
|
|
|
|
runner = get_job_runner(job_name)
|
|
result = runner()
|
|
|
|
end = time.time()
|
|
if isinstance(result, dict):
|
|
result["runtime"] = round(end - start, 2)
|
|
|
|
write_json_file(result_path, {
|
|
"ready": True,
|
|
"data": result
|
|
})
|
|
|
|
update_progress(100)
|
|
|
|
except Exception as e:
|
|
write_json_file(result_path, {
|
|
"ready": True,
|
|
"data": {
|
|
"error": True,
|
|
"message": str(e),
|
|
"details": traceback.format_exc()
|
|
}
|
|
})
|
|
update_progress(0)
|
|
|
|
def refresh_worker_state():
|
|
global is_running, worker_process, current_job
|
|
|
|
with state_lock:
|
|
proc = worker_process
|
|
running = is_running
|
|
|
|
if running and proc is not None and not proc.is_alive():
|
|
try:
|
|
proc.join(timeout=0.2)
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
proc.close()
|
|
except Exception:
|
|
pass
|
|
|
|
with state_lock:
|
|
is_running = False
|
|
worker_process = None
|
|
current_job = None
|
|
|
|
force_memory_cleanup()
|
|
|
|
|
|
|
|
def start_background_job(job_name):
|
|
global is_running, current_job, worker_process
|
|
result_path = get_result_path(job_name)
|
|
|
|
refresh_worker_state()
|
|
|
|
with state_lock:
|
|
if is_running:
|
|
return False
|
|
|
|
is_running = True
|
|
current_job = job_name
|
|
|
|
if os.path.exists(result_path):
|
|
try:
|
|
os.remove(result_path)
|
|
except Exception:
|
|
pass
|
|
|
|
update_progress(0)
|
|
|
|
proc = mp.Process(target=worker_entry, args=(job_name,))
|
|
proc.start()
|
|
|
|
with state_lock:
|
|
worker_process = proc
|
|
|
|
return True
|
|
|
|
|
|
def read_progress():
|
|
try:
|
|
with open(PROGRESS_PATH, "r") as f:
|
|
return json.load(f).get("progress", 0)
|
|
except Exception:
|
|
return 0
|
|
|
|
update_progress(0)
|
|
|
|
|
|
@app.route("/start-iforest", methods=["POST"])
|
|
def start_iforest():
|
|
started = start_background_job("iforest")
|
|
if not started:
|
|
return jsonify({
|
|
"status": "already_running",
|
|
"message": "Another algorithm is already running"
|
|
}), 409
|
|
return jsonify({"status": "started"})
|
|
|
|
|
|
@app.route("/start-iforest-custom", methods=["POST"])
|
|
def start_iforest_custom():
|
|
started = start_background_job("iforest_custom")
|
|
if not started:
|
|
return jsonify({
|
|
"status": "already_running",
|
|
"message": "Another algorithm is already running"
|
|
}), 409
|
|
return jsonify({"status": "started"})
|
|
|
|
@app.route("/start-autoencoder-custom", methods=["POST"])
|
|
def start_autoencoder_custom():
|
|
started = start_background_job("autoencoder_custom")
|
|
if not started:
|
|
return jsonify({
|
|
"status": "already_running",
|
|
"message": "Another algorithm is already running"
|
|
}), 409
|
|
return jsonify({"status": "started"})
|
|
|
|
|
|
@app.route("/start-autoencoder", methods=["POST"])
|
|
def start_autoencoder():
|
|
started = start_background_job("autoencoder")
|
|
if not started:
|
|
return jsonify({
|
|
"status": "already_running",
|
|
"message": "Another algorithm is already running"
|
|
}), 409
|
|
return jsonify({"status": "started"})
|
|
|
|
|
|
@app.route("/get-progress", methods=["GET"])
|
|
def get_progress():
|
|
refresh_worker_state()
|
|
return jsonify(read_json_file(PROGRESS_PATH, {"progress": 0}))
|
|
|
|
|
|
@app.route("/get-status", methods=["GET"])
|
|
def get_status():
|
|
refresh_worker_state()
|
|
|
|
with state_lock:
|
|
running = is_running
|
|
job = current_job
|
|
|
|
progress_data = read_json_file(PROGRESS_PATH, {"progress": 0})
|
|
|
|
return jsonify({
|
|
"running": running,
|
|
"current_job": job,
|
|
"progress": progress_data.get("progress", 0)
|
|
})
|
|
|
|
|
|
@app.route("/get-isolation-forest-result", methods=["GET"])
|
|
def get_iforest_result():
|
|
refresh_worker_state()
|
|
data = read_json_file(RESULT_IFOREST)
|
|
if not data:
|
|
return jsonify({"ready": False})
|
|
return jsonify(data)
|
|
|
|
@app.route("/get-isolation-forest-result-custom", methods=["GET"])
|
|
def get_iforest_result_custom():
|
|
refresh_worker_state()
|
|
data = read_json_file(RESULT_IFOREST_CUSTOM)
|
|
if not data:
|
|
return jsonify({"ready": False})
|
|
return jsonify(data)
|
|
|
|
|
|
@app.route("/get-autoencoder-result", methods=["GET"])
|
|
def get_auto_result():
|
|
refresh_worker_state()
|
|
data = read_json_file(RESULT_AUTOENCODER)
|
|
if not data:
|
|
return jsonify({"ready": False})
|
|
return jsonify(data)
|
|
|
|
@app.route("/get-autoencoder-result-custom", methods=["GET"])
|
|
def get_autoencoder_result_custom():
|
|
refresh_worker_state()
|
|
data = read_json_file(RESULT_AUTOENCODER_CUSTOM)
|
|
if not data:
|
|
return jsonify({"ready": False})
|
|
return jsonify(data)
|
|
|
|
|
|
|
|
def generate_pdf_report(results, algorithm, params):
|
|
filename = os.path.join(REPORT_DIR, "report.pdf")
|
|
c = canvas.Canvas(filename, pagesize=A4)
|
|
|
|
w, h = A4
|
|
y = h - 2*cm
|
|
|
|
|
|
c.setFont("Helvetica-Bold", 24)
|
|
c.drawCentredString(w/2, y, f"{algorithm} Report")
|
|
y -= 1.4*cm
|
|
|
|
c.setFont("Helvetica", 11)
|
|
c.drawCentredString(
|
|
w/2, y,
|
|
f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
|
|
)
|
|
y -= 1.6*cm
|
|
|
|
|
|
c.setFont("Helvetica-Bold", 15)
|
|
c.drawString(2*cm, y, "Algorithm Parameters:")
|
|
y -= 0.9*cm
|
|
|
|
c.setFont("Helvetica", 11)
|
|
t = c.beginText(2.2*cm, y)
|
|
t.setLeading(13)
|
|
for k, v in params.items():
|
|
t.textLine(f"{k}: {v}")
|
|
c.drawText(t)
|
|
y = t.getY() - 1.1*cm
|
|
|
|
|
|
c.setFont("Helvetica-Bold", 15)
|
|
c.drawString(2*cm, y, "Execution Summary:")
|
|
y -= 0.9*cm
|
|
|
|
c.setFont("Helvetica", 11)
|
|
t = c.beginText(2.2*cm, y)
|
|
t.setLeading(13)
|
|
t.textLine(f"Runtime: {results['runtime']} seconds")
|
|
t.textLine(f"Normal Traffic: {int(results['normal_count'])}")
|
|
t.textLine(f"Detected Attacks: {int(results['attack_count'])}")
|
|
t.textLine(f"Accuracy: {results['accuracy']:.4f}")
|
|
c.drawText(t)
|
|
y = t.getY() - 1.3*cm
|
|
|
|
|
|
sys = get_system_info()
|
|
|
|
c.setFont("Helvetica-Bold", 15)
|
|
c.drawString(2 * cm, y, "System Information")
|
|
y -= 0.9 * cm
|
|
|
|
c.setFont("Helvetica", 11)
|
|
info = c.beginText(2.2 * cm, y)
|
|
info.setLeading(13)
|
|
info.textLine(f"CPU: {sys['cpu']}")
|
|
info.textLine(f"Physical Cores: {sys['cpu_physical_cores']}")
|
|
info.textLine(f"Logical Cores: {sys['cpu_total_cores']}")
|
|
info.textLine(f"GPU: {sys['gpu_name']}")
|
|
info.textLine(f"OS: {sys['os']}")
|
|
info.textLine(f"Python Version: {sys['python_version']}")
|
|
info.textLine(f"System RAM: {sys['ram_total_gb']} GB")
|
|
c.drawText(info)
|
|
y = info.getY() - 1.3 * cm
|
|
|
|
|
|
c.setFont("Helvetica-Bold", 15)
|
|
c.drawString(2 * cm, y, "Memory Usage:")
|
|
y -= 0.9 * cm
|
|
|
|
c.setFont("Helvetica", 11)
|
|
t = c.beginText(2.2 * cm, y)
|
|
t.setLeading(13)
|
|
t.textLine(f"RAM Before: {results['ram_before']} MB")
|
|
t.textLine(f"Peak RAM: {results['ram_peak']} MB")
|
|
t.textLine(f"RAM After: {results['ram_after']} MB")
|
|
t.textLine(f"Total Increase: {results['ram_increase']} MB")
|
|
c.drawText(t)
|
|
y = t.getY() - 1.1 * cm
|
|
|
|
|
|
c.setFont("Helvetica-Bold", 15)
|
|
c.drawString(2*cm, y, "Top 5 Most Significant Anomalies:")
|
|
y -= 1*cm
|
|
|
|
c.setFont("Helvetica", 11)
|
|
for anomaly in results["top_anomalies"]:
|
|
t = c.beginText(2.4*cm, y)
|
|
t.setLeading(13)
|
|
for k, v in anomaly.items():
|
|
t.textLine(f"{k}: {v}")
|
|
c.drawText(t)
|
|
y = t.getY() - 0.9*cm
|
|
|
|
if y < 3*cm:
|
|
c.showPage()
|
|
y = h - 2*cm
|
|
c.setFont("Helvetica", 11)
|
|
|
|
c.save()
|
|
return filename
|
|
|
|
@app.route("/download-pdf", methods=["POST"])
|
|
def download_pdf():
|
|
data = request.get_json()
|
|
algorithm = data["algorithm"]
|
|
results = data["results"]
|
|
|
|
raw_algorithm = data["algorithm"]
|
|
algorithm_name = raw_algorithm.lower()
|
|
|
|
if "isolation forest" in algorithm_name:
|
|
algo_title = "Isolation Forest"
|
|
params = {
|
|
"n_estimators": 600,
|
|
"max_samples": 0.3,
|
|
"contamination": results.get("contamination"),
|
|
"max_features": 0.7
|
|
}
|
|
|
|
elif "autoencoder" in algorithm_name:
|
|
algo_title = "Autoencoder"
|
|
params = {
|
|
"epochs": 20,
|
|
"batch_size": 64,
|
|
"latent_dimension": 16,
|
|
"validation_split": 0.1
|
|
}
|
|
|
|
else:
|
|
params = {}
|
|
|
|
pdf_path = generate_pdf_report(results, algo_title, params)
|
|
return send_file(pdf_path, as_attachment=True)
|
|
|
|
@app.route("/get-system-info", methods=["GET"])
|
|
def system_info_endpoint():
|
|
return jsonify(get_system_info())
|
|
|
|
def get_cpu_name():
|
|
try:
|
|
key = winreg.OpenKey(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"HARDWARE\DESCRIPTION\System\CentralProcessor\0"
|
|
)
|
|
cpu_name, _ = winreg.QueryValueEx(key, "ProcessorNameString")
|
|
winreg.CloseKey(key)
|
|
return cpu_name.strip()
|
|
except:
|
|
return platform.processor() or "Unknown CPU"
|
|
|
|
def get_system_info():
|
|
return {
|
|
"cpu": get_cpu_name(),
|
|
"cpu_physical_cores": psutil.cpu_count(logical=False),
|
|
"cpu_total_cores": psutil.cpu_count(logical=True),
|
|
"gpu_available": False,
|
|
"gpu_name": "None",
|
|
"os": platform.platform(),
|
|
"python_version": platform.python_version(),
|
|
"ram_total_gb": round(psutil.virtual_memory().total / (1024**3), 2),
|
|
"ram_used_mb": round(psutil.virtual_memory().used / (1024**2), 2),
|
|
}
|
|
|
|
@app.route("/download-json", methods=["POST"])
|
|
def download_json():
|
|
data = request.get_json()
|
|
|
|
algorithm = data["algorithm"]
|
|
results = data["results"]
|
|
runtime = data["runtime"]
|
|
|
|
raw_algorithm = data["algorithm"]
|
|
algorithm_name = raw_algorithm.lower()
|
|
|
|
if "isolation forest" in algorithm_name:
|
|
algo_title = "Isolation Forest"
|
|
params = {
|
|
"n_estimators": 600,
|
|
"max_samples": 0.3,
|
|
"contamination": results.get("contamination"),
|
|
"max_features": 0.7
|
|
}
|
|
|
|
elif "autoencoder" in algorithm_name:
|
|
algo_title = "Autoencoder"
|
|
params = {
|
|
"epochs": 20,
|
|
"batch_size": 64,
|
|
"latent_dimension": 16,
|
|
"validation_split": 0.1
|
|
}
|
|
|
|
system_info = get_system_info()
|
|
|
|
ram_info = {
|
|
"ram_before": results["ram_before"],
|
|
"ram_peak": results["ram_peak"],
|
|
"ram_after": results["ram_after"],
|
|
"ram_increase": results["ram_increase"]
|
|
}
|
|
|
|
anomalies = results.get("top_anomalies", [])
|
|
|
|
json_report = {
|
|
"algorithm": algorithm,
|
|
"generated": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
|
"algorithm_parameters": params,
|
|
|
|
"execution_summary": {
|
|
"runtime_seconds": runtime,
|
|
"normal_traffic": results["normal_count"],
|
|
"detected_attacks": results["attack_count"],
|
|
"accuracy": results["accuracy"]
|
|
},
|
|
|
|
"system_information": system_info,
|
|
|
|
"memory_usage": ram_info,
|
|
|
|
"top_5_anomalies": anomalies
|
|
}
|
|
|
|
json_path = os.path.join(REPORT_DIR, "report.json")
|
|
with open(json_path, "w", encoding="utf-8") as f:
|
|
json.dump(json_report, f, indent=4)
|
|
|
|
return send_file(json_path, as_attachment=True)
|
|
|
|
|
|
@app.route("/upload-dataset", methods=["POST"])
|
|
def upload_dataset():
|
|
refresh_worker_state()
|
|
|
|
with state_lock:
|
|
if is_running:
|
|
return jsonify({
|
|
"error": "Cannot upload a new dataset while an algorithm is running"
|
|
}), 409
|
|
|
|
if "file" not in request.files:
|
|
return jsonify({"error": "No file provided"}), 400
|
|
|
|
file = request.files["file"]
|
|
|
|
if file.filename == "":
|
|
return jsonify({"error": "Empty filename"}), 400
|
|
|
|
if not file.filename.lower().endswith(".csv"):
|
|
return jsonify({"error": "Only CSV files are allowed"}), 400
|
|
|
|
csv_path = os.path.join(TEMP_DIR, "upload.csv")
|
|
tmp_path = os.path.join(TEMP_DIR, "upload.tmp.csv")
|
|
|
|
try:
|
|
if os.path.exists(tmp_path):
|
|
os.remove(tmp_path)
|
|
|
|
file.save(tmp_path)
|
|
os.replace(tmp_path, csv_path)
|
|
|
|
return jsonify({"status": "uploaded", "path": csv_path})
|
|
finally:
|
|
try:
|
|
file.close()
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
if hasattr(file, "stream") and file.stream:
|
|
file.stream.close()
|
|
except Exception:
|
|
pass
|
|
|
|
force_memory_cleanup()
|
|
|
|
@app.route("/save-config", methods=["POST"])
|
|
def save_config():
|
|
refresh_worker_state()
|
|
|
|
with state_lock:
|
|
if is_running:
|
|
return jsonify({
|
|
"error": "Cannot save config while an algorithm is running"
|
|
}), 409
|
|
|
|
data = request.get_json()
|
|
|
|
if data is None:
|
|
return jsonify({"error": "No JSON provided"}), 400
|
|
|
|
config_path = os.path.join(TEMP_DIR, "config.json")
|
|
|
|
with open(config_path, "w", encoding="utf-8") as f:
|
|
json.dump(data, f, indent=4)
|
|
|
|
force_memory_cleanup()
|
|
return jsonify({"status": "config_saved"})
|
|
|
|
if __name__ == "__main__":
|
|
mp.freeze_support()
|
|
app.run(host="0.0.0.0", port=5000, debug=False, use_reloader=False)
|