wiregui/tests/test_integration_mfa.py

"""Integration tests for MFA — full registration and authentication flows through the database."""

import pyotp
from sqlmodel import func, select

from wiregui.auth.mfa import generate_totp_secret, verify_totp_code
from wiregui.auth.passwords import hash_password, verify_password
from wiregui.auth.session import authenticate_user
from wiregui.models.mfa_method import MFAMethod
from wiregui.models.user import User
from wiregui.utils.time import utcnow


async def test_full_totp_registration_flow(session, monkeypatch):
    """End-to-end: create user → generate secret → verify code → store method → re-verify from DB."""
    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def mock_session():
        yield session

    # Create user with password
    user = User(email="mfa-flow@example.com", password_hash=hash_password("secure123"))
    session.add(user)
    await session.flush()

    # Step 1: Generate TOTP secret (happens in account page)
    secret = generate_totp_secret()

    # Step 2: User scans QR, enters code from their authenticator
    totp = pyotp.TOTP(secret)
    code = totp.now()

    # Step 3: Verify the code is correct before saving
    assert verify_totp_code(secret, code) is True

    # Step 4: Save the MFA method to DB
    method = MFAMethod(
        name="My Authenticator",
        type="totp",
        payload={"secret": secret},
        user_id=user.id,
    )
    session.add(method)
    await session.flush()

    # Step 5: Simulate future login — load method from DB and verify a fresh code
    fetched_methods = (await session.execute(
        select(MFAMethod).where(MFAMethod.user_id == user.id)
    )).scalars().all()

    assert len(fetched_methods) == 1
    stored_secret = fetched_methods[0].payload["secret"]
    fresh_code = pyotp.TOTP(stored_secret).now()
    assert verify_totp_code(stored_secret, fresh_code) is True


async def test_mfa_blocks_login_without_code(session, monkeypatch):
    """User with MFA should not be fully authenticated without completing MFA challenge."""
    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def mock_session():
        yield session

    monkeypatch.setattr("wiregui.auth.session.async_session", mock_session)

    # Create user with MFA
    user = User(email="mfa-block@example.com", password_hash=hash_password("password1"))
    session.add(user)
    await session.flush()

    secret = generate_totp_secret()
    method = MFAMethod(name="Phone", type="totp", payload={"secret": secret}, user_id=user.id)
    session.add(method)
    await session.flush()

    # Password auth succeeds
    authed_user = await authenticate_user("mfa-block@example.com", "password1")
    assert authed_user is not None

    # But MFA methods exist — login page would redirect to /mfa instead of completing login
    mfa_methods = (await session.execute(
        select(MFAMethod).where(MFAMethod.user_id == authed_user.id)
    )).scalars().all()
    assert len(mfa_methods) > 0  # Login flow would check this and redirect to /mfa


async def test_mfa_wrong_code_rejected(session):
    """Wrong TOTP code should be rejected even if method is valid."""
    user = User(email="mfa-wrong@example.com", password_hash=hash_password("pw"))
    session.add(user)
    await session.flush()

    secret = generate_totp_secret()
    method = MFAMethod(name="Auth", type="totp", payload={"secret": secret}, user_id=user.id)
    session.add(method)
    await session.flush()

    # Load from DB and try wrong code
    fetched = (await session.execute(
        select(MFAMethod).where(MFAMethod.user_id == user.id)
    )).scalar_one()

    assert verify_totp_code(fetched.payload["secret"], "000000") is False
    assert verify_totp_code(fetched.payload["secret"], "123456") is False


async def test_mfa_multiple_methods_any_valid_code_works(session):
    """If user has multiple TOTP methods, a valid code from any should work."""
    user = User(email="mfa-multi@example.com")
    session.add(user)
    await session.flush()

    secret1 = generate_totp_secret()
    secret2 = generate_totp_secret()

    session.add(MFAMethod(name="Phone", type="totp", payload={"secret": secret1}, user_id=user.id))
    session.add(MFAMethod(name="Backup", type="totp", payload={"secret": secret2}, user_id=user.id))
    await session.flush()

    methods = (await session.execute(
        select(MFAMethod).where(MFAMethod.user_id == user.id)
    )).scalars().all()

    # Code from method 1 should verify against method 1's secret
    code1 = pyotp.TOTP(secret1).now()
    verified = False
    for m in methods:
        if verify_totp_code(m.payload["secret"], code1):
            verified = True
            break
    assert verified is True

    # Code from method 2 should also work
    code2 = pyotp.TOTP(secret2).now()
    verified2 = False
    for m in methods:
        if verify_totp_code(m.payload["secret"], code2):
            verified2 = True
            break
    assert verified2 is True


async def test_mfa_method_last_used_tracking(session):
    """Verifying MFA should update last_used_at timestamp."""
    user = User(email="mfa-tracking@example.com")
    session.add(user)
    await session.flush()

    secret = generate_totp_secret()
    method = MFAMethod(name="Auth", type="totp", payload={"secret": secret}, user_id=user.id)
    session.add(method)
    await session.flush()

    assert method.last_used_at is None

    # Simulate successful verification and update
    code = pyotp.TOTP(secret).now()
    assert verify_totp_code(secret, code) is True

    method.last_used_at = utcnow()
    session.add(method)
    await session.flush()

    fetched = await session.get(MFAMethod, method.id)
    assert fetched.last_used_at is not None


async def test_mfa_delete_method_allows_login_without_mfa(session, monkeypatch):
    """After removing all MFA methods, user should not be redirected to MFA challenge."""
    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def mock_session():
        yield session

    monkeypatch.setattr("wiregui.auth.session.async_session", mock_session)

    user = User(email="mfa-remove@example.com", password_hash=hash_password("pw"))
    session.add(user)
    await session.flush()

    secret = generate_totp_secret()
    method = MFAMethod(name="Temp", type="totp", payload={"secret": secret}, user_id=user.id)
    session.add(method)
    await session.flush()

    # MFA exists
    count = (await session.execute(
        select(func.count()).select_from(MFAMethod).where(MFAMethod.user_id == user.id)
    )).scalar()
    assert count == 1

    # Delete it
    await session.delete(method)
    await session.flush()

    count = (await session.execute(
        select(func.count()).select_from(MFAMethod).where(MFAMethod.user_id == user.id)
    )).scalar()
    assert count == 0

    # Password auth still works
    authed = await authenticate_user("mfa-remove@example.com", "pw")
    assert authed is not None

    # No MFA methods — login flow would skip MFA challenge
    mfa_check = (await session.execute(
        select(MFAMethod).where(MFAMethod.user_id == authed.id)
    )).scalars().all()
    assert len(mfa_check) == 0


async def test_disabled_user_with_mfa_cannot_login(session, monkeypatch):
    """Disabled user should be rejected at password stage, never reaching MFA."""
    from contextlib import asynccontextmanager

    @asynccontextmanager
    async def mock_session():
        yield session

    monkeypatch.setattr("wiregui.auth.session.async_session", mock_session)

    user = User(
        email="mfa-disabled@example.com",
        password_hash=hash_password("pw"),
        disabled_at=utcnow(),
    )
    session.add(user)
    await session.flush()

    secret = generate_totp_secret()
    session.add(MFAMethod(name="Auth", type="totp", payload={"secret": secret}, user_id=user.id))
    await session.flush()

    # Password auth rejects disabled user before MFA is ever checked
    result = await authenticate_user("mfa-disabled@example.com", "pw")
    assert result is None