TelegramBotAdmin/tests/test_admin_api_auth.py

from __future__ import annotations

from datetime import UTC, datetime, timedelta

from aiohttp import CookieJar
from aiohttp.test_utils import TestClient, TestServer


async def test_login_forced_password_change_csrf_and_session(app_modules):
    admin_api = app_modules.load("wbb.admin.api")
    application = admin_api.build_admin_application()
    await application["admin_api"].initialize()
    client = TestClient(TestServer(application), cookie_jar=CookieJar(unsafe=True))
    await client.start_server()
    try:
        invalid = await client.post(
            "/api/admin/v1/auth/login",
            json={"username": "admin", "password": "wrong"},
        )
        assert invalid.status == 401
        assert (await invalid.json())["error"]["code"] == "invalid_credentials"

        login = await client.post(
            "/api/admin/v1/auth/login",
            json={"username": "admin", "password": "qwe0.123456"},
        )
        assert login.status == 200
        login_data = (await login.json())["data"]
        assert login_data["must_change_password"] is True

        blocked = await client.get("/api/admin/v1/dashboard")
        assert blocked.status == 428
        assert (await blocked.json())["error"]["code"] == "password_change_required"

        no_csrf = await client.put(
            "/api/admin/v1/auth/password",
            json={
                "current_password": "qwe0.123456",
                "new_password": "changed-pass-123",
            },
        )
        assert no_csrf.status == 403

        changed = await client.put(
            "/api/admin/v1/auth/password",
            headers={"X-CSRF-Token": login_data["csrf_token"]},
            json={
                "current_password": "qwe0.123456",
                "new_password": "changed-pass-123",
            },
        )
        assert changed.status == 200
        changed_data = (await changed.json())["data"]
        assert changed_data["must_change_password"] is False

        unconfirmed = await client.patch(
            "/api/admin/v1/chats/-100/profile",
            headers={"X-CSRF-Token": changed_data["csrf_token"]},
            json={"title": "New title"},
        )
        assert unconfirmed.status == 409
        assert (await unconfirmed.json())["error"]["code"] == "confirmation_required"

        dashboard = await client.get("/api/admin/v1/dashboard")
        assert dashboard.status == 200
        assert "counts" in (await dashboard.json())["data"]

        await app_modules.wbb.db.managed_chats.insert_one(
            {
                "bot_id": "primary",
                "chat_id": -100,
                "title": "机器人已退出的群",
                "accessible": False,
            }
        )
        delete_without_csrf = await client.delete("/api/admin/v1/chats/-100")
        assert delete_without_csrf.status == 403
        deleted = await client.delete(
            "/api/admin/v1/chats/-100",
            headers={"X-CSRF-Token": changed_data["csrf_token"]},
        )
        assert deleted.status == 200
        assert (await deleted.json())["data"] == {
            "chat_id": "-100",
            "removed": True,
        }
        assert await app_modules.wbb.db.managed_chats.find_one({"chat_id": -100}) is None

        logout = await client.post(
            "/api/admin/v1/auth/logout",
            headers={"X-CSRF-Token": changed_data["csrf_token"]},
        )
        assert logout.status == 200
        unauthenticated = await client.get("/api/admin/v1/dashboard")
        assert unauthenticated.status == 401
    finally:
        await client.close()


async def test_login_is_rate_limited_with_mongo_naive_datetime(app_modules):
    admin_api = app_modules.load("wbb.admin.api")
    application = admin_api.build_admin_application()
    await application["admin_api"].initialize()
    client = TestClient(TestServer(application), cookie_jar=CookieJar(unsafe=True))
    await client.start_server()
    try:
        for _ in range(5):
            response = await client.post(
                "/api/admin/v1/auth/login",
                json={"username": "admin", "password": "wrong"},
            )
            assert response.status == 401

        await app_modules.wbb.db.admin_users.update_one(
            {"username": "admin"},
            {"$set": {"last_failed_login_at": datetime.now(UTC).replace(tzinfo=None)}},
        )
        limited = await client.post(
            "/api/admin/v1/auth/login",
            json={"username": "admin", "password": "qwe0.123456"},
        )
        assert limited.status == 429
        assert (await limited.json())["error"]["code"] == "login_rate_limited"

        expired_at = datetime.now(UTC) - timedelta(minutes=16)
        await app_modules.wbb.db.admin_users.update_one(
            {"username": "admin"},
            {"$set": {"last_failed_login_at": expired_at.replace(tzinfo=None)}},
        )
        login = await client.post(
            "/api/admin/v1/auth/login",
            json={"username": "admin", "password": "qwe0.123456"},
        )
        assert login.status == 200
    finally:
        await client.close()