webmail-client/backend/tests/push/test_webpush.py

"""Web Push protocol: RFC 8291 encryption (against the RFC's worked example and
by decrypting as a browser would), VAPID signatures, the push-service allow-list."""

from __future__ import annotations

import json

from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import ec
from cryptography.hazmat.primitives.asymmetric.utils import encode_dss_signature
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

from app.push import webpush
from app.push.webpush import b64u_decode, b64u_encode

# RFC 8291 §5 / Appendix A.
RFC_AUTH = "BTBZMqHH6r4Tts7J_aSIgg"
RFC_UA_PUBLIC = (
    "BCVxsr7N_eNgVRqvHtD0zTZsEc6-VV-JvLexhqUzORcxaOzi6-AYWXvTBHm4bjyPjs7Vd8pZGH6SRpkNtoIAiw4"
)
RFC_AS_PRIVATE = "yfWPiYE-n46HLnH0KqZOF1fJJU3MYrct3AELtAQ-oRw"
RFC_SALT = "DGv6ra1nlYgDCS1FRnbzlw"
RFC_BODY = (
    "DGv6ra1nlYgDCS1FRnbzlwAAEABBBP4z9KsN6nGRTbVYI_c7VJSPQTBtkgcy27ml"
    "mlMoZIIgDll6e3vCYLocInmYWAmS6TlzAC8wEqKK6PBru3jl7A_yl95bQpu6cVPT"
    "pK4Mqgkf1CXztLVBSt2Ks3oZwbuwXPXLWyouBWLVWGNWQexSgSxsj_Qulcy4a-fN"
)


def _hkdf(salt: bytes, ikm: bytes, info: bytes, n: int) -> bytes:
    return HKDF(algorithm=hashes.SHA256(), length=n, salt=salt, info=info).derive(ikm)


def browser_decrypt(body: bytes, ua: ec.EllipticCurvePrivateKey, auth: bytes) -> bytes:
    """What the browser does with a push body (RFC 8291 receiver side)."""
    salt, idlen = body[:16], body[20]
    as_public = body[21 : 21 + idlen]
    ua_public = ua.public_key().public_bytes(
        serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
    )
    as_key = ec.EllipticCurvePublicKey.from_encoded_point(ec.SECP256R1(), as_public)
    ikm = _hkdf(
        auth, ua.exchange(ec.ECDH(), as_key), b"WebPush: info\x00" + ua_public + as_public, 32
    )
    cek = _hkdf(salt, ikm, b"Content-Encoding: aes128gcm\x00", 16)
    nonce = _hkdf(salt, ikm, b"Content-Encoding: nonce\x00", 12)
    plain = AESGCM(cek).decrypt(nonce, body[21 + idlen :], None)
    assert plain.endswith(b"\x02")
    return plain[:-1]


def test_rfc8291_worked_example() -> None:
    sender = ec.derive_private_key(
        int.from_bytes(b64u_decode(RFC_AS_PRIVATE), "big"), ec.SECP256R1()
    )
    body = webpush.encrypt(
        b"When I grow up, I want to be a watermelon",
        RFC_UA_PUBLIC,
        RFC_AUTH,
        salt=b64u_decode(RFC_SALT),
        sender=sender,
    )
    assert b64u_encode(body) == RFC_BODY


def test_browser_can_decrypt() -> None:
    ua = ec.generate_private_key(ec.SECP256R1())
    auth = b"0123456789abcdef"
    p256dh = b64u_encode(
        ua.public_key().public_bytes(
            serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
        )
    )
    payload = json.dumps({"title": "Zażółć gęślą jaźń"}).encode()
    body = webpush.encrypt(payload, p256dh, b64u_encode(auth))
    assert browser_decrypt(body, ua, auth) == payload


def test_vapid_authorization_is_a_valid_es256_jwt() -> None:
    key = webpush.Vapid.generate()
    again = webpush.Vapid.from_pem(key.pem())
    assert again.public_b64 == key.public_b64
    header = key.authorization("https://fcm.googleapis.com/fcm/send/abc", "https://ultimo.email")
    token = header.split("t=")[1].split(",")[0]
    assert header.endswith(f"k={key.public_b64}")
    head, claims, sig = token.split(".")
    assert json.loads(b64u_decode(claims)) | {"exp": 0} == {
        "aud": "https://fcm.googleapis.com",
        "sub": "https://ultimo.email",
        "exp": 0,
    }
    raw = b64u_decode(sig)
    der = encode_dss_signature(int.from_bytes(raw[:32], "big"), int.from_bytes(raw[32:], "big"))
    public = ec.EllipticCurvePublicKey.from_encoded_point(
        ec.SECP256R1(), b64u_decode(key.public_b64)
    )
    public.verify(der, f"{head}.{claims}".encode(), ec.ECDSA(hashes.SHA256()))  # raises if bad


def test_only_push_services_are_allowed() -> None:
    ok = [
        "https://fcm.googleapis.com/fcm/send/x",
        "https://jmt17.google.com/fcm/send/x",
        "https://updates.push.services.mozilla.com/wpush/v2/x",
        "https://web.push.apple.com/x",
        "https://wns2-db5p.notify.windows.com/w/?token=x",
    ]
    bad = [
        "http://fcm.googleapis.com/fcm/send/x",
        "https://fcm.googleapis.com:8443/x",
        "https://evil.com/fcm.googleapis.com",
        "https://fcm.googleapis.com.evil.com/x",
        "https://google.com.evil.com/x",
        "https://notapple-push.apple.com.evil.com/x",
        "https://127.0.0.1/x",
        "https://postgres/x",
    ]
    assert all(webpush.allowed_endpoint(u) for u in ok)
    assert not any(webpush.allowed_endpoint(u) for u in bad)