239 lines
8.4 KiB
Python
239 lines
8.4 KiB
Python
"""Browser push notification service using VAPID/pywebpush."""
|
|
import asyncio
|
|
import base64
|
|
import json
|
|
import logging
|
|
from datetime import datetime, timezone
|
|
from functools import lru_cache
|
|
from pathlib import Path
|
|
|
|
from sqlalchemy import delete, select
|
|
|
|
from fabledassistant.config import Config
|
|
from fabledassistant.models import async_session
|
|
from fabledassistant.models.push_subscription import PushSubscription
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Persisted alongside other app data so keys survive container restarts.
|
|
_VAPID_KEYS_FILE = Path(Config.IMAGE_CACHE_DIR).parent / "vapid_keys.json"
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _get_webpush():
|
|
"""Lazy import to avoid startup errors if pywebpush is not installed."""
|
|
try:
|
|
from pywebpush import webpush
|
|
return webpush
|
|
except ImportError:
|
|
return None
|
|
|
|
|
|
def vapid_enabled() -> bool:
|
|
return bool(Config.VAPID_PRIVATE_KEY and Config.VAPID_PUBLIC_KEY)
|
|
|
|
|
|
def ensure_vapid_keys() -> None:
|
|
"""Load or auto-generate VAPID keys, storing them in the data volume.
|
|
|
|
Called once at startup. If VAPID_PRIVATE_KEY / VAPID_PUBLIC_KEY env vars
|
|
are already set they take precedence and nothing is written to disk.
|
|
"""
|
|
if vapid_enabled():
|
|
logger.info("VAPID keys loaded from environment variables")
|
|
return
|
|
|
|
# Try to load previously generated keys from the data volume.
|
|
if _VAPID_KEYS_FILE.exists():
|
|
try:
|
|
data = json.loads(_VAPID_KEYS_FILE.read_text())
|
|
Config.VAPID_PRIVATE_KEY = data["private_key"]
|
|
Config.VAPID_PUBLIC_KEY = data["public_key"]
|
|
logger.info("VAPID keys loaded from %s", _VAPID_KEYS_FILE)
|
|
return
|
|
except Exception:
|
|
logger.warning("Failed to load VAPID keys from %s — regenerating", _VAPID_KEYS_FILE, exc_info=True)
|
|
|
|
# Generate a fresh key pair.
|
|
try:
|
|
from cryptography.hazmat.primitives import serialization
|
|
from py_vapid import Vapid01
|
|
|
|
v = Vapid01()
|
|
v.generate_keys()
|
|
|
|
# pywebpush expects the private key as a base64url-encoded DER blob
|
|
# (passed to Vapid.from_string → from_der), NOT a PEM string.
|
|
private_der = v.private_key.private_bytes(
|
|
serialization.Encoding.DER,
|
|
serialization.PrivateFormat.TraditionalOpenSSL,
|
|
serialization.NoEncryption(),
|
|
)
|
|
private_b64 = base64.urlsafe_b64encode(private_der).rstrip(b"=").decode()
|
|
|
|
pub_bytes = v.public_key.public_bytes(
|
|
serialization.Encoding.X962,
|
|
serialization.PublicFormat.UncompressedPoint,
|
|
)
|
|
public_b64 = base64.urlsafe_b64encode(pub_bytes).rstrip(b"=").decode()
|
|
|
|
# Persist so they survive container restarts.
|
|
_VAPID_KEYS_FILE.parent.mkdir(parents=True, exist_ok=True)
|
|
_VAPID_KEYS_FILE.write_text(json.dumps({"private_key": private_b64, "public_key": public_b64}))
|
|
|
|
Config.VAPID_PRIVATE_KEY = private_b64
|
|
Config.VAPID_PUBLIC_KEY = public_b64
|
|
logger.info("Generated new VAPID keys and saved to %s", _VAPID_KEYS_FILE)
|
|
except Exception:
|
|
logger.warning("Failed to generate VAPID keys — push notifications will be unavailable", exc_info=True)
|
|
|
|
|
|
async def regenerate_vapid_keys() -> bool:
|
|
"""Delete existing VAPID keys, clear all push subscriptions, and generate a fresh pair.
|
|
|
|
All existing browser subscriptions are invalidated when keys rotate, so they
|
|
must be cleared — users will need to re-enable notifications.
|
|
"""
|
|
if _VAPID_KEYS_FILE.exists():
|
|
_VAPID_KEYS_FILE.unlink()
|
|
Config.VAPID_PRIVATE_KEY = ""
|
|
Config.VAPID_PUBLIC_KEY = ""
|
|
|
|
# Clear all push subscriptions — they are bound to the old public key.
|
|
async with async_session() as session:
|
|
await session.execute(delete(PushSubscription))
|
|
await session.commit()
|
|
|
|
ensure_vapid_keys()
|
|
enabled = vapid_enabled()
|
|
if enabled:
|
|
logger.info("VAPID keys regenerated successfully")
|
|
else:
|
|
logger.error("VAPID key regeneration failed")
|
|
return enabled
|
|
|
|
|
|
async def save_subscription(user_id: int, subscription_json: dict, user_agent: str | None = None) -> PushSubscription:
|
|
"""Upsert a push subscription by endpoint."""
|
|
endpoint = subscription_json.get("endpoint", "")
|
|
keys = subscription_json.get("keys", {})
|
|
p256dh = keys.get("p256dh", "")
|
|
auth = keys.get("auth", "")
|
|
|
|
async with async_session() as session:
|
|
result = await session.execute(
|
|
select(PushSubscription).where(
|
|
PushSubscription.user_id == user_id,
|
|
PushSubscription.endpoint == endpoint,
|
|
)
|
|
)
|
|
existing = result.scalars().first()
|
|
if existing:
|
|
existing.p256dh = p256dh
|
|
existing.auth = auth
|
|
existing.user_agent = user_agent
|
|
existing.last_used = datetime.now(timezone.utc)
|
|
await session.commit()
|
|
await session.refresh(existing)
|
|
return existing
|
|
else:
|
|
sub = PushSubscription(
|
|
user_id=user_id,
|
|
endpoint=endpoint,
|
|
p256dh=p256dh,
|
|
auth=auth,
|
|
user_agent=user_agent,
|
|
)
|
|
session.add(sub)
|
|
await session.commit()
|
|
await session.refresh(sub)
|
|
return sub
|
|
|
|
|
|
async def delete_subscription(user_id: int, endpoint: str) -> None:
|
|
async with async_session() as session:
|
|
await session.execute(
|
|
delete(PushSubscription).where(
|
|
PushSubscription.user_id == user_id,
|
|
PushSubscription.endpoint == endpoint,
|
|
)
|
|
)
|
|
await session.commit()
|
|
|
|
|
|
async def _remove_expired_subscription(user_id: int, endpoint: str) -> None:
|
|
"""Remove a subscription that returned 410 Gone."""
|
|
try:
|
|
await delete_subscription(user_id, endpoint)
|
|
logger.info("Removed expired push subscription for user %d", user_id)
|
|
except Exception:
|
|
logger.warning("Failed to remove expired subscription", exc_info=True)
|
|
|
|
|
|
async def send_push_notification(
|
|
user_id: int,
|
|
title: str,
|
|
body: str,
|
|
url: str = "/",
|
|
) -> None:
|
|
"""Send a push notification to all subscriptions for a user.
|
|
|
|
Fire-and-forget — wrap in asyncio.create_task() at call site.
|
|
"""
|
|
if not vapid_enabled():
|
|
logger.debug("VAPID not configured, skipping push notification")
|
|
return
|
|
|
|
webpush = _get_webpush()
|
|
if webpush is None:
|
|
logger.warning("pywebpush not installed, cannot send push notifications")
|
|
return
|
|
|
|
async with async_session() as session:
|
|
result = await session.execute(
|
|
select(PushSubscription).where(PushSubscription.user_id == user_id)
|
|
)
|
|
subscriptions = list(result.scalars().all())
|
|
|
|
if not subscriptions:
|
|
return
|
|
|
|
payload = json.dumps({"title": title, "body": body, "url": url})
|
|
vapid_claims = {
|
|
"sub": Config.VAPID_CLAIMS_SUB,
|
|
}
|
|
|
|
def _send_sync(sub: PushSubscription) -> tuple[int, str]:
|
|
"""Synchronous send — runs in executor."""
|
|
try:
|
|
subscription_info = {
|
|
"endpoint": sub.endpoint,
|
|
"keys": {"p256dh": sub.p256dh, "auth": sub.auth},
|
|
}
|
|
response = webpush(
|
|
subscription_info=subscription_info,
|
|
data=payload,
|
|
vapid_private_key=Config.VAPID_PRIVATE_KEY,
|
|
vapid_claims=vapid_claims,
|
|
)
|
|
return response.status_code, sub.endpoint
|
|
except Exception as e:
|
|
logger.error("Push send failed for sub %d: %s", sub.id, e, exc_info=True)
|
|
return 0, sub.endpoint
|
|
|
|
loop = asyncio.get_event_loop()
|
|
tasks = [loop.run_in_executor(None, _send_sync, sub) for sub in subscriptions]
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
for sub, result in zip(subscriptions, results):
|
|
if isinstance(result, Exception):
|
|
logger.error("Push gather exception for user %d: %s", user_id, result, exc_info=result)
|
|
continue
|
|
status_code, endpoint = result
|
|
if status_code in (200, 201):
|
|
logger.info("Push notification sent to sub %d (status %d)", sub.id, status_code)
|
|
elif status_code == 410:
|
|
asyncio.create_task(_remove_expired_subscription(user_id, endpoint))
|
|
elif status_code:
|
|
logger.error("Push returned unexpected status %d for user %d sub %d", status_code, user_id, sub.id)
|