Add ReportMessagesModel

app.py

@@ -4,8 +4,11 @@ import sys
 import re
 
 import httpx
-from fastapi import FastAPI
+
+from fastapi import FastAPI, Request, status
 from fastapi.responses import JSONResponse, FileResponse
+from fastapi.exceptions import RequestValidationError
+
 from transformers import pipeline
 from phishing_datasets import submit_entry
 from url_tools import extract_urls, resolve_short_url, extract_domain_from_url
@@ -13,7 +16,7 @@ from urlscan_client import UrlscanClient
 import requests
 from mnemonic_attack import find_confusable_brand
 
-from models.models import MessageModel, QueryModel, AppModel, InputModel, OutputModel, ReportModel
+from models.models import MessageModel, QueryModel, AppModel, InputModel, OutputModel, ReportMessagesModel, ReportInputModel
 from models.enums import ActionModel, SubActionModel
 
 app = FastAPI()
@@ -168,5 +171,20 @@ def predict(model: InputModel) -> OutputModel:
     return OutputModel(action=action, sub_action=SubActionModel.NONE)
 
 @app.post("/report")
-def report(model: ReportModel):
-    submit_entry(model.sender, model.message)
+def report(model: ReportInputModel):
+    logging.info(f"[REPORT] {model.classification.messages}")
+
+    for msg in model.classification.messages:
+        logging.info(
+            f"[REPORT] {msg.timestamp=} {msg.sender=} {msg.message=}"
+        )
+
+    return {"status": "received", "count": len(model.classification.messages)}
+
+
+@app.exception_handler(RequestValidationError)
+async def validation_exception_handler(request: Request, exc: RequestValidationError):
+	exc_str = f'{exc}'.replace('\n', ' ').replace('   ', ' ')
+	logging.error(f"{request}: {exc_str}")
+	content = {'status_code': 10422, 'message': exc_str, 'data': None}
+	return JSONResponse(content=content, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY)

models/models.py

@@ -1,5 +1,5 @@
 from pydantic import BaseModel
-from enum import Enum
+from models.enums import ActionModel, SubActionModel
 
 class MessageModel(BaseModel):
     text: str
@@ -17,9 +17,18 @@ class InputModel(BaseModel):
     app: AppModel
 
 class OutputModel(BaseModel):
-    action: 'ActionModel'
-    sub_action: 'SubActionModel'
+    action: ActionModel
+    sub_action: SubActionModel
 
-class ReportModel(BaseModel):
+class ReportMessageModel(BaseModel):
+    timestamp: int
     sender: str
-    message: str 
+    message: str
+
+class ReportMessagesModel(BaseModel):
+    messages: list[ReportMessageModel]
+
+class ReportInputModel(BaseModel):
+    _version: int
+    classification: ReportMessagesModel
+    app: AppModel