app/backend/fluksio/flow/webpush.py

"""Web Push: an alert reaching a phone that has this instance installed.

The browser hands us a subscription — an endpoint URL at its own push service,
plus two keys — and from then on the engine can wake that device without it
holding a connection open. Two pieces of crypto are involved and both are
specified: the payload is encrypted to the subscription's keys (RFC 8291,
aes128gcm) so the push service carries something it cannot read, and the
request is signed with this instance's own keypair (VAPID, RFC 8292) so the
service knows who is sending.

Only ``http-ece`` is new here; the signing is `pyjwt` and the request is
`httpx`, both of which the engine already carries. ``pywebpush`` does all of
this in one call, but brings `requests` *and* `aiohttp` with it — two more HTTP
stacks on a machine that may well be a Raspberry Pi.

The keypair is this instance's identity to the push services and lives with
the subscriptions in one file. Losing it means every browser has to subscribe
again; it is regenerated on the spot if the file goes missing.
"""

from __future__ import annotations

import asyncio
import base64
import json
import logging
import threading
import time
from pathlib import Path
from typing import Any
from urllib.parse import urlparse

import http_ece
import httpx
import jwt
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from pydantic import BaseModel, Field

from fluksio.core.config import settings

logger = logging.getLogger(__name__)

#: How long the push service keeps a notification for a device that is offline.
TTL_S = 86400
#: A VAPID token is good for twelve hours; well under the 24 the RFC allows.
TOKEN_LIFETIME_S = 12 * 60 * 60

# ponytail: one lock over the whole file. Subscriptions change when somebody
# presses a button, and pruning happens once per send — per-record locking
# would be machinery for a file with single-digit rows in it.
_lock = threading.Lock()


class Subscription(BaseModel):
    """What a browser handed us. Opaque apart from the endpoint's host."""

    endpoint: str
    #: The subscription's public key and its authentication secret, both
    #: base64url as `PushSubscription.toJSON()` writes them.
    keys: dict[str, str] = Field(default_factory=dict)


class _Store(BaseModel):
    #: This instance's VAPID private key: the raw P-256 scalar, base64url.
    private_key: str = ""
    subscriptions: list[Subscription] = Field(default_factory=list)


def _b64url(raw: bytes) -> str:
    return base64.urlsafe_b64encode(raw).rstrip(b"=").decode()


def _unb64url(text: str) -> bytes:
    # A browser writes these unpadded; base64 refuses them that way.
    return base64.urlsafe_b64decode(text + "=" * (-len(text) % 4))


def _path() -> Path:
    return settings.WEBPUSH_FILE


def _read() -> _Store:
    """Blocking. Call under `_lock` when a write follows."""
    path = _path()
    if not path.exists():
        return _Store()
    try:
        return _Store.model_validate_json(path.read_text())
    except Exception:
        # A file that no longer parses must not stop the engine alerting
        # through its other channels.
        logger.error("The web push store is unreadable; starting a new one")
        return _Store()


def _write(store: _Store) -> None:
    path = _path()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(store.model_dump_json(indent=2))


def _private_key(store: _Store) -> ec.EllipticCurvePrivateKey:
    return ec.derive_private_key(
        int.from_bytes(_unb64url(store.private_key), "big"), ec.SECP256R1()
    )


def _public_bytes(key: ec.EllipticCurvePrivateKey) -> bytes:
    return key.public_key().public_bytes(
        serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
    )


def public_key() -> str:
    """This instance's VAPID public key, generating the pair on first ask.

    Base64url of the uncompressed point, which is the shape
    `pushManager.subscribe` wants for `applicationServerKey`. Blocking.
    """
    with _lock:
        store = _read()
        if not store.private_key:
            key = ec.generate_private_key(ec.SECP256R1())
            store.private_key = _b64url(
                key.private_numbers().private_value.to_bytes(32, "big")
            )
            _write(store)
            logger.info("Generated this instance's web push keypair")
        return _b64url(_public_bytes(_private_key(store)))


def add_subscription(subscription: Subscription) -> None:
    """Remember a browser, replacing what we knew about that endpoint. Blocking."""
    with _lock:
        store = _read()
        store.subscriptions = [
            s for s in store.subscriptions if s.endpoint != subscription.endpoint
        ]
        store.subscriptions.append(subscription)
        _write(store)


def remove_subscription(endpoint: str) -> None:
    """Blocking."""
    with _lock:
        store = _read()
        store.subscriptions = [s for s in store.subscriptions if s.endpoint != endpoint]
        _write(store)


def subscriptions() -> list[Subscription]:
    """Blocking."""
    with _lock:
        return _read().subscriptions


def _vapid_headers(key: ec.EllipticCurvePrivateKey, endpoint: str) -> dict[str, str]:
    """Prove to the push service which instance is sending (RFC 8292)."""
    origin = urlparse(endpoint)
    token = jwt.encode(
        {
            "aud": f"{origin.scheme}://{origin.netloc}",
            "exp": int(time.time()) + TOKEN_LIFETIME_S,
            # Who to shout at about this traffic. A push service will not read
            # it unless something has gone wrong.
            "sub": f"mailto:{settings.EMAILS_FROM_EMAIL or 'admin@fluksio.invalid'}",
        },
        key,
        algorithm="ES256",
    )
    return {"Authorization": f"vapid t={token},k={_b64url(_public_bytes(key))}"}


def _encrypt(payload: bytes, subscription: Subscription) -> bytes:
    """Encrypt to this subscription's keys, so only that browser can read it."""
    # Ephemeral, per message: the whole point of the exchange is that this key
    # and the subscription's agree on a secret nobody else can derive.
    ours = ec.generate_private_key(ec.SECP256R1())
    encrypted: bytes = http_ece.encrypt(
        payload,
        private_key=ours,
        dh=_unb64url(subscription.keys["p256dh"]),
        auth_secret=_unb64url(subscription.keys["auth"]),
        version="aes128gcm",
    )
    return encrypted


async def send_to_all(payload: dict[str, Any]) -> int:
    """Push one payload to every subscribed browser. Returns how many took it.

    A subscription the push service has retired (404/410) is dropped here —
    that is the only signal we get that a browser is never coming back, and
    without acting on it the store grows a dead entry per reinstalled phone.
    """
    # The keypair exists by now — nothing can have subscribed without asking
    # for it — but generating it here costs nothing and keeps this callable.
    await asyncio.to_thread(public_key)
    with _lock:
        store = _read()
    known = store.subscriptions
    if not known:
        return 0
    key = _private_key(store)

    body = json.dumps(payload).encode()
    delivered = 0
    gone: list[str] = []
    async with httpx.AsyncClient(timeout=10) as client:
        for subscription in known:
            try:
                encrypted = await asyncio.to_thread(_encrypt, body, subscription)
                response = await client.post(
                    subscription.endpoint,
                    content=encrypted,
                    headers={
                        "Content-Encoding": "aes128gcm",
                        "Content-Type": "application/octet-stream",
                        "TTL": str(TTL_S),
                        **_vapid_headers(key, subscription.endpoint),
                    },
                )
            except Exception as exc:
                logger.error("Could not push to a subscription: %s", exc)
                continue
            if response.status_code in (404, 410):
                gone.append(subscription.endpoint)
            elif response.is_success:
                delivered += 1
            else:
                logger.error(
                    "A push service refused the notification: %s %s",
                    response.status_code,
                    response.text[:200],
                )

    for endpoint in gone:
        await asyncio.to_thread(remove_subscription, endpoint)
        logger.info("Dropped a subscription the push service has retired")
    return delivered