app.py

from pathlib import Path

import joblib
import numpy as np
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

app = FastAPI(title="Web Attack Detection API")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

EXPECTED_FEATURES = 35
BASE_DIR = Path(__file__).resolve().parent
MODEL_PATH = BASE_DIR / "traffic_model.pkl"
SCALER_PATH = BASE_DIR / "traffic_scaler.pkl"
THRESHOLD_PATH = BASE_DIR / "traffic_threshold.pkl"


class InputData(BaseModel):
    features: list[float] = Field(
        ...,
        min_length=EXPECTED_FEATURES,
        max_length=EXPECTED_FEATURES,
    )


try:
    model = joblib.load(MODEL_PATH)
    scaler = joblib.load(SCALER_PATH)
    threshold = float(joblib.load(THRESHOLD_PATH))
except Exception as ex:
    raise RuntimeError(f"Startup load failed: {ex}") from ex


@app.get("/")
def home() -> dict[str, str]:
    return {"message": "Web attack detection model running"}


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/predict")
def predict(data: InputData) -> dict[str, float | int | str]:
    x = np.asarray(data.features, dtype=np.float64).reshape(1, -1)

    if not np.isfinite(x).all():
        raise HTTPException(status_code=400, detail="Features contain NaN or Inf")

    x_scaled = scaler.transform(x)
    class_index = list(model.classes_).index(-1)
    attack_probability = float(model.predict_proba(x_scaled)[0][class_index])
    raw_prediction = -1 if attack_probability >= threshold else 1

    return {
        "prediction": "attack detected" if raw_prediction == -1 else "normal request",
        "raw_prediction": raw_prediction,
        "attack_probability": attack_probability,
        "threshold": threshold,
    }