From afe20320852c7c4fc24061628dfd2fd134699075 Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Sat, 20 Jul 2024 09:14:38 -0600
Subject: [PATCH] Implement VAPID key generation

---
 docker/main/requirements-wheels.txt           |  4 ++-
 frigate/api/notification.py                   | 25 +++++++++++++++++--
 frigate/comms/webpush.py                      |  8 +++++-
 .../settings/NotificationsSettingsView.tsx    | 21 ++++++++++++++--
 4 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/docker/main/requirements-wheels.txt b/docker/main/requirements-wheels.txt
index 4b4e13850..028a48925 100644
--- a/docker/main/requirements-wheels.txt
+++ b/docker/main/requirements-wheels.txt
@@ -36,4 +36,6 @@ chromadb == 0.5.0
 # Generative AI
 google-generativeai == 0.6.*
 ollama == 0.2.*
-openai == 1.30.*
\ No newline at end of file
+openai == 1.30.*
+# push notifications
+py-vapid == 1.9.*
\ No newline at end of file
diff --git a/frigate/api/notification.py b/frigate/api/notification.py
index 63cff9d08..76197d436 100644
--- a/frigate/api/notification.py
+++ b/frigate/api/notification.py
@@ -1,14 +1,19 @@
 """Notification apis."""
 
 import logging
+import os
 
 from flask import (
     Blueprint,
+    current_app,
     jsonify,
+    make_response,
     request,
 )
 from peewee import DoesNotExist
+from py_vapid import Vapid01
 
+from frigate.const import CONFIG_DIR
 from frigate.models import User
 
 logger = logging.getLogger(__name__)
@@ -16,6 +21,18 @@ logger = logging.getLogger(__name__)
 NotificationBp = Blueprint("notifications", __name__)
 
 
+@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,
+        )
+
+    key = Vapid01.from_file(os.path.join(CONFIG_DIR, "notifications.pem"))
+    return jsonify(key.public_key), 200
+
+
 @NotificationBp.route("/notifications/register", methods=["POST"])
 def register_notifications():
     username = request.headers.get("remote-user", type=str) or "admin"
@@ -29,6 +46,10 @@ def register_notifications():
         User.update(notification_tokens=User.notification_tokens.append(token)).where(
             User.username == username
         ).execute()
-        return jsonify({"success": True, "message": "Successfully saved token."}), 200
+        return make_response(
+            jsonify({"success": True, "message": "Successfully saved token."}), 200
+        )
     except DoesNotExist:
-        return jsonify({"success": False, "message": "Could not find user."}), 404
+        return make_response(
+            jsonify({"success": False, "message": "Could not find user."}), 404
+        )
diff --git a/frigate/comms/webpush.py b/frigate/comms/webpush.py
index d2cf6ffd0..994881a99 100644
--- a/frigate/comms/webpush.py
+++ b/frigate/comms/webpush.py
@@ -2,10 +2,14 @@
 
 import json
 import logging
+import os
 from typing import Any, Callable
 
+from py_vapid import Vapid01
+
 from frigate.comms.dispatcher import Communicator
 from frigate.config import FrigateConfig
+from frigate.const import CONFIG_DIR
 from frigate.models import User
 
 logger = logging.getLogger(__name__)
@@ -16,7 +20,9 @@ class WebPushClient(Communicator):  # type: ignore[misc]
 
     def __init__(self, config: FrigateConfig) -> None:
         self.config = config
-        # TODO check for VAPID key
+
+        # Pull keys from PEM or generate if they do not exist
+        self.key = Vapid01.from_file(os.path.join(CONFIG_DIR, "notifications.pem"))
 
         self.tokens = []
         self.invalid_tokens = []
diff --git a/web/src/views/settings/NotificationsSettingsView.tsx b/web/src/views/settings/NotificationsSettingsView.tsx
index 765dbc489..b1d6bad4a 100644
--- a/web/src/views/settings/NotificationsSettingsView.tsx
+++ b/web/src/views/settings/NotificationsSettingsView.tsx
@@ -5,11 +5,27 @@ import { Toaster } from "@/components/ui/sonner";
 import { Switch } from "@/components/ui/switch";
 import { FrigateConfig } from "@/types/frigateConfig";
 import axios from "axios";
+import { useEffect, useState } from "react";
 import useSWR from "swr";
 
+const NOTIFICATION_SERVICE_WORKER = "notifications-worker.ts";
+
 export default function NotificationView() {
   const { data: config } = useSWR<FrigateConfig>("config");
 
+  // notification state
+
+  const [notificationsSubscribed, setNotificationsSubscribed] =
+    useState<boolean>();
+
+  useEffect(() => {
+    navigator.serviceWorker
+      .getRegistration(NOTIFICATION_SERVICE_WORKER)
+      .then((worker) => {
+        setNotificationsSubscribed(worker != null);
+      });
+  }, []);
+
   return (
     <>
       <div className="flex size-full flex-col md:flex-row">
@@ -48,12 +64,13 @@ export default function NotificationView() {
                   // TODO make the notifications button show enable / disable depending on current state
                 }
                 <Button
+                  disabled={notificationsSubscribed == undefined}
                   onClick={() => {
                     Notification.requestPermission().then((permission) => {
                       console.log("notification permissions are ", permission);
                       if (permission === "granted") {
                         navigator.serviceWorker
-                          .register("notifications-worker.ts")
+                          .register(NOTIFICATION_SERVICE_WORKER)
                           .then((registration) => {
                             registration.pushManager
                               .subscribe()
@@ -68,7 +85,7 @@ export default function NotificationView() {
                     });
                   }}
                 >
-                  Enable Notifications
+                  {`${notificationsSubscribed ? "Disable" : "Enable"} Notifications`}
                 </Button>
               </div>
             </div>