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

"""M1 identity: 2FA (TOTP), session list/revoke, password reset (SPEC §6.1/§9.1)."""

from __future__ import annotations

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

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.passwords import hash_password
from app.identity.sessions import create_session
from app.main import create_app
from app.models import User
from app.redis_client import get_redis

CSRF = "csrf-token"
PASSWORD = "hunter2pass"


@dataclass
class Ctx:
    app: object
    client: httpx.AsyncClient
    user_id: uuid.UUID
    handle: str
    session: AsyncSession


def _hdr() -> dict[str, str]:
    return {"X-CSRF": CSRF}


@pytest_asyncio.fixture
async def ctx() -> AsyncIterator[Ctx]:
    engine = create_async_engine(get_settings().database_url, poolclass=NullPool)
    sm = async_sessionmaker(engine, expire_on_commit=False)
    session = sm()
    handle = f"id{uuid.uuid4().hex[:8]}@test.local"
    user = User(handle=handle, password_hash=hash_password(PASSWORD), recovery_email="r@x.com")
    session.add(user)
    await session.commit()
    token = await create_session(user.id, "pytest")
    app = create_app()
    client = httpx.AsyncClient(
        transport=ASGITransport(app=app),
        base_url="http://testserver",
        cookies={"session": token, "csrf": CSRF},
    )
    try:
        yield Ctx(app=app, client=client, user_id=user.id, handle=handle, session=session)
    finally:
        await client.aclose()
        await session.execute(delete(User).where(User.id == user.id))
        await session.commit()
        await session.close()
        await engine.dispose()


def _fresh_client(app: object) -> httpx.AsyncClient:
    return httpx.AsyncClient(transport=ASGITransport(app=app), base_url="http://testserver")


async def test_totp_setup_enable_and_login(ctx: Ctx) -> None:
    setup = await ctx.client.post("/api/auth/totp/setup", headers=_hdr())
    assert setup.status_code == 200
    secret = setup.json()["secret"]
    assert setup.json()["provisioning_uri"].startswith("otpauth://")

    # wrong code rejected
    bad = await ctx.client.post("/api/auth/totp/enable", json={"code": "000000"}, headers=_hdr())
    assert bad.status_code == 400

    code = pyotp.TOTP(secret).now()
    ok = await ctx.client.post("/api/auth/totp/enable", json={"code": code}, headers=_hdr())
    assert ok.status_code == 200

    # login now requires the 2FA code
    async with _fresh_client(ctx.app) as anon:
        no_totp = await anon.post(
            "/api/auth/login", json={"handle": ctx.handle, "password": PASSWORD}
        )
        assert no_totp.status_code == 401
        assert no_totp.json()["error"]["code"] == "totp_required"

        with_totp = await anon.post(
            "/api/auth/login",
            json={"handle": ctx.handle, "password": PASSWORD, "totp": pyotp.TOTP(secret).now()},
        )
        assert with_totp.status_code == 200


async def test_totp_disable(ctx: Ctx) -> None:
    # enable first
    secret = (await ctx.client.post("/api/auth/totp/setup", headers=_hdr())).json()["secret"]
    await ctx.client.post(
        "/api/auth/totp/enable", json={"code": pyotp.TOTP(secret).now()}, headers=_hdr()
    )
    assert (await ctx.client.get("/api/auth/me")).json()["totp_enabled"] is True

    # wrong password is rejected
    bad = await ctx.client.post("/api/auth/totp/disable", json={"password": "nope"}, headers=_hdr())
    assert bad.status_code == 401
    assert (await ctx.client.get("/api/auth/me")).json()["totp_enabled"] is True

    # correct password disables it
    ok = await ctx.client.post(
        "/api/auth/totp/disable", json={"password": PASSWORD}, headers=_hdr()
    )
    assert ok.status_code == 200
    assert (await ctx.client.get("/api/auth/me")).json()["totp_enabled"] is False

    # login now works without a code
    async with _fresh_client(ctx.app) as anon:
        resp = await anon.post("/api/auth/login", json={"handle": ctx.handle, "password": PASSWORD})
        assert resp.status_code == 200


async def test_sessions_list_and_revoke(ctx: Ctx) -> None:
    # create a second session for the same user (a different device)
    other_token = await create_session(ctx.user_id, "other-device")

    listing = await ctx.client.get("/api/auth/sessions")
    assert listing.status_code == 200
    rows = listing.json()
    assert len(rows) == 2
    current = [r for r in rows if r["current"]]
    assert len(current) == 1

    other_sid = next(r["sid"] for r in rows if not r["current"])
    revoked = await ctx.client.request("DELETE", f"/api/auth/sessions/{other_sid}", headers=_hdr())
    assert revoked.status_code == 200

    # the revoked token no longer resolves
    async with _fresh_client(ctx.app) as anon:
        anon.cookies.set("session", other_token)
        me = await anon.get("/api/auth/me")
        assert me.status_code == 401


async def test_password_reset_flow(ctx: Ctx) -> None:
    # request a reset; token is stored in redis (not returned)
    req = await ctx.client.post("/api/auth/password/reset-request", json={"handle": ctx.handle})
    assert req.status_code == 200

    redis = get_redis()
    keys = await redis.keys("pwreset:*")
    # Pick the token bound to THIS user (redis DB may hold others from prior runs).
    token = None
    for key in keys:
        if await redis.get(key) == str(ctx.user_id):
            token = key.split(":", 1)[1]
            break
    assert token is not None, "a reset token should have been issued for this user"

    new_password = "brandNewPass9"
    reset = await ctx.client.post(
        "/api/auth/password/reset", json={"token": token, "password": new_password}
    )
    assert reset.status_code == 200

    # token is single-use
    again = await ctx.client.post(
        "/api/auth/password/reset", json={"token": token, "password": new_password}
    )
    assert again.status_code == 400

    # new password works, old one does not
    async with _fresh_client(ctx.app) as anon:
        good = await anon.post(
            "/api/auth/login", json={"handle": ctx.handle, "password": new_password}
        )
        assert good.status_code == 200
        bad = await anon.post("/api/auth/login", json={"handle": ctx.handle, "password": PASSWORD})
        assert bad.status_code == 401


async def test_reset_request_unknown_handle_still_ok(ctx: Ctx) -> None:
    # must not leak whether a handle exists
    resp = await ctx.client.post("/api/auth/password/reset-request", json={"handle": "nobody-here"})
    assert resp.status_code == 200