From 76ee4ef4286e1c745b3c332ef0d7e5ad6df71400 Mon Sep 17 00:00:00 2001
From: Matthew Raymer <matthew.raymer@anomalist.design>
Date: Mon, 2 Oct 2023 09:17:58 -0400
Subject: [PATCH] Some adjustments to my first attempt.  Try tomorow!

---
 app.py           | 32 ++++++++++++++++++++++++++++++--
 requirements.txt |  1 +
 2 files changed, 31 insertions(+), 2 deletions(-)

diff --git a/app.py b/app.py
index 6f25dfd..b79f378 100644
--- a/app.py
+++ b/app.py
@@ -1,10 +1,13 @@
+from cryptography.hazmat.primitives import serialization
 from flask import Flask, request, jsonify
 from models import db, VAPIDKey, Subscription
 from py_vapid import Vapid
 from pywebpush import webpush, WebPushException
 
+import base64
 import json
 import os
+import re
 
 def create_app(config_name):
     app = Flask(__name__)
@@ -45,12 +48,36 @@ def create_app(config_name):
         except WebPushException as e:
             print(f"Failed to send notification: {e}")
 
+    def is_valid_base64_url(s):
+        return re.match(r'^[A-Za-z0-9_-]*$', s) is not None
+
+    def is_valid_server_key(server_key):
+        if not is_valid_base64_url(server_key):
+            return False
+
+        return len(server_key) == 88
 
     @app.route('/web-push/vapid', methods=['GET'])
     def get_vapid():
         key = VAPIDKey.query.first()
-        if key:
-            return jsonify(public_key=key.public_key)
+        pem_bytes = key.private_key.encode("utf-8")
+        private_key = serialization.load_pem_private_key(pem_bytes, password=None)
+        public_key = private_key.public_key()
+        public_key_der = public_key.public_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo
+        )
+        
+        base64_data = base64.urlsafe_b64encode(public_key_der)
+        base64_str = base64_data.decode('utf-8')
+        base64_str = base64_str.rstrip('=')
+
+        if is_valid_server_key(base64_str):
+            return jsonify(vapidKey=base64_str)
+        
+        else:
+            return jsonify(error=f'VAPID server key is not valid. {len(base64_str)} {is_valid_base64_url(base64_str)}'), 422
+        
         return jsonify(error='No VAPID keys found'), 404
 
 
@@ -93,6 +120,7 @@ def create_app(config_name):
             db.session.delete(subscription)
             db.session.commit()
             return jsonify(success=True, message="Subscription deleted successfully")
+        
         else:
             return jsonify(success=False, error="Subscription not found"), 404
 
diff --git a/requirements.txt b/requirements.txt
index bc3eee2..69fcbd3 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+cryptography
 flask>=2.0.0
 flask_sqlalchemy
 py_vapid