webmail-client/backend/tests/identity/test_trusted_devices.py

"""Trusted devices: "don't ask for the code on this device" at a 2FA sign-in."""

from __future__ import annotations

import uuid
from collections.abc import AsyncIterator
from dataclasses import dataclass, field

import httpx
import pyotp
import pytest_asyncio
from httpx import ASGITransport
from sqlalchemy import delete
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
from sqlalchemy.pool import NullPool

from app.config import get_settings
from app.identity import trusted_devices
from app.identity.passwords import hash_password
from app.identity.totp import encrypt_totp_secret
from app.main import create_app
from app.models import User

PASSWORD = "hunter2pass"


@dataclass
class Ctx:
    session: AsyncSession
    created: list[uuid.UUID] = field(default_factory=list)

    async def user(self) -> tuple[User, str]:
        secret = pyotp.random_base32()
        user = User(
            handle=f"td{uuid.uuid4().hex[:8]}@test.local",
            password_hash=hash_password(PASSWORD),
            totp_secret=encrypt_totp_secret(secret),
        )
        self.session.add(user)
        await self.session.commit()
        self.created.append(user.id)
        return user, secret


def _browser() -> httpx.AsyncClient:
    """One browser: keeps its cookies between requests."""
    return httpx.AsyncClient(transport=ASGITransport(app=create_app()), base_url="http://t")


@pytest_asyncio.fixture
async def ctx() -> AsyncIterator[Ctx]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    session = async_sessionmaker(engine, expire_on_commit=False)()
    c = Ctx(session=session)
    try:
        yield c
    finally:
        for uid in c.created:
            await trusted_devices.revoke_all(uid)
        await session.execute(delete(User).where(User.id.in_(c.created)))
        await session.commit()
        await session.close()
        await engine.dispose()


async def _login(
    browser: httpx.AsyncClient, user: User, totp: str | None = None, remember: bool = False
) -> httpx.Response:
    body: dict[str, object] = {"handle": user.handle, "password": PASSWORD}
    if totp:
        body["totp"] = totp
    if remember:
        body["remember_device"] = True
    return await browser.post("/api/auth/login", json=body)


async def test_a_remembered_browser_signs_in_without_the_code(ctx: Ctx) -> None:
    user, secret = await ctx.user()
    async with _browser() as browser:
        first = await _login(browser, user, pyotp.TOTP(secret).now(), remember=True)
        assert first.status_code == 200
        assert "trusted_device" in browser.cookies
        again = await _login(browser, user)
        assert again.status_code == 200
        # The password is still checked.
        wrong = await browser.post(
            "/api/auth/login", json={"handle": user.handle, "password": "not-it-at-all"}
        )
        assert wrong.json()["error"]["code"] == "invalid_credentials"


async def test_without_the_box_the_code_is_asked_every_time(ctx: Ctx) -> None:
    user, secret = await ctx.user()
    async with _browser() as browser:
        assert (await _login(browser, user, pyotp.TOTP(secret).now())).status_code == 200
        assert "trusted_device" not in browser.cookies
        assert (await _login(browser, user)).json()["error"]["code"] == "totp_required"


async def test_trust_is_per_account_and_per_browser(ctx: Ctx) -> None:
    alice, alice_secret = await ctx.user()
    bob, _ = await ctx.user()
    async with _browser() as browser:
        await _login(browser, alice, pyotp.TOTP(alice_secret).now(), remember=True)
        # Another account on the same browser still needs its own code.
        assert (await _login(browser, bob)).json()["error"]["code"] == "totp_required"
    async with _browser() as other_browser:
        assert (await _login(other_browser, alice)).json()["error"]["code"] == "totp_required"


async def test_turning_2fa_off_revokes_the_trust(ctx: Ctx) -> None:
    user, secret = await ctx.user()
    async with _browser() as browser:
        await _login(browser, user, pyotp.TOTP(secret).now(), remember=True)
        token = browser.cookies["trusted_device"]
        await trusted_devices.revoke_all(user.id)  # what totp/disable, reset and block call
        assert (await _login(browser, user)).json()["error"]["code"] == "totp_required"
        assert browser.cookies.get("trusted_device") == token  # stale cookie, no longer honoured


async def test_totp_disable_endpoint_revokes(ctx: Ctx) -> None:
    user, secret = await ctx.user()
    async with _browser() as browser:
        await _login(browser, user, pyotp.TOTP(secret).now(), remember=True)
        old_token = browser.cookies["trusted_device"]
        csrf = browser.cookies["csrf"]
        resp = await browser.post(
            "/api/auth/totp/disable", json={"password": PASSWORD}, headers={"X-CSRF": csrf}
        )
        assert resp.status_code == 200
    # Turn 2FA back on: the old trust must not come back with it.
    user.totp_secret = encrypt_totp_secret(pyotp.random_base32())
    ctx.session.add(user)
    await ctx.session.commit()
    async with _browser() as browser:
        browser.cookies.set("trusted_device", old_token)
        assert (await _login(browser, user)).json()["error"]["code"] == "totp_required"