api_server/app.py

import os
import threading
import time
from collections import deque

import joblib
import pandas as pd
from flask import Flask, jsonify, request
from flask_cors import CORS
from flask_socketio import SocketIO

import sys

sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from ml_models.ensemble_pipeline import load_ensemble
from ollama_ai.ai_analysis import analyze_network_traffic


if getattr(sys, "frozen", False):
    BASE_DIR = sys._MEIPASS
else:
    BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

app = Flask(__name__, static_folder=os.path.join(BASE_DIR, "dashboard"), static_url_path="/")
CORS(app)
socketio = SocketIO(app, cors_allowed_origins="*", async_mode="threading")

MODEL_PATH = os.path.join(BASE_DIR, "ml_models", "device_classifier.pkl")
SCALER_PATH = os.path.join(BASE_DIR, "datasets", "scaler.pkl")
LE_PATH = os.path.join(BASE_DIR, "datasets", "label_encoder.pkl")
ENSEMBLE_META_PATH = os.path.join(BASE_DIR, "ml_models", "ensemble_artifacts", "metadata.json")

model = None
scaler = None
label_encoder = None
ensemble_pipeline = None
sniffer_instance = None
recent_alerts = deque(maxlen=50)
latest_snapshot = None

_broadcast_started = False
_broadcast_lock = threading.Lock()
_snapshot_lock = threading.Lock()


def load_models():
    global model, scaler, label_encoder, ensemble_pipeline

    try:
        if model is None and os.path.exists(MODEL_PATH):
            model = joblib.load(MODEL_PATH)
        if scaler is None and os.path.exists(SCALER_PATH):
            scaler = joblib.load(SCALER_PATH)
        if label_encoder is None and os.path.exists(LE_PATH):
            label_encoder = joblib.load(LE_PATH)
        if ensemble_pipeline is None and os.path.exists(ENSEMBLE_META_PATH):
            ensemble_pipeline = load_ensemble()
    except Exception as exc:
        print(f"Error loading models: {exc}")


def _protocol_to_numeric(protocol: str) -> int:
    protocol = str(protocol).upper()
    if protocol == "TCP":
        return 0
    if protocol == "UDP":
        return 1
    return 2


def _append_alert(ip: str, threat_type: str, severity: str, confidence: float):
    alert = {
        "ip": ip,
        "threat_type": threat_type,
        "severity": severity,
        "confidence": round(confidence, 2),
        "timestamp": time.time(),
    }
    recent_alerts.appendleft(alert)
    return alert


def _build_legacy_predictions(features_list):
    results = []
    anomaly_count = 0

    batch_data = []
    for feature in features_list:
        protocol_val = _protocol_to_numeric(feature["protocol"])
        batch_data.append(
            {
                "packet_size": feature["packet_size"],
                "protocol": protocol_val,
                "packet_rate": feature["packet_rate"],
                "flow_duration": feature["flow_duration"],
            }
        )

    df = pd.DataFrame(batch_data)

    if scaler:
        expected_cols = getattr(scaler, "feature_names_in_", df.columns)
        for col in expected_cols:
            if col not in df.columns:
                df[col] = 0
        df = df[expected_cols]
        df_scaled = scaler.transform(df)
    else:
        df_scaled = df.values

    if model:
        preds_encoded = model.predict(df_scaled)
        pred_probs = model.predict_proba(df_scaled)
    else:
        preds_encoded = [0] * len(df)
        pred_probs = [[1.0]] * len(df)

    if label_encoder:
        device_types = label_encoder.inverse_transform(preds_encoded)
    else:
        classes = ["smartphone", "laptop", "smart_tv", "iot_device", "tablet"]
        device_types = [classes[p % len(classes)] for p in preds_encoded]

    for index, feature in enumerate(features_list):
        confidence = float(max(pred_probs[index]) * 100)
        device_type = device_types[index]
        is_suspicious = confidence < 60.0

        if is_suspicious:
            anomaly_count += 1
            _append_alert(
                feature["ip"],
                device_type,
                "Medium" if confidence >= 40 else "High",
                confidence,
            )

        results.append(
            {
                "ip": feature["ip"],
                "threat_type": device_type,
                "device_type": device_type,
                "confidence": round(confidence, 2),
                "is_suspicious": is_suspicious,
                "severity": "Suspicious" if is_suspicious else "Normal",
                "metrics": {
                    "packet_rate": round(feature["packet_rate"], 2),
                    "packet_size": round(feature["packet_size"], 2),
                    "byte_rate": round(feature.get("byte_rate", 0.0), 2),
                    "flow_duration": round(feature["flow_duration"], 2),
                },
            }
        )

    return results, anomaly_count


def _build_ensemble_predictions(features_list):
    results = []
    anomaly_count = 0

    for feature in features_list:
        prediction = ensemble_pipeline.predict(feature)
        confidence = float(prediction["confidence"] * 100.0)

        if prediction["is_anomaly"]:
            anomaly_count += 1
            _append_alert(
                feature["ip"],
                prediction["threat_type"],
                prediction["severity"],
                confidence,
            )

        results.append(
            {
                "ip": feature["ip"],
                "threat_type": prediction["threat_type"],
                "device_type": prediction["threat_type"],
                "confidence": round(confidence, 2),
                "is_suspicious": prediction["is_anomaly"],
                "severity": prediction["severity"],
                "anomaly_score": float(prediction["anomaly_score"]),
                "autoencoder_error": float(prediction["autoencoder_error"]),
                "metrics": {
                    "packet_rate": round(feature["packet_rate"], 2),
                    "packet_size": round(feature["packet_size"], 2),
                    "byte_rate": round(feature.get("byte_rate", 0.0), 2),
                    "flow_duration": round(feature["flow_duration"], 2),
                },
            }
        )

    return results, anomaly_count


def _build_snapshot():
    load_models()

    if sniffer_instance is None:
        return {
            "health": {"health_score": 100, "status": "Unknown", "anomalies": 0},
            "devices": [],
            "arp_devices": [],
            "alerts": list(recent_alerts)[:10],
            "total_devices": 0,
            "anomalies": 0,
            "timestamp": time.time(),
            "capture": {"mode": "idle", "running": False, "interface": None, "last_error": None},
        }

    features_list = sniffer_instance.get_latest_features()
    arp_devices = sniffer_instance.arp_scan()
    capture_status = sniffer_instance.get_status() if hasattr(sniffer_instance, "get_status") else {
        "mode": "unknown",
        "running": bool(getattr(sniffer_instance, "is_sniffing", False)),
        "interface": None,
        "last_error": getattr(sniffer_instance, "last_error", None),
    }

    if not features_list:
        return {
            "health": {"health_score": 100, "status": "Unknown", "anomalies": 0},
            "devices": [],
            "arp_devices": arp_devices,
            "alerts": list(recent_alerts)[:10],
            "total_devices": 0,
            "anomalies": 0,
            "timestamp": time.time(),
            "capture": capture_status,
        }

    if ensemble_pipeline is not None and os.path.exists(ENSEMBLE_META_PATH):
        devices, anomaly_count = _build_ensemble_predictions(features_list)
    else:
        devices, anomaly_count = _build_legacy_predictions(features_list)

    health_score = max(0, 100 - int(anomaly_count) * 10)
    status = "Good"
    if health_score < 80:
        status = "Warning"
    if health_score < 50:
        status = "Critical"

    total_traffic = round(float(sum(feature.get("byte_rate", 0.0) for feature in features_list)), 2)

    return {
        "health": {
            "health_score": health_score,
            "status": status,
            "anomalies": anomaly_count,
            "total_devices": len(features_list),
            "traffic_load": total_traffic,
        },
        "devices": devices,
        "arp_devices": arp_devices,
        "alerts": list(recent_alerts)[:10],
        "total_devices": len(devices),
        "anomalies": anomaly_count,
        "timestamp": time.time(),
        "capture": capture_status,
    }


def _store_snapshot(snapshot):
    global latest_snapshot
    with _snapshot_lock:
        latest_snapshot = snapshot


def _get_snapshot():
    with _snapshot_lock:
        if latest_snapshot is not None:
            return latest_snapshot

    snapshot = _build_snapshot()
    _store_snapshot(snapshot)
    return snapshot


def _broadcast_loop():
    while True:
        snapshot = _build_snapshot()
        _store_snapshot(snapshot)
        socketio.emit("network_update", snapshot)
        socketio.sleep(2)


def start_socket_broadcasts():
    global _broadcast_started

    with _broadcast_lock:
        if _broadcast_started:
            return
        _broadcast_started = True
        socketio.start_background_task(_broadcast_loop)


@app.route("/")
def index():
    return app.send_static_file("index.html")


@app.route("/scan-network", methods=["GET"])
def scan_network():
    snapshot = _get_snapshot()
    return jsonify(
        {
            "devices": snapshot["devices"],
            "arp_devices": snapshot["arp_devices"],
            "alerts": snapshot["alerts"],
            "total_devices": snapshot["total_devices"],
            "anomalies": snapshot["anomalies"],
            "health": snapshot["health"],
            "timestamp": snapshot["timestamp"],
            "capture": snapshot.get("capture", {}),
        }
    )


@app.route("/arp-devices", methods=["GET"])
def arp_devices():
    if sniffer_instance is None:
        return jsonify({"devices": []})

    return jsonify({"devices": sniffer_instance.arp_scan()})


@app.route("/predict-device", methods=["POST"])
def predict_device():
    load_models()
    data = request.json

    if not data:
        return jsonify({"error": "No input data provided"}), 400

    required = ["packet_size", "protocol", "packet_rate", "flow_duration"]
    if not all(key in data for key in required):
        return jsonify({"error": f"Missing required fields. Expected: {required}"}), 400

    try:
        protocol_val = data["protocol"]
        if isinstance(protocol_val, str):
            protocol_val = _protocol_to_numeric(protocol_val)

        df = pd.DataFrame(
            [
                {
                    "packet_size": float(data["packet_size"]),
                    "protocol": protocol_val,
                    "packet_rate": float(data["packet_rate"]),
                    "flow_duration": float(data["flow_duration"]),
                }
            ]
        )

        if scaler:
            expected_cols = getattr(scaler, "feature_names_in_", df.columns)
            for col in expected_cols:
                if col not in df.columns:
                    df[col] = 0
            df = df[expected_cols]
            df_scaled = scaler.transform(df)
        else:
            df_scaled = df.values

        if not model:
            return jsonify({"error": "Model not trained yet."}), 503

        pred_encoded = model.predict(df_scaled)[0]
        confidence = float(max(model.predict_proba(df_scaled)[0]) * 100)

        if label_encoder:
            device_type = label_encoder.inverse_transform([pred_encoded])[0]
        else:
            classes = ["smartphone", "laptop", "smart_tv", "iot_device", "tablet"]
            device_type = classes[pred_encoded % len(classes)]

        return jsonify({"device_type": device_type, "confidence": round(confidence, 2)})

    except Exception as exc:
        import traceback

        return jsonify({"error": str(exc), "trace": traceback.format_exc()}), 500


@app.route("/network-health", methods=["GET"])
def network_health():
    snapshot = _get_snapshot()
    return jsonify(snapshot["health"])


@app.route("/capture-status", methods=["GET"])
def capture_status():
    snapshot = _get_snapshot()
    return jsonify(snapshot.get("capture", {"mode": "idle", "running": False, "interface": None, "last_error": None}))


@app.route("/ai-analysis", methods=["POST"])
def ai_analysis():
    data = request.json
    if not data:
        return jsonify({"error": "No data provided"}), 400

    result = analyze_network_traffic(data)
    return jsonify(result)


if __name__ == "__main__":
    app.run(debug=True, host="0.0.0.0", port=5000)