py-push-server/app.py

from typing import Dict, Tuple, Union, Optional
from flask import Flask, request, jsonify, Response
from models import db, VAPIDKey, Subscription
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from pywebpush import webpush, WebPushException

import base64
import threading
import time

class WebPushService:
    
    def __init__(self, config_name: str) -> None:
        self.app = Flask(__name__)
        self.app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///data/webpush.db'
        db.init_app(self.app)
        with self.app.app_context():
            self._initialize()
        
        self.daily_notification_thread = threading.Thread(target=self._send_daily_notifications)
        self.daily_notification_thread.start()

        
    def _generate_and_save_vapid_keys(self) -> None:
        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_bytes = private_key.private_bytes(
                encoding=serialization.Encoding.DER,
                format=serialization.PrivateFormat.PKCS8,
                encryption_algorithm=serialization.NoEncryption()
            )
            private_key_base64 = base64.b64encode(private_key_bytes).decode()

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

        except Exception as e:
            print(f"Error generating VAPID keys: {e}")

            
    def _initialize(self) -> None:
        if not VAPIDKey.query.first():
            self._generate_and_save_vapid_keys()

            
    def _send_push_notification(self, subscription_info: Dict, message: Dict, vapid_key: VAPIDKey) -> bool:
        try:
            webpush(
                subscription_info=subscription_info,
                data=json.dumps(message),
                vapid_private_key=vapid_key.private_key,
                vapid_claims={"sub": "mailto:admin@yourdomain.com"}
            )
            return True
        
        except WebPushException as ex:
            print(f"Failed to send push notification: {ex}")
            return False

        
    def _send_daily_notifications(self) -> None:
        while True:
            with self.app.app_context():
                all_subscriptions = Subscription.query.all()
                vapid_key = VAPIDKey.query.first()
                message = {"title": "Daily Update", "message": "Here's your daily update!"}
                for subscription in all_subscriptions:
                    subscription_info = {
                        "endpoint": subscription.endpoint,
                        "keys": {
                            "p256dh": subscription.p256dh,
                            "auth": subscription.auth
                        }
                    }
                    self._send_push_notification(subscription_info, message, vapid_key)

            time.sleep(24 * 60 * 60)

    # Route handlers and other methods would go here...

    @app.route('/regenerate_vapid', methods=['POST'])
    def regenerate_vapid(self) -> Tuple[str, int]:
        try:
            with self.app.app_context():
                VAPIDKey.query.delete()
                db.session.commit()
                self._generate_and_save_vapid_keys()
                
            return jsonify(success=True, message="VAPID keys regenerated successfully"), 200
        
        except Exception as e:
            return jsonify(error=f'Error regenerating VAPID keys: {str(e)}'), 500


    @app.route('/get_vapid')
    def get_vapid(self) -> Response:
        key = VAPIDKey.query.first()
        return jsonify(vapidKey=key.public_key)


    @app.route('/subscribe', methods=['POST'])
    def subscribe(self) -> Tuple[str, int]:
        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()

        time.sleep(5)

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

        message = {"title": "Subscription Successful", "message": "Thank you for subscribing!"}
        success = self._send_push_notification(subscription_info, message, vapid_key)

        return jsonify(success=success, message=vapid_key.private_key)

    
    @app.route('/unsubscribe', methods=['POST'])
    def unsubscribe(self) -> Tuple[str, int]:
        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