243 lines
9.5 KiB
Python
243 lines
9.5 KiB
Python
import json
|
|
|
|
import pyotp
|
|
import pytest
|
|
from rest_framework import status
|
|
from rest_framework.test import APIClient
|
|
|
|
from infrasynth.security.models import TwoFactorConfig
|
|
from infrasynth.security.two_factor.services import RecoveryCodeService, TOTPService
|
|
|
|
|
|
@pytest.fixture
|
|
def totp_service():
|
|
return TOTPService()
|
|
|
|
|
|
@pytest.fixture
|
|
def recovery_service():
|
|
return RecoveryCodeService()
|
|
|
|
|
|
class TestTOTPService:
|
|
def test_generate_secret_is_base32(self, totp_service):
|
|
secret = totp_service.generate_secret()
|
|
assert len(secret) >= 16
|
|
pyotp.TOTP(secret).now()
|
|
|
|
def test_generate_secret_unique(self, totp_service):
|
|
assert totp_service.generate_secret() != totp_service.generate_secret()
|
|
|
|
def test_provisioning_uri_contains_issuer_and_email(self, totp_service):
|
|
uri = totp_service.get_provisioning_uri("SECRETBASE32", "user@example.com")
|
|
assert "user%40example.com" in uri
|
|
assert "InfraSynth" in uri
|
|
|
|
def test_qr_base64_is_valid_png(self, totp_service):
|
|
import base64
|
|
|
|
qr = totp_service.generate_qr_base64("SECRETBASE32", "user@example.com")
|
|
raw = base64.b64decode(qr)
|
|
assert raw[:8] == b"\x89PNG\r\n\x1a\n"
|
|
|
|
def test_verify_correct_code(self, totp_service):
|
|
secret = totp_service.generate_secret()
|
|
code = pyotp.TOTP(secret).now()
|
|
assert totp_service.verify(secret, code) is True
|
|
|
|
def test_verify_wrong_code(self, totp_service):
|
|
secret = totp_service.generate_secret()
|
|
assert totp_service.verify(secret, "000000") is False
|
|
|
|
def test_encrypt_decrypt_roundtrip(self, totp_service):
|
|
secret = "MYTESTSECRET"
|
|
encrypted = totp_service.encrypt_secret(secret)
|
|
assert encrypted != secret
|
|
assert totp_service.decrypt_secret(encrypted) == secret
|
|
|
|
|
|
class TestRecoveryCodeService:
|
|
def test_generate_codes_count(self, recovery_service):
|
|
codes = recovery_service.generate_codes()
|
|
assert len(codes) == 8
|
|
assert len(set(codes)) == 8
|
|
|
|
def test_codes_have_prefix(self, recovery_service):
|
|
codes = recovery_service.generate_codes()
|
|
assert all(c.startswith("RC-") for c in codes)
|
|
|
|
def test_encrypt_decrypt_roundtrip(self, recovery_service):
|
|
codes = ["RC-AAA-BBB"]
|
|
encrypted = recovery_service.encrypt_codes(codes)
|
|
assert encrypted != json.dumps(codes)
|
|
assert recovery_service.decrypt_codes(encrypted) == codes
|
|
|
|
def test_verify_code(self, recovery_service):
|
|
codes = recovery_service.generate_codes()
|
|
encrypted = recovery_service.encrypt_codes(codes)
|
|
assert recovery_service.verify_code(codes[0], encrypted) is True
|
|
assert recovery_service.verify_code("RC-WRONG", encrypted) is False
|
|
|
|
def test_verify_corrupted_data_returns_false(self, recovery_service):
|
|
assert recovery_service.verify_code("RC-X", "garbage-data") is False
|
|
|
|
def test_remove_used_code(self, recovery_service):
|
|
codes = recovery_service.generate_codes()
|
|
encrypted = recovery_service.encrypt_codes(codes)
|
|
used = codes[0]
|
|
updated = recovery_service.remove_used_code(used, encrypted)
|
|
remaining = recovery_service.decrypt_codes(updated)
|
|
assert used not in remaining
|
|
assert len(remaining) == len(codes) - 1
|
|
|
|
def test_remove_nonexistent_returns_none(self, recovery_service):
|
|
codes = recovery_service.generate_codes()
|
|
encrypted = recovery_service.encrypt_codes(codes)
|
|
assert recovery_service.remove_used_code("RC-NOPE", encrypted) is None
|
|
|
|
def test_remove_from_corrupted_returns_none(self, recovery_service):
|
|
assert recovery_service.remove_used_code("RC-X", "garbage") is None
|
|
|
|
|
|
@pytest.fixture
|
|
def two_factor_client(user):
|
|
client = APIClient()
|
|
client.force_authenticate(user=user)
|
|
return client
|
|
|
|
|
|
class TestTwoFactorSetup:
|
|
def test_setup_returns_secret_and_qr(self, two_factor_client):
|
|
resp = two_factor_client.post("/api/auth/2fa/setup/")
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
data = resp.json()
|
|
assert "secret" in data
|
|
assert "qr_base64" in data
|
|
assert "provisioning_uri" in data
|
|
|
|
def test_setup_requires_auth(self, api_client):
|
|
resp = api_client.post("/api/auth/2fa/setup/")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
|
|
class TestTwoFactorVerifySetup:
|
|
def test_verify_setup_complete_flow(self, two_factor_client, user):
|
|
setup = two_factor_client.post("/api/auth/2fa/setup/")
|
|
secret = setup.json()["secret"]
|
|
code = pyotp.TOTP(secret).now()
|
|
resp = two_factor_client.post("/api/auth/2fa/verify-setup/", {"code": code}, format="json")
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
assert len(resp.json()["recovery_codes"]) == 8
|
|
|
|
config = TwoFactorConfig.objects.get(user=user)
|
|
assert config.is_enabled is True
|
|
assert config.is_configured is True
|
|
assert config.secret_key_encrypted
|
|
|
|
def test_verify_setup_wrong_code(self, two_factor_client):
|
|
two_factor_client.post("/api/auth/2fa/setup/")
|
|
resp = two_factor_client.post("/api/auth/2fa/verify-setup/", {"code": "000000"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
def test_verify_setup_without_session(self, two_factor_client):
|
|
resp = two_factor_client.post("/api/auth/2fa/verify-setup/", {"code": "123456"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
|
|
class TestTwoFactorDisable:
|
|
def test_disable_without_config(self, two_factor_client):
|
|
resp = two_factor_client.post("/api/auth/2fa/disable/")
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
|
|
def test_disable_with_config(self, two_factor_client, user):
|
|
TwoFactorConfig.objects.create(
|
|
user=user,
|
|
is_enabled=True,
|
|
is_configured=True,
|
|
secret_key_encrypted="enc",
|
|
recovery_codes_encrypted="enc",
|
|
)
|
|
resp = two_factor_client.post("/api/auth/2fa/disable/")
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
user.two_factor_config.refresh_from_db()
|
|
assert user.two_factor_config.is_enabled is False
|
|
assert user.two_factor_config.secret_key_encrypted is None
|
|
|
|
|
|
class TestTwoFactorVerifyAndRecovery:
|
|
@pytest.fixture
|
|
def configured_user(self, db, two_factor_client, user):
|
|
secret = TOTPService().generate_secret()
|
|
rcs = RecoveryCodeService()
|
|
TwoFactorConfig.objects.create(
|
|
user=user,
|
|
is_enabled=True,
|
|
is_configured=True,
|
|
secret_key_encrypted=TOTPService().encrypt_secret(secret),
|
|
recovery_codes_encrypted=rcs.encrypt_codes(rcs.generate_codes()),
|
|
)
|
|
return user, secret
|
|
|
|
def test_verify_requires_pre_auth_session(self, two_factor_client, configured_user):
|
|
resp = two_factor_client.post("/api/auth/2fa/verify/", {"code": "123456"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
def test_recovery_requires_pre_auth_session(self, two_factor_client, configured_user):
|
|
resp = two_factor_client.post("/api/auth/2fa/recovery/", {"recovery_code": "RC-X"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
def test_verify_with_correct_code(self, two_factor_client, configured_user):
|
|
user, secret = configured_user
|
|
session = two_factor_client.session
|
|
session["_2fa_user_id"] = user.pk
|
|
session["_2fa_pre_auth_token"] = "pre-auth-token"
|
|
session.save()
|
|
|
|
code = pyotp.TOTP(secret).now()
|
|
resp = two_factor_client.post("/api/auth/2fa/verify/", {"code": code}, format="json")
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
|
|
def test_verify_with_wrong_code(self, two_factor_client, configured_user):
|
|
user, secret = configured_user
|
|
session = two_factor_client.session
|
|
session["_2fa_user_id"] = user.pk
|
|
session.save()
|
|
|
|
resp = two_factor_client.post("/api/auth/2fa/verify/", {"code": "000000"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
def test_verify_without_2fa_config(self, two_factor_client, user):
|
|
session = two_factor_client.session
|
|
session["_2fa_user_id"] = user.pk
|
|
session.save()
|
|
|
|
resp = two_factor_client.post("/api/auth/2fa/verify/", {"code": "123456"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|
|
|
|
def test_recovery_with_valid_code(self, two_factor_client, configured_user):
|
|
user, secret = configured_user
|
|
config = TwoFactorConfig.objects.get(user=user)
|
|
recovery_codes = RecoveryCodeService().decrypt_codes(config.recovery_codes_encrypted)
|
|
session = two_factor_client.session
|
|
session["_2fa_user_id"] = user.pk
|
|
session.save()
|
|
|
|
resp = two_factor_client.post(
|
|
"/api/auth/2fa/recovery/",
|
|
{"recovery_code": recovery_codes[0]},
|
|
format="json",
|
|
)
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
|
|
config.refresh_from_db()
|
|
remaining = RecoveryCodeService().decrypt_codes(config.recovery_codes_encrypted)
|
|
assert recovery_codes[0] not in remaining
|
|
|
|
def test_recovery_with_invalid_code(self, two_factor_client, configured_user):
|
|
user, secret = configured_user
|
|
session = two_factor_client.session
|
|
session["_2fa_user_id"] = user.pk
|
|
session.save()
|
|
|
|
resp = two_factor_client.post("/api/auth/2fa/recovery/", {"recovery_code": "RC-INVALID"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|