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

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

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

# Assuming models are in the same directory as the script for deployment
# For Colab, you might need to adjust paths or upload files.
MODEL_PATH = "web_attack_detection_model.pkl"
VECTORIZER_PATH = "url_vectorizer.pkl"

class InputData(BaseModel):
    features: list[float] # This assumes features are already numerical
    # If you want to input URLs, you'd change this to:
    # url: str


try:
    model = joblib.load(MODEL_PATH)
    vectorizer = joblib.load(VECTORIZER_PATH)
    EXPECTED_FEATURES = len(vectorizer.get_feature_names_out())
except Exception as ex:
    raise RuntimeError(f"Model or vectorizer failed to load: {ex}") from ex


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


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


@app.post("/predict")
def predict(data: InputData) -> dict[str, int | float | str]:
    # If InputData was 'url: str', you'd do:
    # x = vectorizer.transform([data.url]).toarray()
    
    if len(data.features) != EXPECTED_FEATURES:
        raise HTTPException(
            status_code=400,
            detail=f"Expected {EXPECTED_FEATURES} features, but got {len(data.features)}",
        )

    try:
        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")

        pred = int(model.predict(x)[0])
        # IsolationForest predicts -1 for anomalies (attacks), 1 for normal
        result = "attack_detected" if pred == -1 else "normal_request"

        response: dict[str, int | float | str] = {
            "prediction": result,
            "raw_prediction": pred,
        }

        # IsolationForest does not have predict_proba, use decision_function for anomaly score
        # Lower scores typically mean more anomalous
        anomaly_score = float(model.decision_function(x)[0])
        response["anomaly_score"] = anomaly_score

        return response
    except HTTPException:
        raise
    except Exception as ex:
        raise HTTPException(status_code=500, detail=f"Prediction failed: {ex}") from ex