"""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"