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

"""Push notifications about new mail: the API devices subscribe through, and the
job that decides what to tell them — delivered to a mocked push service and
decrypted the way the browser would."""

from __future__ import annotations

import datetime as dt
import json
import uuid
from collections.abc import AsyncIterator
from dataclasses import dataclass

import httpx
import pytest_asyncio
import respx
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ec
from httpx import ASGITransport
from sqlalchemy import delete, select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool
from test_webpush import browser_decrypt

from app.config import get_settings
from app.identity.passwords import hash_password
from app.identity.sessions import create_session, destroy_session
from app.main import create_app
from app.models import Account, AppSettings, Mailbox, Message, PushSubscription, User
from app.push import service
from app.push.webpush import b64u_encode

CSRF = "csrf"
ENDPOINT = "https://fcm.googleapis.com/fcm/send/test-device"


@dataclass
class Ctx:
    client: httpx.AsyncClient
    db: AsyncSession
    user_id: uuid.UUID
    token: str
    inbox: Mailbox
    archive: Mailbox
    ua: ec.EllipticCurvePrivateKey
    auth: bytes


@pytest_asyncio.fixture
async def ctx() -> AsyncIterator[Ctx]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    db = async_sessionmaker(engine, expire_on_commit=False)()
    # The server's key pair lives in the one app_settings row: put it back as it was.
    row = await db.get(AppSettings, 1)
    saved = (row.vapid_public, row.vapid_private_ct) if row else None
    user = User(handle=f"push{uuid.uuid4().hex[:8]}@test.local", password_hash=hash_password("pw"))
    db.add(user)
    await db.flush()
    account = Account(
        user_id=user.id,
        email="me@example.com",
        imap_config={},
        smtp_config={},
        secret_ct=b"x",
        discovery_src="manual",
        state="live",
    )
    db.add(account)
    await db.flush()
    inbox = Mailbox(account_id=account.id, user_id=user.id, path="INBOX", special_use="inbox")
    archive = Mailbox(account_id=account.id, user_id=user.id, path="Archive", special_use="archive")
    db.add_all([inbox, archive])
    await db.commit()
    token = await create_session(user.id, "pytest")
    client = httpx.AsyncClient(
        transport=ASGITransport(app=create_app()),
        base_url="http://testserver",
        cookies={"session": token, "csrf": CSRF},
    )
    try:
        yield Ctx(
            client,
            db,
            user.id,
            token,
            inbox,
            archive,
            ec.generate_private_key(ec.SECP256R1()),
            b"0123456789abcdef",
        )
    finally:
        await client.aclose()
        await db.execute(delete(User).where(User.id == user.id))
        if saved is not None:
            await db.execute(
                update(AppSettings)
                .where(AppSettings.id == 1)
                .values(vapid_public=saved[0], vapid_private_ct=saved[1])
            )
        await db.commit()
        await db.close()
        await engine.dispose()


def _keys(c: Ctx) -> dict[str, str]:
    public = c.ua.public_key().public_bytes(
        serialization.Encoding.X962, serialization.PublicFormat.UncompressedPoint
    )
    return {"p256dh": b64u_encode(public), "auth": b64u_encode(c.auth)}


async def _subscribe(c: Ctx, endpoint: str = ENDPOINT) -> httpx.Response:
    return await c.client.post(
        "/api/push/subscriptions",
        json={"endpoint": endpoint, "keys": _keys(c)},
        headers={"X-CSRF": CSRF},
    )


_uid = iter(range(1, 10_000))


async def _mail(
    c: Ctx,
    *,
    subject: str = "Hello",
    sender: str = "alice@example.org",
    name: str = "Alice",
    box: Mailbox | None = None,
    flags: list[str] | None = None,
    age: dt.timedelta = dt.timedelta(minutes=1),
) -> uuid.UUID:
    m = Message(
        user_id=c.user_id,
        account_id=c.inbox.account_id,
        mailbox_id=(box or c.inbox).id,
        uid=next(_uid),
        thread_id=None,
        subject=subject,
        from_addr=[{"name": name, "email": sender}],
        internal_date=dt.datetime.now(dt.UTC) - age,
        flags=flags or [],
        snippet="  First   line of the message ",
        blob_sha256="0" * 64,
    )
    c.db.add(m)
    await c.db.commit()
    return m.id


def _decrypted(c: Ctx, route: respx.Route) -> list[dict[str, object]]:
    return [json.loads(browser_decrypt(call.request.content, c.ua, c.auth)) for call in route.calls]


async def test_subscribe_validates_and_hands_out_the_key(ctx: Ctx) -> None:
    key = await ctx.client.get("/api/push/key")
    assert key.status_code == 200 and len(key.json()["public_key"]) == 87
    assert (await ctx.client.get("/api/push/key")).json() == key.json()  # made once

    assert (await _subscribe(ctx, "https://169.254.169.254/latest")).status_code == 400
    bad = await ctx.client.post(
        "/api/push/subscriptions",
        json={"endpoint": ENDPOINT, "keys": {"p256dh": "AAAA", "auth": "AAAA"}},
        headers={"X-CSRF": CSRF},
    )
    assert bad.status_code == 400
    assert (await _subscribe(ctx)).status_code == 200
    assert (await _subscribe(ctx)).status_code == 200  # again: refreshed, not doubled
    rows = (
        await ctx.db.scalars(
            select(PushSubscription).where(PushSubscription.user_id == ctx.user_id)
        )
    ).all()
    assert len(rows) == 1

    gone = await ctx.client.post(
        "/api/push/subscriptions/remove", json={"endpoint": ENDPOINT}, headers={"X-CSRF": CSRF}
    )
    assert gone.status_code == 200
    assert (
        await ctx.db.scalar(select(PushSubscription).where(PushSubscription.endpoint == ENDPOINT))
    ) is None


@respx.mock
async def test_new_inbox_mail_is_pushed_encrypted_and_signed(ctx: Ctx) -> None:
    await _subscribe(ctx)
    route = respx.post(ENDPOINT).mock(return_value=httpx.Response(201))
    news = await _mail(ctx, subject="Lunch?")
    skipped = [
        await _mail(ctx, flags=["\\Seen"]),  # read on another device already
        await _mail(ctx, box=ctx.archive),  # filed away by a filter
        await _mail(ctx, age=dt.timedelta(days=1)),  # old mail moved back to the Inbox
        await _mail(ctx, sender="ME@example.com"),  # from yourself
    ]
    assert await service.notify(ctx.user_id, [news, *skipped]) == 1
    request = route.calls[0].request
    assert request.headers["Content-Encoding"] == "aes128gcm"
    assert request.headers["Authorization"].startswith("vapid t=")
    [payload] = _decrypted(ctx, route)
    assert payload["title"] == "Alice"
    assert payload["body"] == "Lunch?\nFirst line of the message"
    assert payload["url"] == "/mail" and payload["type"] == "mail"
    assert payload["unread"] == 3  # Inbox unread: news, old, from-me


@respx.mock
async def test_many_at_once_become_one_summary(ctx: Ctx) -> None:
    await _subscribe(ctx)
    route = respx.post(ENDPOINT).mock(return_value=httpx.Response(201))
    ids = [
        await _mail(ctx, name=n, sender=f"{n.lower()}@example.org")
        for n in ("Ann", "Bob", "Cy", "Dee", "Ann")
    ]
    assert await service.notify(ctx.user_id, ids) == 1
    [payload] = _decrypted(ctx, route)
    assert payload["title"] == "5 new messages"
    assert payload["body"] == "Ann, Bob, Cy…"
    assert payload["tag"] == "new-mail"


@respx.mock
async def test_signed_out_or_expired_devices_are_forgotten(ctx: Ctx) -> None:
    await _subscribe(ctx)
    other = "https://updates.push.services.mozilla.com/wpush/v2/other"
    await _subscribe(ctx, other)
    respx.post(ENDPOINT).mock(return_value=httpx.Response(410))  # browser unsubscribed
    firefox = respx.post(other).mock(return_value=httpx.Response(201))
    assert await service.notify(ctx.user_id, [await _mail(ctx)]) == 1
    left = (
        await ctx.db.scalars(
            select(PushSubscription.endpoint).where(PushSubscription.user_id == ctx.user_id)
        )
    ).all()
    assert left == [other]

    await destroy_session(ctx.token)  # signed out: no subject lines on that device any more
    assert await service.notify(ctx.user_id, [await _mail(ctx)]) == 0
    assert firefox.call_count == 1
    assert (
        await ctx.db.scalar(select(PushSubscription).where(PushSubscription.user_id == ctx.user_id))
    ) is None


async def test_no_devices_no_work(ctx: Ctx) -> None:
    assert await service.notify(ctx.user_id, [await _mail(ctx)]) == 0