kite/backend/tests/test_auth.py

"""
Wave 0 xfail stubs for CR-01, CR-02, CR-03 — session revocation on privilege change.

These tests lock the behavioral contract expected by the session-revocation
implementation already present in backend/api/auth.py. They are marked xfail
with strict=False so that XPASS (already passing) or XFAIL both succeed during
Wave 0. Plan 08-03 removes the @pytest.mark.xfail decorator and promotes them to
strict passing tests.

CR-01: change_password revokes all other refresh tokens; current session preserved
CR-02: enable_totp revokes all other refresh tokens; current session preserved
CR-03: disable_totp revokes all other refresh tokens; current session preserved
"""
from __future__ import annotations

from datetime import datetime, timezone
from unittest.mock import patch

import pyotp
import pytest
import pytest_asyncio
from httpx import ASGITransport, AsyncClient
from sqlalchemy.ext.asyncio import AsyncSession


# ── Helpers ───────────────────────────────────────────────────────────────────

VALID_PASSWORD = "StrongPass12!"


class FakeRedis:
    """In-memory fake Redis for testing. Mirrors test_auth_api.py's FakeRedis."""

    def __init__(self):
        self._store: dict = {}

    async def get(self, key):
        entry = self._store.get(key)
        if entry is None:
            return None
        val, exp = entry
        if exp is not None and datetime.now(timezone.utc).timestamp() > exp:
            del self._store[key]
            return None
        return val

    async def incr(self, key):
        entry = self._store.get(key)
        if entry is None:
            self._store[key] = (1, None)
            return 1
        val, exp = entry
        new_val = val + 1
        self._store[key] = (new_val, exp)
        return new_val

    async def expire(self, key, seconds):
        if key in self._store:
            val, _ = self._store[key]
            deadline = datetime.now(timezone.utc).timestamp() + seconds
            self._store[key] = (val, deadline)

    async def set(self, key, value, ex=None):
        deadline = None
        if ex is not None:
            deadline = datetime.now(timezone.utc).timestamp() + ex
        self._store[key] = (value, deadline)

    async def close(self):
        pass


@pytest_asyncio.fixture
async def revoke_client(db_session: AsyncSession):
    """Async HTTP test client with DB override and fresh FakeRedis.

    All session-revocation tests use this fixture so they share a clean
    in-memory state for both the DB and Redis (per-account rate limiter).

    The client sends a fixed User-Agent ("docuvault-test/1.0") matching
    the conftest._TEST_USER_AGENT constant. This is required because the
    access token's `fgp` claim is bound to the User-Agent at login time —
    subsequent requests using that access token must send the same User-Agent
    or `get_current_user` returns 401 "Token fingerprint mismatch".
    """
    from deps.db import get_db
    from main import app
    from api.auth import limiter as auth_limiter

    app.dependency_overrides[get_db] = lambda: db_session
    fake_redis = FakeRedis()
    app.state.redis = fake_redis

    try:
        auth_limiter._storage.reset()
    except Exception:
        pass

    async with AsyncClient(
        transport=ASGITransport(app=app),
        base_url="http://test",
        headers={"User-Agent": "docuvault-test/1.0"},
    ) as c:
        yield c

    app.dependency_overrides.clear()
    app.state.redis = None


async def _register_user(client: AsyncClient, handle: str, email: str) -> None:
    """Register a user account; assert 201."""
    resp = await client.post(
        "/api/auth/register",
        json={"handle": handle, "email": email, "password": VALID_PASSWORD},
    )
    assert resp.status_code == 201, f"Register failed: {resp.text}"


async def _login_session(
    client: AsyncClient, email: str, user_agent: str | None = None
) -> tuple[str, str]:
    """Login and return (access_token, refresh_token_cookie_value).

    Each call to login creates a new RefreshToken row in the DB, so two
    sequential calls yield two separate sessions even with the same User-Agent.
    The user_agent parameter controls the fgp claim embedded in the access token —
    if supplied, the same User-Agent must be used for all subsequent requests
    that present that access token.
    """
    headers: dict = {}
    if user_agent:
        headers["User-Agent"] = user_agent
    resp = await client.post(
        "/api/auth/login",
        json={"email": email, "password": VALID_PASSWORD},
        headers=headers,
    )
    assert resp.status_code == 200, f"Login failed: {resp.text}"
    data = resp.json()
    access_token = data["access_token"]
    # The refresh cookie is set with path=/api/auth/refresh — extract the raw value
    # so we can inject it explicitly when needed.
    refresh_cookie = resp.cookies.get("refresh_token") or ""
    return access_token, refresh_cookie


async def _try_refresh(
    client: AsyncClient, refresh_token: str
) -> int:
    """POST /api/auth/refresh with the given refresh_token cookie; return status code.

    Patches the security alert email task so presenting a revoked token does not
    attempt a real Celery/Redis broker connection (which is unavailable in unit tests).
    """
    from unittest.mock import MagicMock
    mock_task = MagicMock()
    mock_task.delay = MagicMock()
    with patch.dict(
        "sys.modules",
        {"tasks.email_tasks": MagicMock(send_security_alert_email=mock_task)},
    ):
        resp = await client.post(
            "/api/auth/refresh",
            cookies={"refresh_token": refresh_token},
        )
    return resp.status_code


# ── Tests ─────────────────────────────────────────────────────────────────────


@pytest.mark.asyncio
async def test_change_password_revokes_other_sessions(revoke_client, db_session):
    """CR-01: change_password revokes other sessions; current session preserved.

    Setup: register user; login twice (sessions A and B).
    Action: using session A's access token, POST /api/auth/change-password with
            session A's refresh_token cookie so skip_token_hash preserves session A.
    Assert:
      - response 200 with sessions_revoked == 1
      - session A refresh token still valid (200 on /api/auth/refresh)
      - session B refresh token revoked (401 on /api/auth/refresh)
    """
    email = "cr01_test@example.com"
    await _register_user(revoke_client, handle="cr01user", email=email)

    # Two sequential logins create two separate RefreshToken rows (sessions A and B).
    # The client uses a fixed User-Agent so the access token fgp claim matches
    # subsequent requests that present token_a.
    token_a, refresh_a = await _login_session(revoke_client, email)
    _token_b, refresh_b = await _login_session(revoke_client, email)

    with patch("services.auth.check_hibp", return_value=False):
        resp = await revoke_client.post(
            "/api/auth/change-password",
            json={"current_password": VALID_PASSWORD, "new_password": "NewValidPass99!"},
            headers={"Authorization": f"Bearer {token_a}"},
            cookies={"refresh_token": refresh_a},
        )

    assert resp.status_code == 200, f"change_password failed: {resp.text}"
    data = resp.json()
    assert data["sessions_revoked"] == 1, (
        f"Expected exactly 1 other session revoked, got {data['sessions_revoked']}"
    )

    # Session A's refresh token must still be usable
    status_a = await _try_refresh(revoke_client, refresh_a)
    assert status_a == 200, f"Session A refresh should still be valid, got {status_a}"

    # Session B's refresh token must now be revoked
    status_b = await _try_refresh(revoke_client, refresh_b)
    assert status_b == 401, f"Session B refresh should be revoked (401), got {status_b}"


@pytest.mark.asyncio
async def test_enable_totp_revokes_other_sessions(revoke_client, db_session):
    """CR-02: enable_totp revokes other sessions; current session preserved.

    Setup: register user; login twice (sessions A and B).
    Action: using session A, GET /api/auth/totp/setup to obtain secret,
            derive a valid TOTP code via pyotp.TOTP(secret).now(),
            POST /api/auth/totp/enable with that code and session A's refresh cookie.
    Assert:
      - response 200 with sessions_revoked == 1
      - session A refresh token still valid (200 on /api/auth/refresh)
      - session B refresh token revoked (401 on /api/auth/refresh)
    """
    email = "cr02_test@example.com"
    await _register_user(revoke_client, handle="cr02user", email=email)

    # Two sequential logins create two separate RefreshToken rows (sessions A and B).
    token_a, refresh_a = await _login_session(revoke_client, email)
    _token_b, refresh_b = await _login_session(revoke_client, email)

    # Obtain TOTP secret for this user
    setup_resp = await revoke_client.get(
        "/api/auth/totp/setup",
        headers={"Authorization": f"Bearer {token_a}"},
    )
    assert setup_resp.status_code == 200, f"TOTP setup failed: {setup_resp.text}"
    secret = setup_resp.json()["secret"]

    # Generate a valid TOTP code from the provisioned secret
    totp_code = pyotp.TOTP(secret).now()

    resp = await revoke_client.post(
        "/api/auth/totp/enable",
        json={"code": totp_code},
        headers={"Authorization": f"Bearer {token_a}"},
        cookies={"refresh_token": refresh_a},
    )

    assert resp.status_code == 200, f"enable_totp failed: {resp.text}"
    data = resp.json()
    assert data["sessions_revoked"] == 1, (
        f"Expected exactly 1 other session revoked, got {data['sessions_revoked']}"
    )

    # Session A's refresh token must still be usable
    status_a = await _try_refresh(revoke_client, refresh_a)
    assert status_a == 200, f"Session A refresh should still be valid, got {status_a}"

    # Session B's refresh token must now be revoked
    status_b = await _try_refresh(revoke_client, refresh_b)
    assert status_b == 401, f"Session B refresh should be revoked (401), got {status_b}"


@pytest.mark.asyncio
async def test_disable_totp_revokes_other_sessions(revoke_client, db_session):
    """CR-03: disable_totp revokes other sessions; current session preserved.

    Setup: register user, enable TOTP directly in DB to avoid a setup session,
           then log in twice with TOTP codes (sessions A and B).
    Action: using session A, DELETE /api/auth/totp with session A's refresh cookie.
    Assert:
      - response 200 with sessions_revoked == 1
      - session A refresh token still valid (200 on /api/auth/refresh)
      - session B refresh token revoked (401 on /api/auth/refresh)
    """
    from sqlalchemy import select
    from db.models import User

    email = "cr03_test@example.com"
    await _register_user(revoke_client, handle="cr03user", email=email)

    # Retrieve the user and enable TOTP directly in the DB so there is no
    # "setup session" refresh token that would skew the sessions_revoked count.
    # Using a known fixed secret makes TOTP code generation deterministic.
    known_secret = "JBSWY3DPEHPK3PXP"
    result = await db_session.execute(select(User).where(User.email == email))
    user = result.scalar_one()
    user.totp_enabled = True
    user.totp_secret = known_secret
    await db_session.commit()

    # Login twice with TOTP, patching verify_totp to bypass the Redis replay
    # prevention that would block reuse of the same code within 90 seconds.
    # We test session revocation here, not TOTP code validation.
    with patch("services.auth.verify_totp", return_value=True):
        # Step 1: submit password — server responds with requires_totp (no TOTP provided)
        step1_a = await revoke_client.post(
            "/api/auth/login",
            json={"email": email, "password": VALID_PASSWORD},
        )
        assert step1_a.status_code == 200, f"Session A login step 1 failed: {step1_a.text}"
        assert step1_a.json().get("requires_totp"), f"Expected requires_totp: {step1_a.json()}"

        # Step 2A: complete session A login with a TOTP code
        step2_a = await revoke_client.post(
            "/api/auth/login",
            json={"email": email, "password": VALID_PASSWORD, "totp_code": "000000"},
        )
        assert step2_a.status_code == 200, f"Session A login step 2 failed: {step2_a.text}"
        token_a = step2_a.json()["access_token"]
        refresh_a = step2_a.cookies.get("refresh_token") or ""

        # Step 2B: complete session B login with a TOTP code (same code OK — verify_totp patched)
        step2_b = await revoke_client.post(
            "/api/auth/login",
            json={"email": email, "password": VALID_PASSWORD, "totp_code": "000000"},
        )
        assert step2_b.status_code == 200, f"Session B login step 2 failed: {step2_b.text}"
        _token_b = step2_b.json()["access_token"]
        refresh_b = step2_b.cookies.get("refresh_token") or ""

    # Disable TOTP using session A — DELETE /api/auth/totp requires no TOTP code in body
    resp = await revoke_client.delete(
        "/api/auth/totp",
        headers={"Authorization": f"Bearer {token_a}"},
        cookies={"refresh_token": refresh_a},
    )

    assert resp.status_code == 200, f"disable_totp failed: {resp.text}"
    data = resp.json()
    assert data["sessions_revoked"] == 1, (
        f"Expected exactly 1 other session revoked, got {data['sessions_revoked']}"
    )

    # Session A's refresh token must still be usable
    status_a = await _try_refresh(revoke_client, refresh_a)
    assert status_a == 200, f"Session A refresh should still be valid, got {status_a}"

    # Session B's refresh token must now be revoked
    status_b = await _try_refresh(revoke_client, refresh_b)
    assert status_b == 401, f"Session B refresh should be revoked (401), got {status_b}"