Skip to content

Commit 280f255

Browse files
Raven95676Raven95676
committed
test: add unit tests for TOTP functionality and recovery code validation
1 parent e1d0501 commit 280f255

2 files changed

Lines changed: 445 additions & 0 deletions

File tree

‎tests/test_dashboard.py‎

Lines changed: 347 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from types import SimpleNamespace
1313
from urllib.parse import parse_qs, urlsplit, urlunsplit
1414

15+
import pyotp
1516
import pytest
1617
import pytest_asyncio
1718
from quart import Quart, jsonify
@@ -28,6 +29,10 @@
2829
verify_dashboard_password,
2930
)
3031
from astrbot.core.utils.pip_installer import PipInstallError
32+
from astrbot.core.utils.totp import (
33+
TOTP_TRUSTED_DEVICE_COOKIE_NAME,
34+
generate_recovery_code,
35+
)
3136
from astrbot.dashboard.password_state import (
3237
get_dashboard_password_hash,
3338
is_password_change_required,
@@ -356,6 +361,348 @@ async def test_auth_login_secure_cookie_override(
356361
assert "SameSite=Strict" in jwt_cookie_header
357362

358363

364+
@pytest.mark.asyncio
365+
async def test_auth_login_requires_totp_when_enabled_and_not_trusted(
366+
app: Quart,
367+
core_lifecycle_td: AstrBotCoreLifecycle,
368+
):
369+
original_dashboard_config = copy.deepcopy(
370+
core_lifecycle_td.astrbot_config["dashboard"]
371+
)
372+
test_client = app.test_client()
373+
_, recovery_code_hash = generate_recovery_code()
374+
secret = pyotp.random_base32()
375+
376+
try:
377+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
378+
"enable": True,
379+
"secret": secret,
380+
"recovery_code_hash": recovery_code_hash,
381+
}
382+
response = await test_client.post(
383+
"/api/auth/login",
384+
json={
385+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
386+
"password": _resolve_dashboard_password(core_lifecycle_td),
387+
},
388+
)
389+
data = await response.get_json()
390+
assert response.status_code == 401
391+
assert data["status"] == "error"
392+
assert data["data"]["totp_required"] is True
393+
finally:
394+
await _restore_dashboard_password_state(
395+
core_lifecycle_td,
396+
original_dashboard_config,
397+
)
398+
399+
400+
@pytest.mark.asyncio
401+
async def test_auth_login_accepts_valid_totp_code(
402+
app: Quart,
403+
core_lifecycle_td: AstrBotCoreLifecycle,
404+
):
405+
original_dashboard_config = copy.deepcopy(
406+
core_lifecycle_td.astrbot_config["dashboard"]
407+
)
408+
test_client = app.test_client()
409+
_, recovery_code_hash = generate_recovery_code()
410+
secret = pyotp.random_base32()
411+
412+
try:
413+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
414+
"enable": True,
415+
"secret": secret,
416+
"recovery_code_hash": recovery_code_hash,
417+
}
418+
response = await test_client.post(
419+
"/api/auth/login",
420+
json={
421+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
422+
"password": _resolve_dashboard_password(core_lifecycle_td),
423+
"code": pyotp.TOTP(secret).now(),
424+
},
425+
)
426+
data = await response.get_json()
427+
assert data["status"] == "ok"
428+
assert "token" in data["data"]
429+
finally:
430+
await _restore_dashboard_password_state(
431+
core_lifecycle_td,
432+
original_dashboard_config,
433+
)
434+
435+
436+
@pytest.mark.asyncio
437+
async def test_auth_login_rejects_invalid_totp_code(
438+
app: Quart,
439+
core_lifecycle_td: AstrBotCoreLifecycle,
440+
):
441+
original_dashboard_config = copy.deepcopy(
442+
core_lifecycle_td.astrbot_config["dashboard"]
443+
)
444+
test_client = app.test_client()
445+
_, recovery_code_hash = generate_recovery_code()
446+
secret = pyotp.random_base32()
447+
448+
try:
449+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
450+
"enable": True,
451+
"secret": secret,
452+
"recovery_code_hash": recovery_code_hash,
453+
}
454+
valid_code = pyotp.TOTP(secret).now()
455+
invalid_code = str((int(valid_code) + 1) % 1_000_000).zfill(6)
456+
response = await test_client.post(
457+
"/api/auth/login",
458+
json={
459+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
460+
"password": _resolve_dashboard_password(core_lifecycle_td),
461+
"code": invalid_code,
462+
},
463+
)
464+
data = await response.get_json()
465+
assert response.status_code == 401
466+
assert data["status"] == "error"
467+
finally:
468+
await _restore_dashboard_password_state(
469+
core_lifecycle_td,
470+
original_dashboard_config,
471+
)
472+
473+
474+
@pytest.mark.asyncio
475+
async def test_auth_login_with_recovery_code_disables_totp(
476+
app: Quart,
477+
core_lifecycle_td: AstrBotCoreLifecycle,
478+
):
479+
original_dashboard_config = copy.deepcopy(
480+
core_lifecycle_td.astrbot_config["dashboard"]
481+
)
482+
test_client = app.test_client()
483+
recovery_code, recovery_code_hash = generate_recovery_code()
484+
secret = pyotp.random_base32()
485+
486+
try:
487+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
488+
"enable": True,
489+
"secret": secret,
490+
"recovery_code_hash": recovery_code_hash,
491+
}
492+
response = await test_client.post(
493+
"/api/auth/login",
494+
json={
495+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
496+
"password": _resolve_dashboard_password(core_lifecycle_td),
497+
"code": recovery_code,
498+
},
499+
)
500+
data = await response.get_json()
501+
assert data["status"] == "ok"
502+
assert core_lifecycle_td.astrbot_config["dashboard"]["totp"] == {
503+
"enable": False,
504+
"secret": "",
505+
"recovery_code_hash": "",
506+
}
507+
finally:
508+
await _restore_dashboard_password_state(
509+
core_lifecycle_td,
510+
original_dashboard_config,
511+
)
512+
513+
514+
@pytest.mark.asyncio
515+
async def test_auth_login_sets_trusted_device_cookie_when_flag_true(
516+
app: Quart,
517+
core_lifecycle_td: AstrBotCoreLifecycle,
518+
):
519+
original_dashboard_config = copy.deepcopy(
520+
core_lifecycle_td.astrbot_config["dashboard"]
521+
)
522+
test_client = app.test_client()
523+
_, recovery_code_hash = generate_recovery_code()
524+
secret = pyotp.random_base32()
525+
526+
try:
527+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
528+
"enable": True,
529+
"secret": secret,
530+
"recovery_code_hash": recovery_code_hash,
531+
}
532+
response = await test_client.post(
533+
"/api/auth/login",
534+
json={
535+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
536+
"password": _resolve_dashboard_password(core_lifecycle_td),
537+
"code": pyotp.TOTP(secret).now(),
538+
"trust_device_flag": True,
539+
},
540+
)
541+
data = await response.get_json()
542+
assert data["status"] == "ok"
543+
set_cookie_headers = response.headers.getlist("Set-Cookie")
544+
trusted_cookie_header = next(
545+
(
546+
value
547+
for value in set_cookie_headers
548+
if TOTP_TRUSTED_DEVICE_COOKIE_NAME in value
549+
),
550+
"",
551+
)
552+
assert trusted_cookie_header
553+
assert "HttpOnly" in trusted_cookie_header
554+
assert "SameSite=Strict" in trusted_cookie_header
555+
assert "Path=/api/auth" in trusted_cookie_header
556+
finally:
557+
await _restore_dashboard_password_state(
558+
core_lifecycle_td,
559+
original_dashboard_config,
560+
)
561+
562+
563+
@pytest.mark.asyncio
564+
async def test_auth_login_skips_totp_when_trusted_cookie_valid(
565+
app: Quart,
566+
core_lifecycle_td: AstrBotCoreLifecycle,
567+
):
568+
original_dashboard_config = copy.deepcopy(
569+
core_lifecycle_td.astrbot_config["dashboard"]
570+
)
571+
test_client = app.test_client()
572+
_, recovery_code_hash = generate_recovery_code()
573+
secret = pyotp.random_base32()
574+
575+
try:
576+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
577+
"enable": True,
578+
"secret": secret,
579+
"recovery_code_hash": recovery_code_hash,
580+
}
581+
first_login = await test_client.post(
582+
"/api/auth/login",
583+
json={
584+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
585+
"password": _resolve_dashboard_password(core_lifecycle_td),
586+
"code": pyotp.TOTP(secret).now(),
587+
"trust_device_flag": True,
588+
},
589+
)
590+
first_data = await first_login.get_json()
591+
assert first_data["status"] == "ok"
592+
593+
second_login = await test_client.post(
594+
"/api/auth/login",
595+
json={
596+
"username": core_lifecycle_td.astrbot_config["dashboard"]["username"],
597+
"password": _resolve_dashboard_password(core_lifecycle_td),
598+
},
599+
)
600+
second_data = await second_login.get_json()
601+
assert second_login.status_code == 200
602+
assert second_data["status"] == "ok"
603+
finally:
604+
await _restore_dashboard_password_state(
605+
core_lifecycle_td,
606+
original_dashboard_config,
607+
)
608+
609+
610+
@pytest.mark.asyncio
611+
async def test_auth_totp_disable_by_totp_code(
612+
app: Quart,
613+
authenticated_header: dict,
614+
core_lifecycle_td: AstrBotCoreLifecycle,
615+
):
616+
original_dashboard_config = copy.deepcopy(
617+
core_lifecycle_td.astrbot_config["dashboard"]
618+
)
619+
test_client = app.test_client()
620+
_, recovery_code_hash = generate_recovery_code()
621+
secret = pyotp.random_base32()
622+
623+
try:
624+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
625+
"enable": True,
626+
"secret": secret,
627+
"recovery_code_hash": recovery_code_hash,
628+
}
629+
response = await test_client.post(
630+
"/api/auth/totp/disable",
631+
headers=authenticated_header,
632+
json={"code": pyotp.TOTP(secret).now()},
633+
)
634+
data = await response.get_json()
635+
assert data["status"] == "ok"
636+
assert core_lifecycle_td.astrbot_config["dashboard"]["totp"] == {
637+
"enable": False,
638+
"secret": "",
639+
"recovery_code_hash": "",
640+
}
641+
finally:
642+
await _restore_dashboard_password_state(
643+
core_lifecycle_td,
644+
original_dashboard_config,
645+
)
646+
647+
648+
@pytest.mark.asyncio
649+
async def test_auth_totp_verify_setup_with_valid_code_returns_recovery_code(
650+
app: Quart,
651+
authenticated_header: dict,
652+
):
653+
test_client = app.test_client()
654+
secret = pyotp.random_base32()
655+
response = await test_client.post(
656+
"/api/auth/totp/verify-setup",
657+
headers=authenticated_header,
658+
json={"secret": secret, "code": pyotp.TOTP(secret).now()},
659+
)
660+
data = await response.get_json()
661+
assert data["status"] == "ok"
662+
assert isinstance(data["data"]["recovery_code"], str)
663+
assert isinstance(data["data"]["recovery_code_hash"], str)
664+
assert data["data"]["recovery_code"]
665+
assert data["data"]["recovery_code_hash"]
666+
667+
668+
@pytest.mark.asyncio
669+
async def test_auth_totp_disable_by_recovery_code(
670+
app: Quart,
671+
authenticated_header: dict,
672+
core_lifecycle_td: AstrBotCoreLifecycle,
673+
):
674+
original_dashboard_config = copy.deepcopy(
675+
core_lifecycle_td.astrbot_config["dashboard"]
676+
)
677+
test_client = app.test_client()
678+
recovery_code, recovery_code_hash = generate_recovery_code()
679+
secret = pyotp.random_base32()
680+
681+
try:
682+
core_lifecycle_td.astrbot_config["dashboard"]["totp"] = {
683+
"enable": True,
684+
"secret": secret,
685+
"recovery_code_hash": recovery_code_hash,
686+
}
687+
response = await test_client.post(
688+
"/api/auth/totp/disable",
689+
headers=authenticated_header,
690+
json={"code": recovery_code},
691+
)
692+
data = await response.get_json()
693+
assert data["status"] == "ok"
694+
assert core_lifecycle_td.astrbot_config["dashboard"]["totp"] == {
695+
"enable": False,
696+
"secret": "",
697+
"recovery_code_hash": "",
698+
}
699+
finally:
700+
await _restore_dashboard_password_state(
701+
core_lifecycle_td,
702+
original_dashboard_config,
703+
)
704+
705+
359706
@pytest.mark.asyncio
360707
async def test_legacy_md5_dashboard_password_keeps_legacy_auth_until_edit(
361708
app: Quart,

0 commit comments

Comments
 (0)