TelegramBotAdmin/tests/test_admin_api_auth.py

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from types import SimpleNamespace

from aiohttp import CookieJar
from aiohttp.test_utils import TestClient, TestServer
from pyrogram.enums import ChatMemberStatus


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

        giveaway_db = app_modules.load("wbb.utils.dbgiveaway")
        now = datetime.now(UTC)
        giveaway = await giveaway_db.create_giveaway(
            chat_id=-100,
            creator_id=1,
            creator_name="Admin",
            title="Original",
            description="",
            prizes=[{"name": "Winner", "count": 1}],
            starts_at=now + timedelta(hours=1),
            ends_at=now + timedelta(hours=2),
        )
        await giveaway_db.attach_giveaway_message(
            giveaway["giveaway_id"], -100, 88
        )
        giveaway = await giveaway_db.get_giveaway(giveaway["giveaway_id"])
        app_modules.app.members[(-100, 999)] = SimpleNamespace(
            status=ChatMemberStatus.OWNER
        )
        expected_updated_at = giveaway["updated_at"]
        if expected_updated_at.tzinfo is None:
            expected_updated_at = expected_updated_at.replace(tzinfo=UTC)
        update_body = {
            "title": "Updated",
            "description": "Valid for one month",
            "starts_at": (now + timedelta(hours=1)).isoformat(),
            "ends_at": (now + timedelta(hours=3)).isoformat(),
            "expected_updated_at": expected_updated_at.isoformat(),
            "confirm": True,
        }
        update_without_csrf = await client.patch(
            f"/api/admin/v1/giveaways/{giveaway['giveaway_id']}",
            json=update_body,
        )
        assert update_without_csrf.status == 403
        update_without_confirmation = await client.patch(
            f"/api/admin/v1/giveaways/{giveaway['giveaway_id']}",
            headers={"X-CSRF-Token": changed_data["csrf_token"]},
            json={**update_body, "confirm": False},
        )
        assert update_without_confirmation.status == 409
        updated = await client.patch(
            f"/api/admin/v1/giveaways/{giveaway['giveaway_id']}",
            headers={"X-CSRF-Token": changed_data["csrf_token"]},
            json=update_body,
        )
        assert updated.status == 200
        updated_data = (await updated.json())["data"]
        assert updated_data["title"] == "Updated"
        assert updated_data["description"] == "Valid for one month"
        assert app_modules.app.edited_messages[-1][:2] == (-100, 88)
        audit = await app_modules.wbb.db.admin_audit_logs.find_one(
            {"action": "giveaway.update", "target_id": giveaway["giveaway_id"]}
        )
        assert audit["success"] is True

        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()