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