binoubinks/ADSP_finalProjectBack

1#!/usr/bin/env python2# encoding: utf-83 4from fastapi import FastAPI, Form, Depends, Request5from fastapi.encoders import jsonable_encoder6from fastapi.responses import JSONResponse7from fastapi.middleware.cors import CORSMiddleware8from pydantic import BaseModel9import pickle10import joblib11import pandas as pd12from sklearn.preprocessing import LabelEncoder13from extraction_features import extract_features 14 15app = FastAPI()16 17# Add CORS middleware18app.add_middleware(19    CORSMiddleware,20    allow_origins=["*"], 21    allow_credentials=True,22    allow_methods=["*"],23    allow_headers=["*"],24)25 26# model_file = open('logistic_regression_model.pkl', 'rb')27# model = pickle.load(model_file, encoding='bytes')28model = joblib.load('logistic_regression_model.pkl')29label_encoders = joblib.load('label_encoders.pkl')30# Columns used in the model31selected_columns = [32    'URLLength', 'Domain', 'DomainLength', 'TLD', 33    'CharContinuationRate', 'TLDLength', 'NoOfSubDomain', 34    'DegitRatioInURL', 'SpacialCharRatioInURL', 'IsHTTPS'35]36 37# Function to manage values for encoding (giving a new number for url which have never been seen)38def safe_transform(encoder, value):39    if value in encoder.classes_:40        return encoder.transform([value])[0]41    else:42        return -1 43 44class Msg(BaseModel):45    msg: str46 47class Req(BaseModel):48    url: str49    50class Resp(BaseModel):51    url: str52    label: str53 54 55@app.get("/")56async def root():57    return {"message": "Hello, Welcome to the final project from Albin Tardivel"}58 59def form_req(url: str = Form(...)):60    return Req(url=str(url))61 62 63@app.get("/path")64async def demo_get():65    return {"message": "This is /path endpoint, use a post request to transform the text to uppercase"}66 67 68@app.post("/path")69async def demo_post(inp: Msg):70    return {"message": inp.msg.upper()}71 72 73@app.get("/path/{path_id}")74async def demo_get_path_id(path_id: int):75    return {"message": f"This is /path/{path_id} endpoint, use post request to retrieve result"}76 77 78@app.get("/predict/{path_id}")79async def predict(path_id: int):80    return {"message":  f"This is /predict/{path_id} endpoint, use post request to retrieve result"}81 82@app.post("/predict")83async def predict(request: Request, requess: Req = Depends(form_req)):84    '''85    Predict if url is phishing or legitimate86    and render the result to the html page87    '''88    url = requess.url89 90    features = extract_features(str(url))91    dataFrame_features = pd.DataFrame([features])92 93    # Apply features encoding (convert everything into int64)94    for column in ['Domain', 'TLD']:95        encoder = label_encoders[column]96        dataFrame_features[column] = dataFrame_features[column].apply(lambda x: safe_transform(encoder, x))97 98    data = dataFrame_features[selected_columns].values.reshape(1, -1)99 100    prediction_proba = model.predict_proba(data)[:, 1]101    threshold = 0.9102    print("prediction_proba:", prediction_proba)103    output = 1 if prediction_proba >= threshold else 0104    105    output_text = "Legitimate" if output == 1 else "Phishing"106 107    # Render index.html with prediction results108    json_compatible_resp_data = jsonable_encoder(Resp(url=requess.url, label=output_text))109    return JSONResponse(content=json_compatible_resp_data)