From 023f1431f28edbffee6808de74d4dcf6b3979435 Mon Sep 17 00:00:00 2001
From: Rui Alves <alvesrui951@gmail.com>
Date: Sun, 15 Sep 2024 17:08:41 +0100
Subject: [PATCH] Convert all notifications API endpoints to FastAPI

---
 frigate/api/app.py          |  2 --
 frigate/api/defs/tags.py    |  1 +
 frigate/api/fastapi_app.py  |  3 ++-
 frigate/api/notification.py | 53 ++++++++++++++++++-------------------
 4 files changed, 29 insertions(+), 30 deletions(-)

diff --git a/frigate/api/app.py b/frigate/api/app.py
index b3baec8c3..ee518e640 100644
--- a/frigate/api/app.py
+++ b/frigate/api/app.py
@@ -23,7 +23,6 @@ from frigate.api.auth import AuthBp, get_jwt_secret, limiter
 from frigate.api.defs.tags import Tags
 from frigate.api.event import EventBp
 from frigate.api.export import ExportBp
-from frigate.api.notification import NotificationBp
 from frigate.api.review import ReviewBp
 from frigate.config import FrigateConfig
 from frigate.const import CONFIG_DIR
@@ -50,7 +49,6 @@ bp.register_blueprint(EventBp)
 bp.register_blueprint(ExportBp)
 bp.register_blueprint(ReviewBp)
 bp.register_blueprint(AuthBp)
-bp.register_blueprint(NotificationBp)
 
 router = APIRouter()
 
diff --git a/frigate/api/defs/tags.py b/frigate/api/defs/tags.py
index 8a47de1c0..27cb6cb74 100644
--- a/frigate/api/defs/tags.py
+++ b/frigate/api/defs/tags.py
@@ -6,3 +6,4 @@ class Tags(Enum):
     preview = "Preview"
     logs = "Logs"
     media = "Media"
+    notifications = "Notifications"
diff --git a/frigate/api/fastapi_app.py b/frigate/api/fastapi_app.py
index 9dd7b0610..3ec96dfee 100644
--- a/frigate/api/fastapi_app.py
+++ b/frigate/api/fastapi_app.py
@@ -3,7 +3,7 @@ import logging
 from fastapi import FastAPI
 
 from frigate.api import app as main_app
-from frigate.api import media, preview
+from frigate.api import media, notification, preview
 from frigate.plus import PlusApi
 from frigate.ptz.onvif import OnvifController
 from frigate.stats.emitter import StatsEmitter
@@ -29,6 +29,7 @@ def create_fastapi_app(
     app.include_router(main_app.router)
     app.include_router(media.router)
     app.include_router(preview.router)
+    app.include_router(notification.router)
     # App Properties
     app.frigate_config = frigate_config
     app.detected_frames_processor = detected_frames_processor
diff --git a/frigate/api/notification.py b/frigate/api/notification.py
index 2fe61882f..907b4f350 100644
--- a/frigate/api/notification.py
+++ b/frigate/api/notification.py
@@ -4,62 +4,61 @@ import logging
 import os
 
 from cryptography.hazmat.primitives import serialization
-from flask import (
-    Blueprint,
-    current_app,
-    jsonify,
-    make_response,
-    request,
-)
+from fastapi import APIRouter, Request
+from fastapi.responses import JSONResponse
 from peewee import DoesNotExist
 from py_vapid import Vapid01, utils
 
+from frigate.api.defs.tags import Tags
 from frigate.const import CONFIG_DIR
 from frigate.models import User
 
 logger = logging.getLogger(__name__)
 
-NotificationBp = Blueprint("notifications", __name__)
+router = APIRouter(tags=[Tags.notifications])
 
 
-@NotificationBp.route("/notifications/pubkey", methods=["GET"])
-def get_vapid_pub_key():
-    if not current_app.frigate_config.notifications.enabled:
-        return make_response(
-            jsonify({"success": False, "message": "Notifications are not enabled."}),
-            400,
+@router.get("/notifications/pubkey")
+def get_vapid_pub_key(request: Request):
+    if not request.app.frigate_config.notifications.enabled:
+        return JSONResponse(
+            content=({"success": False, "message": "Notifications are not enabled."}),
+            status_code=400,
         )
 
     key = Vapid01.from_file(os.path.join(CONFIG_DIR, "notifications.pem"))
     raw_pub = key.public_key.public_bytes(
         serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
     )
-    return jsonify(utils.b64urlencode(raw_pub)), 200
+    return JSONResponse(content=utils.b64urlencode(raw_pub), status_code=200)
 
 
-@NotificationBp.route("/notifications/register", methods=["POST"])
-def register_notifications():
-    if current_app.frigate_config.auth.enabled:
-        username = request.headers.get("remote-user", type=str) or "admin"
+@router.post("/notifications/register")
+def register_notifications(request: Request, body: dict = None):
+    if request.app.frigate_config.auth.enabled:
+        username = request.headers.get("remote-user") or "admin"
     else:
         username = "admin"
 
-    json: dict[str, any] = request.get_json(silent=True) or {}
+    json: dict[str, any] = body or {}
     sub = json.get("sub")
 
     if not sub:
-        return jsonify(
-            {"success": False, "message": "Subscription must be provided."}
-        ), 400
+        return JSONResponse(
+            content={"success": False, "message": "Subscription must be provided."},
+            status_code=400,
+        )
 
     try:
         User.update(notification_tokens=User.notification_tokens.append(sub)).where(
             User.username == username
         ).execute()
-        return make_response(
-            jsonify({"success": True, "message": "Successfully saved token."}), 200
+        return JSONResponse(
+            content=({"success": True, "message": "Successfully saved token."}),
+            status_code=200,
         )
     except DoesNotExist:
-        return make_response(
-            jsonify({"success": False, "message": "Could not find user."}), 404
+        return JSONResponse(
+            content=({"success": False, "message": "Could not find user."}),
+            status_code=404,
         )