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