Gan, Jimmy
f371099a86
- Add comprehensive tests for 2FA setup workflow - Test secret generation, verification, and disabling - Test complete 2FA lifecycle - Add permission tests for authenticated endpoints Results: - TOTP router coverage: 63% → 100% - Total coverage: 45.14% → 45.57% - Total tests: 156 → 164 (8 new tests) - All tests passing
121 lines
4.2 KiB
Python
121 lines
4.2 KiB
Python
"""
|
|
Integration tests for TOTP 2FA operations.
|
|
"""
|
|
import pytest
|
|
import pyotp
|
|
|
|
|
|
class TestTOTPSetup:
|
|
"""Test TOTP 2FA setup endpoints."""
|
|
|
|
def test_generate_2fa_secret(self, test_app, admin_headers):
|
|
"""Test generating a new 2FA secret."""
|
|
response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "secret" in data
|
|
assert "uri" in data
|
|
assert len(data["secret"]) == 32 # Base32 secret length
|
|
assert "otpauth://totp/" in data["uri"]
|
|
assert "NAS" in data["uri"] and "Dashboard" in data["uri"]
|
|
|
|
def test_generate_2fa_without_auth(self, test_app):
|
|
"""Test generating 2FA secret without authentication."""
|
|
response = test_app.post("/api/auth/setup-2fa/generate")
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_verify_2fa_setup_success(self, test_app, admin_headers):
|
|
"""Test verifying and enabling 2FA with valid code."""
|
|
# First generate a secret
|
|
gen_response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
|
|
secret = gen_response.json()["secret"]
|
|
|
|
# Generate a valid TOTP code
|
|
totp = pyotp.TOTP(secret)
|
|
valid_code = totp.now()
|
|
|
|
# Verify the code
|
|
response = test_app.post(
|
|
"/api/auth/setup-2fa/verify",
|
|
headers=admin_headers,
|
|
json={"secret": secret, "code": valid_code}
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "message" in data
|
|
assert "enabled" in data["message"].lower()
|
|
|
|
def test_verify_2fa_setup_invalid_code(self, test_app, admin_headers):
|
|
"""Test verifying 2FA with invalid code."""
|
|
response = test_app.post(
|
|
"/api/auth/setup-2fa/verify",
|
|
headers=admin_headers,
|
|
json={"secret": "JBSWY3DPEHPK3PXP", "code": "000000"}
|
|
)
|
|
|
|
assert response.status_code == 400
|
|
data = response.json()
|
|
assert "detail" in data
|
|
assert "invalid" in data["detail"].lower()
|
|
|
|
def test_verify_2fa_without_auth(self, test_app):
|
|
"""Test verifying 2FA without authentication."""
|
|
response = test_app.post(
|
|
"/api/auth/setup-2fa/verify",
|
|
json={"secret": "JBSWY3DPEHPK3PXP", "code": "123456"}
|
|
)
|
|
|
|
assert response.status_code == 401
|
|
|
|
def test_disable_2fa(self, test_app, admin_headers):
|
|
"""Test disabling 2FA."""
|
|
response = test_app.post("/api/auth/setup-2fa/disable", headers=admin_headers)
|
|
|
|
assert response.status_code == 200
|
|
data = response.json()
|
|
assert "message" in data
|
|
assert "disabled" in data["message"].lower()
|
|
|
|
def test_disable_2fa_without_auth(self, test_app):
|
|
"""Test disabling 2FA without authentication."""
|
|
response = test_app.post("/api/auth/setup-2fa/disable")
|
|
|
|
assert response.status_code == 401
|
|
|
|
|
|
class TestTOTPWorkflow:
|
|
"""Test complete 2FA setup workflow."""
|
|
|
|
def test_complete_2fa_workflow(self, test_app, admin_headers, mock_config, temp_auth_file):
|
|
"""Test the complete 2FA setup and disable workflow."""
|
|
# Step 1: Generate secret
|
|
gen_response = test_app.post("/api/auth/setup-2fa/generate", headers=admin_headers)
|
|
assert gen_response.status_code == 200
|
|
secret = gen_response.json()["secret"]
|
|
|
|
# Step 2: Verify and enable with valid code
|
|
totp = pyotp.TOTP(secret)
|
|
valid_code = totp.now()
|
|
verify_response = test_app.post(
|
|
"/api/auth/setup-2fa/verify",
|
|
headers=admin_headers,
|
|
json={"secret": secret, "code": valid_code}
|
|
)
|
|
assert verify_response.status_code == 200
|
|
|
|
# Step 3: Verify 2FA is enabled by checking auth_service
|
|
from auth_service import load_totp_secret
|
|
saved_secret = load_totp_secret()
|
|
assert saved_secret == secret
|
|
|
|
# Step 4: Disable 2FA
|
|
disable_response = test_app.post("/api/auth/setup-2fa/disable", headers=admin_headers)
|
|
assert disable_response.status_code == 200
|
|
|
|
# Step 5: Verify 2FA is disabled
|
|
saved_secret = load_totp_secret()
|
|
assert saved_secret == ""
|