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, }