app.py

from preditormodels import PhishingPredictor
import re
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, field_validator
from config import XGB_MODEL_PATH, BERT_MODEL_PATH
class EmailRequest(BaseModel):
    subject: str
    body: str

    @field_validator("subject", "body", mode="before")
    @classmethod
    def clean_text(cls, v):
        if isinstance(v, str):
            # Remove illegal ASCII control characters
            v = re.sub(r"[\x00-\x08\x0B\x0C\x0E-\x1F]", "", v)
        return v


# FastAPI App


app = FastAPI(title="Phishing Detection API")

# Define the specific origins that are allowed to call this API


app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],        # Only allow the domains in the list above
    allow_credentials=False,
    allow_methods=["*"],          # Allows all HTTP methods (POST, GET, etc.)
    allow_headers=["*"],          # Allows all headers (Content-Type, etc.)
)


# Load models once
predictor = PhishingPredictor(BERT_MODEL_PATH, XGB_MODEL_PATH)


@app.get("/")
def read_root():
    return {"message": "Phishing Detection API is running."}


@app.post("/predict")
async def get_prediction(email: EmailRequest):
    """
    Predict whether an email is phishing or legitimate.
    """
    try:
        return predictor.predict(email.subject, email.body)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))



# Run Server

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="127.0.0.1", port=8500)