notification support

ktvmanager/account_checker.py

@@ -2,6 +2,9 @@ import os
 import sys
 from dotenv import load_dotenv
 import mysql.connector
+from datetime import datetime, timedelta
+from routes.api import send_notification
+from ktvmanager.lib.database import get_push_subscriptions, _execute_query
 
 # Add the project root to the Python path
 project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
@@ -30,6 +33,43 @@ def get_all_accounts(db_connection: MySQLConnection) -> List[Dict[str, Any]]:
     return accounts
 
 
+def send_expiry_notifications() -> None:
+    """
+    Sends notifications to users with accounts expiring in the next 30 days.
+    """
+    now = datetime.now()
+    thirty_days_later = now + timedelta(days=30)
+    now_timestamp = int(now.timestamp())
+    thirty_days_later_timestamp = int(thirty_days_later.timestamp())
+
+    query = """
+    SELECT u.id as user_id, ua.username, ua.expiaryDate
+    FROM users u
+    JOIN userAccounts ua ON u.id = ua.userID
+    WHERE ua.expiaryDate BETWEEN %s AND %s
+    """
+    expiring_accounts = _execute_query(query, (now_timestamp, thirty_days_later_timestamp))
+
+    for account in expiring_accounts:
+        user_id = account['user_id']
+        subscriptions = get_push_subscriptions(user_id)
+        for sub in subscriptions:
+            # Check if a notification has been sent recently
+            last_notified_query = "SELECT last_notified FROM push_subscriptions WHERE id = %s"
+            last_notified_result = _execute_query(last_notified_query, (sub['id'],))
+            last_notified = last_notified_result[0]['last_notified'] if last_notified_result and last_notified_result[0]['last_notified'] else None
+
+            if last_notified and last_notified > now - timedelta(days=1):
+                continue
+
+            message = f"Your account {account['username']} is due to expire on {datetime.fromtimestamp(account['expiaryDate']).strftime('%d-%m-%Y')}."
+            send_notification(sub['subscription_json'], message)
+
+            # Update the last notified timestamp
+            update_last_notified_query = "UPDATE push_subscriptions SET last_notified = %s WHERE id = %s"
+            _execute_query(update_last_notified_query, (now, sub['id']))
+
+
 def main() -> None:
     """
     Checks the validity of all accounts in the database against available stream URLs.

ktvmanager/config.py

@@ -14,6 +14,16 @@ class Config:
     DATABASE = os.getenv("DATABASE")
     DBPORT = os.getenv("DBPORT")
     STREAM_URLS = ["http://example.com", "http://example.org"]
+    VAPID_PRIVATE_KEY = """-----BEGIN PRIVATE KEY-----
+MIGHAgEAMBMGByqGSM49AgEGCCqGSM49AwEHBG0wawIBAQQg6vkDnUOnpMUZ+DAv
+gEge20aPDmffv1rYTADnaNP5NvGhRANCAATZvXvlV0QyvzvgOdsEMSt07n5qgbBn
+ICQ0s1x364rGswAcVVJuu8q5XgZQrBLk/lkhQBcyyuuAjc4OvJLADqEk
+-----END PRIVATE KEY-----"""
+    VAPID_PUBLIC_KEY = """-----BEGIN PUBLIC KEY-----
+MFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE2b175VdEMr874DnbBDErdO5+aoGw
+ZyAkNLNcd+uKxrMAHFVSbrvKuV4GUKwS5P5ZIUAXMsrrgI3ODrySwA6hJA==
+-----END PUBLIC KEY-----"""
+    SECRET_KEY = "a_very_secret_key"
 
 class DevelopmentConfig(Config):
     DEBUG = True

ktvmanager/lib/database.py

@@ -20,6 +20,21 @@ def initialize_db_pool() -> None:
         database=current_app.config["DATABASE"],
         port=current_app.config["DBPORT"],
     )
+    _create_push_subscriptions_table()
+
+
+def _create_push_subscriptions_table() -> None:
+    """Creates the push_subscriptions table if it doesn't exist."""
+    query = """
+    CREATE TABLE IF NOT EXISTS push_subscriptions (
+        id INT AUTO_INCREMENT PRIMARY KEY,
+        user_id INT NOT NULL,
+        subscription_json TEXT NOT NULL,
+        last_notified TIMESTAMP NULL,
+        FOREIGN KEY (user_id) REFERENCES users(id)
+    )
+    """
+    _execute_query(query)
 
 
 def _execute_query(query: str, params: Optional[tuple] = None) -> List[Dict[str, Any]] | Dict[str, int]:
@@ -193,3 +208,28 @@ def delete_account(user_id: int) -> Response:
     params = (data["user"], data["stream"], user_id)
     result = _execute_query(query, params)
     return jsonify(result)
+
+
+def save_push_subscription(user_id: int, subscription_json: str) -> None:
+    """Saves a push subscription to the database.
+
+    Args:
+        user_id: The ID of the user.
+        subscription_json: The push subscription information as a JSON string.
+    """
+    query = "INSERT INTO push_subscriptions (user_id, subscription_json) VALUES (%s, %s)"
+    params = (user_id, subscription_json)
+    _execute_query(query, params)
+
+
+def get_push_subscriptions(user_id: int) -> List[Dict[str, Any]]:
+    """Retrieves all push subscriptions for a given user ID.
+
+    Args:
+        user_id: The ID of the user.
+
+    Returns:
+        A list of push subscriptions.
+    """
+    query = "SELECT * FROM push_subscriptions WHERE user_id = %s"
+    return _execute_query(query, (user_id,))

requirements.txt

@@ -33,3 +33,9 @@ typing_extensions==4.14.1
 urllib3==2.4.0
 Werkzeug==3.1.3
 bump-my-version
+requests
+mysql-connector-python
+python-dotenv
+python-dotenv
+pywebpush==1.13.0
+stem==1.8.2

routes/api.py

@@ -1,4 +1,4 @@
-from flask import Blueprint, jsonify, Response
+from flask import Blueprint, jsonify, Response, request, current_app
 from ktvmanager.lib.database import (
     get_user_accounts,
     get_stream_names,
@@ -6,11 +6,15 @@ from ktvmanager.lib.database import (
     add_account,
     delete_account,
     get_user_id_from_username,
+    save_push_subscription,
+    get_push_subscriptions,
 )
 from ktvmanager.lib.get_urls import get_latest_urls_from_dns
 from ktvmanager.lib.auth import requires_basic_auth, check_login
 from ktvmanager.lib.checker import validate_account
 from typing import Tuple
+import json
+from pywebpush import webpush, WebPushException
 
 api_blueprint = Blueprint("api", __name__)
 
@@ -138,3 +142,60 @@ def login_route(username: str, password: str) -> Response:
         A Flask JSON response with the result of the login attempt.
     """
     return check_login(username, password)
+
+
+@api_blueprint.route("/save-subscription", methods=["POST"])
+@requires_basic_auth
+def save_subscription(username: str, password: str) -> Response:
+    """Saves a push notification subscription.
+
+    Args:
+        username: The username of the user.
+        password: The password of the user.
+
+    Returns:
+        A Flask JSON response.
+    """
+    user_id = get_user_id_from_username(username)
+    if not user_id:
+        return jsonify({"message": "User not found"}), 404
+
+    subscription_data = request.get_json()
+    if not subscription_data:
+        return jsonify({"message": "No subscription data provided"}), 400
+
+    save_push_subscription(user_id, json.dumps(subscription_data))
+    return jsonify({"message": "Subscription saved."})
+
+
+def send_notification(subscription_info, message_body):
+    try:
+        webpush(
+            subscription_info=subscription_info,
+            data=message_body,
+            vapid_private_key=current_app.config["VAPID_PRIVATE_KEY"],
+            vapid_claims={"sub": "mailto:your-email@example.com"},
+        )
+    except WebPushException as ex:
+        print(f"Web push error: {ex}")
+        # You might want to remove the subscription if it's invalid
+        if ex.response and ex.response.status_code == 410:
+            print("Subscription is no longer valid, removing from DB.")
+            # Add logic to remove the subscription from your database
+
+
+@api_blueprint.route("/send-expiry-notifications", methods=["POST"])
+@requires_basic_auth
+def send_expiry_notifications_route(username: str, password: str) -> Response:
+    """Triggers the sending of expiry notifications.
+
+    Args:
+        username: The username of the user.
+        password: The password of the user.
+
+    Returns:
+        A Flask JSON response.
+    """
+    from ktvmanager.account_checker import send_expiry_notifications
+    send_expiry_notifications()
+    return jsonify({"message": "Expiry notifications sent."})