push.py

"""
Server-side Web Push — "✅ Your result is ready" notifications that arrive even when the customer's tab is in
the background, the phone screen is off, or they are in another app (browsers pause background pages, so a
notification fired from the page itself can be delayed until they return).

* The browser subscribes once (after the customer allows notifications) and the subscription is stored here,
  linked to a random browser id (cookie `cs_cid`) and, when signed in, to the account email.
* When a job finishes, _stream() calls send_to(); the push service (Google/Mozilla/Apple) delivers it to the
  service worker (static/cs-notify-sw.js), which shows the notification.
* Standard Web Push: VAPID (RFC 8292) + aes128gcm payload encryption (RFC 8291), implemented with
  `cryptography` — no other dependency. VAPID keys are created once and kept in the user database
  (so they survive restarts through the USERS_REPO backup), or set VAPID_PRIVATE_KEY yourself.
"""
import base64
import hashlib
import hmac
import json
import os
import struct
import threading
import time
from urllib.parse import urlparse

import auth as A

try:
    from cryptography.hazmat.primitives import hashes, serialization
    from cryptography.hazmat.primitives.asymmetric import ec
    from cryptography.hazmat.primitives.asymmetric.utils import decode_dss_signature
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    ENABLED = os.getenv("PUSH_ENABLED", "1") != "0"
except Exception:                                    # cryptography missing -> page-level notifications only
    ENABLED = False

CONTACT = os.getenv("VAPID_CONTACT", "mailto:" + (os.getenv("ADMIN_EMAILS", "admin@example.com").split(",")[0].strip()
                                                  or "admin@example.com"))
PUBLIC_KEY = ""                                      # base64url, given to the browser as applicationServerKey
_PRIV = None


def b64u(b):
    return base64.urlsafe_b64encode(b).rstrip(b"=").decode()


def ub64u(s):
    s = (s or "").strip()
    return base64.urlsafe_b64decode(s + "=" * (-len(s) % 4))


def _pub_bytes(key):
    return key.public_key().public_bytes(serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint)


def _load_keys():
    """VAPID key pair: env VAPID_PRIVATE_KEY (base64url raw 32 bytes) -> stored in DB -> generate once."""
    global _PRIV, PUBLIC_KEY
    raw = os.getenv("VAPID_PRIVATE_KEY", "").strip()
    if not raw:
        with A._db() as con:
            row = con.execute("SELECT v FROM meta WHERE k='vapid_private'").fetchone()
            if row:
                raw = row[0]
            else:
                key = ec.generate_private_key(ec.SECP256R1())
                raw = b64u(key.private_numbers().private_value.to_bytes(32, "big"))
                con.execute("INSERT INTO meta (k, v) VALUES ('vapid_private', ?)", (raw,))
        A.request_sync()
    _PRIV = ec.derive_private_key(int.from_bytes(ub64u(raw), "big"), ec.SECP256R1())
    PUBLIC_KEY = b64u(_pub_bytes(_PRIV))


def _db():
    con = A._db()
    con.execute("CREATE TABLE IF NOT EXISTS push_subs (endpoint TEXT PRIMARY KEY, cid TEXT, email TEXT, "
                "sub TEXT, ts REAL)")
    return con


if ENABLED:
    try:
        _load_keys()
    except Exception as err:
        print(f"[push] disabled: {err}", flush=True)
        ENABLED = False


# ---------------------------------------------------------------- subscriptions
def save(cid, sub_json, user=None):
    """Store/refresh a browser's push subscription."""
    if not ENABLED or not cid:
        return
    try:
        sub = json.loads(sub_json or "{}")
        ep, keys = sub.get("endpoint", ""), sub.get("keys") or {}
        if not ep.startswith("https://") or not keys.get("p256dh") or not keys.get("auth"):
            return
    except Exception:
        return
    email = (user or {}).get("email")
    STATS["subs_saved"] += 1
    with _db() as con:
        con.execute("INSERT INTO push_subs (endpoint, cid, email, sub, ts) VALUES (?,?,?,?,?) "
                    "ON CONFLICT(endpoint) DO UPDATE SET cid=excluded.cid, email=COALESCE(excluded.email, email), "
                    "sub=excluded.sub, ts=excluded.ts", (ep, cid[:64], email, json.dumps(sub), time.time()))
        con.execute("DELETE FROM push_subs WHERE ts < ?", (time.time() - 90 * 86400,))


def _targets(cid, user):
    email = (user or {}).get("email")
    with _db() as con:
        rows = con.execute("SELECT endpoint, sub FROM push_subs WHERE cid=? OR (email IS NOT NULL AND email=?)",
                           (cid or "-", email or "-")).fetchall()
    return rows


# ---------------------------------------------------------------- RFC 8291 / 8292
def _hkdf(salt, ikm, info, length):
    prk = hmac.new(salt, ikm, hashlib.sha256).digest()
    return hmac.new(prk, info + b"\x01", hashlib.sha256).digest()[:length]


def encrypt(sub, payload, _eph=None, _salt=None):
    """aes128gcm body for one subscription (RFC 8291). _eph/_salt only for the RFC test vector."""
    ua_pub = ub64u(sub["keys"]["p256dh"])
    auth = ub64u(sub["keys"]["auth"])
    eph = _eph or ec.generate_private_key(ec.SECP256R1())
    as_pub = _pub_bytes(eph)
    shared = eph.exchange(ec.ECDH(), ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), ua_pub))
    ikm = _hkdf(auth, shared, b"WebPush: info\x00" + ua_pub + as_pub, 32)
    salt = _salt or os.urandom(16)
    cek = _hkdf(salt, ikm, b"Content-Encoding: aes128gcm\x00", 16)
    nonce = _hkdf(salt, ikm, b"Content-Encoding: nonce\x00", 12)
    ct = AESGCM(cek).encrypt(nonce, payload + b"\x02", None)
    return salt + struct.pack("!IB", 4096, len(as_pub)) + as_pub + ct


def _vapid(endpoint):
    u = urlparse(endpoint)
    head = b64u(json.dumps({"typ": "JWT", "alg": "ES256"}, separators=(",", ":")).encode())
    body = b64u(json.dumps({"aud": f"{u.scheme}://{u.netloc}", "exp": int(time.time()) + 12 * 3600,
                            "sub": CONTACT}, separators=(",", ":")).encode())
    r, s = decode_dss_signature(_PRIV.sign(f"{head}.{body}".encode(), ec.ECDSA(hashes.SHA256())))
    sig = b64u(r.to_bytes(32, "big") + s.to_bytes(32, "big"))
    return f"vapid t={head}.{body}.{sig}, k={PUBLIC_KEY}"


def send_one(sub, data):
    import httpx
    body = encrypt(sub, json.dumps(data).encode())
    r = httpx.post(sub["endpoint"], content=body, timeout=20, headers={
        "Authorization": _vapid(sub["endpoint"]), "Content-Encoding": "aes128gcm",
        "Content-Type": "application/octet-stream", "TTL": "86400", "Urgency": "high"})
    return r.status_code


def send_to(cid, user, data):
    """Push `data` ({title, body, tag, url}) to every browser of this visitor, in the background."""
    if not ENABLED:
        return

    def run():
        try:
            for ep, sub in _targets(cid, user):
                try:
                    STATS["sent"] += 1
                    code = send_one(json.loads(sub), data)
                    STATS["ok" if code < 300 else "failed"] += 1
                    STATS["last"] = f"{urlparse(ep).netloc} → HTTP {code}"
                    if code in (404, 410):                  # browser unsubscribed -> forget it
                        with _db() as con:
                            con.execute("DELETE FROM push_subs WHERE endpoint=?", (ep,))
                    elif code >= 300:
                        print(f"[push] {urlparse(ep).netloc}: HTTP {code}", flush=True)
                except Exception as err:
                    STATS["failed"] += 1
                    STATS["last"] = f"send failed: {err}"
                    print(f"[push] send failed: {err}", flush=True)
        except Exception as err:
            print(f"[push] {err}", flush=True)

    threading.Thread(target=run, daemon=True).start()


SW_JS = r"""// Service worker for "✅ Your result is ready" notifications (v2).
// Receives server Web Push messages (works while the tab is in the background, the phone is locked, or the
// customer is in another app) and focuses / opens the site when the notification is tapped.
self.addEventListener('install', function () { self.skipWaiting(); });
self.addEventListener('activate', function (e) { e.waitUntil(self.clients.claim()); });

self.addEventListener('push', function (e) {
  var d = {};
  try { d = e.data ? e.data.json() : {}; } catch (err) { d = { body: e.data ? e.data.text() : '' }; }
  var title = d.title || '✅ Your result is ready';
  e.waitUntil(self.registration.showNotification(title, {
    body: d.body || 'Tap to open your result.',
    icon: d.icon, badge: d.icon, tag: d.tag || 'cs-ready', renotify: false,
    vibrate: [120, 60, 120], requireInteraction: false, data: { url: d.url || '/' }
  }));
});

self.addEventListener('notificationclick', function (e) {
  e.notification.close();
  var url = (e.notification.data || {}).url || '/';
  e.waitUntil(self.clients.matchAll({ type: 'window', includeUncontrolled: true }).then(function (list) {
    for (var i = 0; i < list.length; i++) {
      try { if (new URL(list[i].url).origin === self.location.origin && 'focus' in list[i]) return list[i].focus(); }
      catch (err) {}
    }
    if (self.clients.openWindow) return self.clients.openWindow(url);
  }));
});
"""


def web_dir(icon_src=None):
    """Folder served to browsers (service worker + icon), created at start-up so nothing has to be uploaded."""
    import shutil
    import tempfile
    from pathlib import Path
    d = Path(tempfile.gettempdir()) / "cs_web"
    d.mkdir(parents=True, exist_ok=True)
    (d / "cs-notify-sw.js").write_text(SW_JS, encoding="utf-8")
    here = Path(__file__).parent
    for src in (here / "static" / "cs-icon.png", Path(icon_src) if icon_src else None, here / "favicon.png"):
        if src and src.exists():
            shutil.copyfile(src, d / "cs-icon.png")
            break
    return d


STATS = {"subs_saved": 0, "sent": 0, "ok": 0, "failed": 0, "last": ""}


def status():
    if not ENABLED:
        return "Notifications from the server are OFF (the `cryptography` package is missing — add it to requirements.txt)."
    try:
        with _db() as con:
            n = con.execute("SELECT COUNT(*) FROM push_subs").fetchone()[0]
    except Exception:
        n = "?"
    s = STATS
    return (f"Server notifications are ON. Browsers subscribed: {n}. Since restart: {s['subs_saved']} subscriptions "
            f"saved, {s['sent']} sent, {s['ok']} delivered to the push service, {s['failed']} failed."
            + (f" Last: {s['last']}" if s["last"] else ""))


def cid_from(request):
    """Browser id from the `cs_cid` cookie set by the page."""
    try:
        cookie = (request.headers.get("cookie") if request is not None else "") or ""
    except Exception:
        return ""
    for part in cookie.split(";"):
        k, _, v = part.strip().partition("=")
        if k == "cs_cid":
            return v[:64]
    return ""