Files
wecom_it_smart_desk/backend/tests/test_otp_unified.py
T

393 lines
14 KiB
Python
Raw Normal View History

# =============================================================================
# 三端认证重构 AUTH-03 — 统一 OTP 路由测试
# =============================================================================
# 验证 otp.py 的 6 个端点:
# GET /api/auth/otp-status — 查询绑定状态
# POST /api/auth/otp-bind — 生成 secret + 二维码
# POST /api/auth/otp-verify — 验证 OTP(写 Redis 30 分钟)
# POST /api/auth/otp-unbind — 用户主动关闭 OTP
# POST /api/auth/otp-admin-reset/{id} — 管理员重置
# GET /api/auth/otp-admin-users — 管理员查看全部坐席状态
# =============================================================================
import pyotp
import pytest
from sqlalchemy import select
from app.models.agent import Agent
from app.services.mfa_service import MFA_VERIFIED_TTL_SECONDS
from tests.conftest import create_test_agent
# -----------------------------------------------------------------------------
# 辅助函数
# -----------------------------------------------------------------------------
async def _login_and_get_token(client, user_id: str, name: str, otp_code: str = None) -> str:
"""调用 /agents/login 获取 token
Args:
client: 测试客户端
user_id: 用户 ID
name: 用户名
otp_code: 可选的 OTP 验证码(当登录返回 require_otp: true 时需要提供)
"""
payload = {"user_id": user_id, "name": name}
if otp_code:
payload["otp_code"] = otp_code
response = await client.post(
"/agents/login",
json=payload,
)
assert response.status_code == 200, f"登录失败: {response.text}"
body = response.json()
assert body.get("code") == 0, f"登录业务码非 0: {body}"
# 如果返回 require_otp,说明需要 OTP 验证
if body["data"].get("require_otp"):
# 需要先绑定/启用 OTP,这里返回 None 表示需要 OTP
return None
assert "token" in body["data"], f"登录响应没有 token: {body}"
return body["data"]["token"]
def _bearer(token: str) -> dict:
return {"Authorization": f"Bearer {token}"}
async def _seed_admin_role(db_session, employee_id: str) -> str:
"""为用户分配 admin 角色"""
from app.models.role import Role
from app.models.user_role import UserRole
import uuid
from datetime import datetime
# 1. 确保 admin 角色存在
stmt = select(Role).where(Role.name == "admin")
role = (await db_session.execute(stmt)).scalars().first()
if not role:
role = Role(
id=str(uuid.uuid4()),
name="admin",
display_name="管理员",
is_default=False,
permissions=[],
)
db_session.add(role)
await db_session.flush()
# 2. 建立关联
stmt = select(UserRole).where(
UserRole.employee_id == employee_id,
UserRole.role_id == role.id,
)
existing = (await db_session.execute(stmt)).scalars().first()
if not existing:
user_role = UserRole(
id=str(uuid.uuid4()),
employee_id=employee_id,
role_id=role.id,
source="manual",
assigned_at=datetime.now(),
)
db_session.add(user_role)
await db_session.flush()
return role.id
# =============================================================================
# 1. GET /api/auth/otp-status
# =============================================================================
class TestOTPSatus:
"""GET /api/auth/otp-status 测试"""
@pytest.mark.asyncio
async def test_new_user_status_unbound(self, client, db_session):
"""全新用户 → bound=false, enabled=false"""
agent = create_test_agent(user_id="alice_001", name="Alice")
db_session.add(agent)
await db_session.flush()
token = await _login_and_get_token(client, "alice_001", "Alice")
resp = await client.get("/auth/otp-status", headers=_bearer(token))
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
data = body["data"]
assert data["bound"] is False
assert data["enabled"] is False
assert data["verified"] is False
@pytest.mark.asyncio
async def test_bound_user_status(self, client, db_session):
"""已绑定用户 → bound=true, enabled=true
注意:已绑定用户的 mfa_enabled=True, mfa_secret 有值
这样的用户登录时会返回 require_otp: true
"""
from datetime import datetime
agent = create_test_agent(user_id="bob_001", name="Bob")
# 直接设置已绑定状态(mfa_enabled=True, mfa_secret 有值)
agent.mfa_secret = pyotp.random_base32()
agent.mfa_enabled = True
agent.mfa_bound_at = datetime.now()
db_session.add(agent)
await db_session.flush()
# 使用 OTP 验证码登录获取 token
otp_code = pyotp.TOTP(agent.mfa_secret).now()
token = await _login_and_get_token(client, "bob_001", "Bob", otp_code)
# 查询状态
resp = await client.get("/auth/otp-status", headers=_bearer(token))
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
data = body["data"]
assert data["bound"] is True
assert data["enabled"] is True
# =============================================================================
# 2. POST /api/auth/otp-bind
# =============================================================================
class TestOTPBind:
"""POST /api/auth/otp-bind 测试"""
@pytest.mark.asyncio
async def test_bind_returns_secret_and_qrcode(self, client, db_session):
"""返回 secret + otpauth_url + qr_code_base64"""
agent = create_test_agent(user_id="carol_001", name="Carol")
db_session.add(agent)
await db_session.flush()
token = await _login_and_get_token(client, "carol_001", "Carol")
resp = await client.post("/auth/otp-bind", headers=_bearer(token))
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
data = body["data"]
assert "secret" in data
assert "otpauth_url" in data
assert "qr_code_base64" in data
assert len(data["secret"]) == 32
@pytest.mark.asyncio
async def test_bind_already_enabled_rejected(self, client, db_session):
"""已启用则拒绝重新绑定
已绑定 MFA 的用户需要通过 otp-verify 流程来获取 token
"""
from datetime import datetime
agent = create_test_agent(user_id="dave_001", name="Dave")
secret = pyotp.random_base32()
agent.mfa_secret = secret
agent.mfa_enabled = True
agent.mfa_bound_at = datetime.now()
db_session.add(agent)
await db_session.flush()
# 使用 OTP 验证码登录
otp_code = pyotp.TOTP(secret).now()
token = await _login_and_get_token(client, "dave_001", "Dave", otp_code)
resp = await client.post("/auth/otp-bind", headers=_bearer(token))
assert resp.status_code == 200
body = resp.json()
assert body["code"] != 0
# =============================================================================
# 3. POST /api/auth/otp-verify
# =============================================================================
class TestOTPVerify:
"""POST /api/auth/otp-verify 测试"""
@pytest.mark.asyncio
async def test_verify_correct_code(self, client, db_session):
"""正确码 → verified=True + 返回 expires_in"""
from datetime import datetime
agent = create_test_agent(user_id="eve_001", name="Eve")
secret = pyotp.random_base32()
agent.mfa_secret = secret
agent.mfa_enabled = True
agent.mfa_bound_at = datetime.now()
db_session.add(agent)
await db_session.flush()
# 使用 OTP 验证码登录
otp_code = pyotp.TOTP(secret).now()
token = await _login_and_get_token(client, "eve_001", "Eve", otp_code)
# 验证 OTP(使用新的验证码)
new_otp_code = pyotp.TOTP(secret).now()
resp = await client.post(
"/auth/otp-verify",
headers=_bearer(token),
json={"otp_code": new_otp_code},
)
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
data = body["data"]
assert data["verified"] is True
assert data["expires_in"] == MFA_VERIFIED_TTL_SECONDS
@pytest.mark.asyncio
async def test_verify_wrong_code(self, client, db_session):
"""错误码 → verified=False"""
from datetime import datetime
agent = create_test_agent(user_id="frank_001", name="Frank")
secret = pyotp.random_base32()
agent.mfa_secret = secret
agent.mfa_enabled = True
agent.mfa_bound_at = datetime.now()
db_session.add(agent)
await db_session.flush()
# 使用 OTP 验证码登录
otp_code = pyotp.TOTP(secret).now()
token = await _login_and_get_token(client, "frank_001", "Frank", otp_code)
# 使用错误的 OTP 验证
resp = await client.post(
"/auth/otp-verify",
headers=_bearer(token),
json={"otp_code": "000000"},
)
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
assert body["data"]["verified"] is False
# =============================================================================
# 4. POST /api/auth/otp-unbind
# =============================================================================
class TestOTPUnbind:
"""POST /api/auth/otp-unbind 测试"""
@pytest.mark.asyncio
async def test_unbind_correct_code(self, client, db_session):
"""正确 OTP → 清空 secret + enabled=False"""
from datetime import datetime
agent = create_test_agent(user_id="grace_001", name="Grace")
secret = pyotp.random_base32()
agent.mfa_secret = secret
agent.mfa_enabled = True
agent.mfa_bound_at = datetime.now()
db_session.add(agent)
await db_session.flush()
# 使用 OTP 验证码登录
otp_code = pyotp.TOTP(secret).now()
token = await _login_and_get_token(client, "grace_001", "Grace", otp_code)
# 解绑 OTP
resp = await client.post(
"/auth/otp-unbind",
headers=_bearer(token),
json={"otp_code": otp_code},
)
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
assert body["data"]["success"] is True
# DB 验证
stmt = select(Agent).where(Agent.user_id == "grace_001")
db_agent = (await db_session.execute(stmt)).scalars().first()
assert db_agent.mfa_secret is None
assert db_agent.mfa_enabled is False
# =============================================================================
# 5. POST /api/auth/otp-admin-reset/{employee_id}
# =============================================================================
class TestOTPAdminReset:
"""POST /api/auth/otp-admin-reset/{employee_id} 测试"""
@pytest.mark.asyncio
async def test_admin_reset_target_user(self, client, db_session):
"""管理员重置 → 目标用户清空"""
from datetime import datetime
# 目标用户
target = create_test_agent(user_id="henry_001", name="Henry")
target.mfa_secret = pyotp.random_base32()
target.mfa_enabled = True
target.mfa_bound_at = datetime.now()
db_session.add(target)
# 管理员
admin = create_test_agent(user_id="admin_001", name="Admin")
db_session.add(admin)
await db_session.flush()
await _seed_admin_role(db_session, "admin_001")
admin_token = await _login_and_get_token(client, "admin_001", "Admin")
resp = await client.post(
f"/auth/otp-admin-reset/henry_001",
headers=_bearer(admin_token),
)
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
assert body["data"]["success"] is True
# DB 验证
stmt = select(Agent).where(Agent.user_id == "henry_001")
db_target = (await db_session.execute(stmt)).scalars().first()
assert db_target.mfa_secret is None
assert db_target.mfa_enabled is False
# =============================================================================
# 6. GET /api/auth/otp-admin-users
# =============================================================================
class TestOTPAdminUsers:
"""GET /api/auth/otp-admin-users 测试"""
@pytest.mark.asyncio
async def test_admin_list_users(self, client, db_session):
"""管理员查看 → 返回所有用户"""
# 添加两个用户
agent1 = create_test_agent(user_id="user_001", name="User1")
agent2 = create_test_agent(user_id="user_002", name="User2")
db_session.add(agent1)
db_session.add(agent2)
# 管理员
admin = create_test_agent(user_id="admin_002", name="Admin2")
db_session.add(admin)
await db_session.flush()
await _seed_admin_role(db_session, "admin_002")
admin_token = await _login_and_get_token(client, "admin_002", "Admin2")
resp = await client.get("/auth/otp-admin-users", headers=_bearer(admin_token))
assert resp.status_code == 200
body = resp.json()
assert body["code"] == 0
# 应该包含所有用户
users = body["data"]
assert len(users) >= 2
if __name__ == "__main__":
pytest.main([__file__, "-v"])