py-push-server/app.py

from cryptography.hazmat.primitives import serialization
from flask import Flask, request, jsonify
from models import db, VAPIDKey, Subscription
from pywebpush import webpush, WebPushException

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec

import binascii
import base64
import json
import os
import re

def create_app(config_name):
    app = Flask(__name__)
    app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///data/webpush.db'
    db.init_app(app)


    def generate_and_save_vapid_keys():
        try:
            private_key = ec.generate_private_key(ec.SECP256R1(), default_backend())
            public_key_bytes = private_key.public_key().public_bytes(
                encoding=serialization.Encoding.X962,
                format=serialization.PublicFormat.UncompressedPoint
            )
            public_key_base64 = base64.b64encode(public_key_bytes).decode()
            
            private_key_hex = binascii.hexlify(private_key.private_bytes(
                encoding=serialization.Encoding.DER,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )).decode()
            
        except Exception as e:
            print(f"Error generating VAPID keys: {e}")

        key = VAPIDKey(public_key=public_key_base64, private_key=private_key_hex)
        db.session.add(key)
        db.session.commit()


    def initialize():
        if not VAPIDKey.query.first():
            generate_and_save_vapid_keys()


    def send_push_notification(subscription_info, message, vapid_key):
        result = True
        try:
            private_key_hex = vapid_key.private_key
            private_key_der = bytes.fromhex(private_key_hex)

            private_key_base64 = base64.b64encode(private_key_der).decode()
            
            webpush(
                subscription_info=subscription_info,
                data=json.dumps(message),
                vapid_private_key=private_key_base64,
                vapid_claims={
                    "sub": "mailto:your-email@example.com"
                }
            )
            
        except Exception as e:
            result = False
            print(f"Failed to send notification: {e}")
            
        return result

            
    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/clear_subscriptions', methods=['POST'])
    def clear_subscriptions():
        try:
            Subscription.query.delete()
            return jsonify(message='Subscriptions cleared')
        
        except Exception as e:
            return jsonify(error=f'Error clearing subscriptions: {str(e)}'), 500

        
    @app.route('/web-push/regenerate_vapid_keys', methods=['POST'])
    def regenerate_vapid_keys():
        try:
            # Generate new VAPID keys
            VAPIDKey.query.delete()
            generate_and_save_vapid_keys()
            return jsonify(message='VAPID keys regenerated successfully')
        
        except Exception as e:
            return jsonify(error=f'Error regenerating VAPID keys: {str(e)}'), 500

    
    @app.route('/web-push/vapid', methods=['GET'])
    def get_vapid():
        key = VAPIDKey.query.first()
        return jsonify(vapidKey=key.public_key)


    @app.route('/web-push/subscribe', methods=['POST'])
    def subscribe():
        content = request.json
        vapid_key = VAPIDKey.query.first()

        if not vapid_key:
            return jsonify(success=False, error="No VAPID keys available"), 500

        subscription = Subscription(endpoint=content['endpoint'],
                                    p256dh=content['keys']['p256dh'],
                                    auth=content['keys']['auth'],
                                    vapid_key_id=vapid_key.id)
        db.session.add(subscription)
        db.session.commit()

        subscription_info = {
            "endpoint": subscription.endpoint,
            "keys": {
                "p256dh": subscription.p256dh,
                "auth": subscription.auth
            }
        }

        message = {"title": "Subscription Successful", "body": "Thank you for subscribing!"}
        success = send_push_notification(subscription_info, message, vapid_key)
        
        return jsonify(success=success, message=vapid_key.private_key)


    @app.route('/web-push/unsubscribe', methods=['DELETE'])
    def unsubscribe():
        content = request.json
        endpoint = content['endpoint']
        subscription = Subscription.query.filter_by(endpoint=endpoint).first()

        if subscription:
            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

    with app.app_context():
        initialize()
        
    return app