Make access gating a first-class, pip-consumable extension point so a consuming app can gate any of its own views behind 2FA / ALTCHA / entitlement / feature flag / permission, or gate nothing, without editing the kit. - infrasynth.gates: Gate, GateResult, GatePermission, @gated and built-ins TwoFactorGate, AltchaGate, EntitlementGate, FeatureGate, PermissionGate; denials raise the correct namespaced error/status (per-endpoint, opt-in, default is no gating) - mint a `2fa` JWT claim only after verification (preserved across workspace selection) so TwoFactorGate is meaningful for API/multi-workspace clients - GatePermission added to DEFAULT_PERMISSION_CLASSES; HybridPermission evaluates declared gates so kit permissions gate automatically - document the extension surface and stable import paths in README
320 lines
13 KiB
Python
320 lines
13 KiB
Python
import json
|
|
|
|
import pyotp
|
|
import pytest
|
|
from django.utils import timezone
|
|
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 TestTwoFactorLoginFlow:
|
|
@staticmethod
|
|
def _enable(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 secret
|
|
|
|
def test_login_challenges_for_second_factor(self, db, user):
|
|
self._enable(user)
|
|
client = APIClient()
|
|
resp = client.post(
|
|
"/api/v1/auth/login/",
|
|
{"username": "testuser", "password": "testpass123"},
|
|
format="json",
|
|
)
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
assert resp.json()["requires2fa"] is True
|
|
assert "access_token" not in resp.cookies
|
|
|
|
def test_verify_after_login_issues_tokens(self, db, user):
|
|
secret = self._enable(user)
|
|
client = APIClient()
|
|
client.post("/api/v1/auth/login/", {"username": "testuser", "password": "testpass123"}, format="json")
|
|
resp = client.post(
|
|
"/api/v1/auth/2fa/verify/",
|
|
{"code": pyotp.TOTP(secret).now()},
|
|
format="json",
|
|
)
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
assert "access_token" in resp.cookies
|
|
|
|
def test_no_second_factor_logs_in_directly(self, db, user):
|
|
client = APIClient()
|
|
resp = client.post(
|
|
"/api/v1/auth/login/",
|
|
{"username": "testuser", "password": "testpass123"},
|
|
format="json",
|
|
)
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
assert "access_token" in resp.cookies
|
|
|
|
def test_verify_mints_second_factor_claim(self, db, user):
|
|
from rest_framework_simplejwt.tokens import AccessToken
|
|
|
|
from infrasynth.shared.crypto import decrypt
|
|
|
|
secret = self._enable(user)
|
|
client = APIClient()
|
|
client.post("/api/v1/auth/login/", {"username": "testuser", "password": "testpass123"}, format="json")
|
|
client.post("/api/v1/auth/2fa/verify/", {"code": pyotp.TOTP(secret).now()}, format="json")
|
|
token = AccessToken(decrypt(client.cookies["access_token"].value))
|
|
assert token["2fa"] is True
|
|
|
|
def test_plain_login_has_no_second_factor_claim(self, db, user):
|
|
from rest_framework_simplejwt.tokens import AccessToken
|
|
|
|
from infrasynth.shared.crypto import decrypt
|
|
|
|
client = APIClient()
|
|
client.post("/api/v1/auth/login/", {"username": "testuser", "password": "testpass123"}, format="json")
|
|
token = AccessToken(decrypt(client.cookies["access_token"].value))
|
|
assert token["2fa"] is False
|
|
|
|
|
|
class TestTwoFactorSetup:
|
|
def test_setup_returns_secret_and_qr(self, two_factor_client):
|
|
resp = two_factor_client.post("/api/v1/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/v1/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/v1/auth/2fa/setup/")
|
|
secret = setup.json()["secret"]
|
|
code = pyotp.TOTP(secret).now()
|
|
resp = two_factor_client.post("/api/v1/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/v1/auth/2fa/setup/")
|
|
resp = two_factor_client.post("/api/v1/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/v1/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/v1/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/v1/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/v1/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/v1/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["_2fa_started_at"] = timezone.now().isoformat()
|
|
session.save()
|
|
|
|
code = pyotp.TOTP(secret).now()
|
|
resp = two_factor_client.post("/api/v1/auth/2fa/verify/", {"code": code}, format="json")
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
assert "access_token" in resp.cookies
|
|
|
|
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/v1/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/v1/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["_2fa_pre_auth_token"] = "pre-auth-token"
|
|
session["_2fa_started_at"] = timezone.now().isoformat()
|
|
session.save()
|
|
|
|
resp = two_factor_client.post(
|
|
"/api/v1/auth/2fa/recovery/",
|
|
{"recovery_code": recovery_codes[0]},
|
|
format="json",
|
|
)
|
|
assert resp.status_code == status.HTTP_200_OK
|
|
assert "access_token" in resp.cookies
|
|
|
|
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/v1/auth/2fa/recovery/", {"recovery_code": "RC-INVALID"}, format="json")
|
|
assert resp.status_code == status.HTTP_401_UNAUTHORIZED
|