apps/security/tests.py

"""
Tests for the security app — Phase 2.

Covers MFA setup/verify, session management, rate limiting,
backup codes, and trusted devices.
"""
from django.test import TestCase, override_settings
from django.urls import reverse
from django.contrib.auth import get_user_model
import jwt

from apps.accounts.services import register_user
from apps.security.models import (
    BackupCode,
    MFAFactor,
    RateLimit,
    SecurityEvent,
    Session,
    TrustedDevice,
)
from apps.security.services import (
    check_rate_limit,
    create_session,
    disable_totp,
    generate_backup_codes,
    get_remaining_backup_codes_count,
    global_logout,
    is_mfa_enabled,
    is_trusted_device,
    record_attempt,
    record_success,
    register_trusted_device,
    revoke_all_trusted_devices,
    revoke_session,
    setup_totp,
    verify_backup_code,
    jwt_service,
)
from core.exceptions import RateLimitError, AuthenticationError
from core.services.encryption import decrypt_field, encrypt_field

User = get_user_model()


# ─── Encryption Tests ────────────────────────────────────────────────────────


class EncryptionTests(TestCase):
    """Test core encryption service."""

    def test_encrypt_decrypt_roundtrip(self):
        """Encrypting then decrypting returns original value."""
        original = "my_secret_totp_key_abc123"
        encrypted = encrypt_field(original)
        self.assertNotEqual(encrypted, original)
        decrypted = decrypt_field(encrypted)
        self.assertEqual(decrypted, original)

    def test_tampered_ciphertext_fails(self):
        """Tampered ciphertext raises ValueError."""
        encrypted = encrypt_field("test_secret")
        tampered = encrypted[:-3] + "XXX"
        with self.assertRaises(Exception):
            decrypt_field(tampered)


# ─── MFA Tests ───────────────────────────────────────────────────────────────


class MFASetupTests(TestCase):
    """Test TOTP MFA setup flow."""

    def setUp(self):
        self.user = register_user("mfa@example.com", "StrongPass123!", org_name="Test Org", business_type="SaaS")

    def test_setup_totp_returns_secret(self):
        """Setup returns secret and factor ID."""
        data = setup_totp(self.user)
        self.assertIn("secret", data)
        self.assertIn("factor_id", data)
        self.assertEqual(data["account"], "mfa@example.com")

    def test_mfa_not_enabled_after_setup(self):
        """MFA is not enabled until verified."""
        setup_totp(self.user)
        self.assertFalse(is_mfa_enabled(self.user))

    def test_disable_totp(self):
        """Disabling removes enabled flag."""
        setup_totp(self.user)
        disable_totp(self.user)
        self.assertFalse(is_mfa_enabled(self.user))


# ─── Backup Codes Tests ─────────────────────────────────────────────────────


class BackupCodeTests(TestCase):
    """Test backup code generation and verification."""

    def setUp(self):
        self.user = register_user("backup@example.com", "StrongPass123!", org_name="Test Org", business_type="SaaS")

    def test_generate_codes_returns_10(self):
        """Generates 10 backup codes."""
        codes = generate_backup_codes(self.user)
        self.assertEqual(len(codes), 10)

    def test_verify_backup_code_success(self):
        """Valid backup code is accepted and consumed."""
        codes = generate_backup_codes(self.user)
        self.assertTrue(verify_backup_code(self.user, codes[0]))
        # Same code should not work again
        self.assertFalse(verify_backup_code(self.user, codes[0]))

    def test_invalid_backup_code_rejected(self):
        """Invalid backup code is rejected."""
        generate_backup_codes(self.user)
        self.assertFalse(verify_backup_code(self.user, "XXXX-YYYY"))

    def test_remaining_count_decreases(self):
        """Remaining count goes down after using a code."""
        codes = generate_backup_codes(self.user)
        initial = get_remaining_backup_codes_count(self.user)
        verify_backup_code(self.user, codes[0])
        self.assertEqual(get_remaining_backup_codes_count(self.user), initial - 1)

    def test_regenerate_invalidates_old(self):
        """Regenerating creates new codes and invalidates old ones."""
        old_codes = generate_backup_codes(self.user)
        new_codes = generate_backup_codes(self.user)
        self.assertFalse(verify_backup_code(self.user, old_codes[0]))
        self.assertTrue(verify_backup_code(self.user, new_codes[0]))


# ─── Session Tests ───────────────────────────────────────────────────────────


class SessionTests(TestCase):
    """Test session management."""

    def setUp(self):
        self.user = register_user("session@example.com", "StrongPass123!", org_name="Test Org", business_type="SaaS")

    def test_create_session(self):
        """Creating a session returns tokens."""
        result = create_session(
            self.user,
            ip_address="127.0.0.1",
            user_agent="TestBrowser",
        )
        self.assertIn("session_id", result)
        self.assertIn("access_token", result)
        self.assertIn("refresh_token", result)
        
        # Verify access_token is a valid JWT
        payload = jwt_service.verify_access_token(result["access_token"])
        self.assertEqual(payload["sub"], str(self.user.id))

    def test_session_is_active(self):
        """New session is active."""
        result = create_session(self.user)
        session = Session.objects.get(pk=result["session_id"])
        self.assertTrue(session.is_active)

    def test_revoke_session(self):
        """Revoking a session marks it inactive."""
        result = create_session(self.user)
        revoke_session(result["session_id"], self.user)
        session = Session.objects.get(pk=result["session_id"])
        self.assertFalse(session.is_active)

    def test_global_logout(self):
        """Global logout revokes all sessions."""
        create_session(self.user)
        create_session(self.user)
        count = global_logout(self.user)
        self.assertEqual(count, 2)
        active = Session.objects.filter(user=self.user, revoked_at__isnull=True).count()
        self.assertEqual(active, 0)

    def test_security_event_logged_on_session_create(self):
        """Session creation logs a security event."""
        create_session(self.user, ip_address="10.0.0.1")
        event = SecurityEvent.objects.filter(
            user=self.user,
            event_type=SecurityEvent.EventType.LOGIN_SUCCESS,
        ).first()
        self.assertIsNotNone(event)
        self.assertEqual(event.metadata["ip_address"], "10.0.0.1")


# ─── Rate Limit Tests ───────────────────────────────────────────────────────


class RateLimitTests(TestCase):
    """Test rate limiting logic."""

    def test_no_block_under_threshold(self):
        """No error when under the limit."""
        check_rate_limit("test:key:1")  # Should not raise

    def test_block_after_threshold(self):
        """Blocked after exceeding max attempts."""
        key = "test:key:2"
        for _ in range(5):
            record_attempt(key, max_attempts=5)
        with self.assertRaises(RateLimitError):
            check_rate_limit(key)

    def test_success_resets_counter(self):
        """Successful action resets the counter."""
        key = "test:key:3"
        for _ in range(3):
            record_attempt(key, max_attempts=5)
        record_success(key)
        check_rate_limit(key)  # Should not raise


# ─── Trusted Device Tests ───────────────────────────────────────────────────


class TrustedDeviceTests(TestCase):
    """Test trusted device management."""

    def setUp(self):
        self.user = register_user("device@example.com", "StrongPass123!", org_name="Test Org", business_type="SaaS")

    def test_register_and_check_device(self):
        """Registering a device makes it trusted."""
        register_trusted_device(self.user, "device-abc-123")
        self.assertTrue(is_trusted_device(self.user, "device-abc-123"))

    def test_unknown_device_not_trusted(self):
        """Unknown device is not trusted."""
        self.assertFalse(is_trusted_device(self.user, "unknown-device"))

    def test_revoke_all_devices(self):
        """Revoking all devices clears the list."""
        register_trusted_device(self.user, "dev-1")
        register_trusted_device(self.user, "dev-2")
        count = revoke_all_trusted_devices(self.user)
        self.assertEqual(count, 2)
        self.assertFalse(is_trusted_device(self.user, "dev-1"))


# ─── URL Tests ───────────────────────────────────────────────────────────────


class SecurityURLTests(TestCase):
    """Test URL routing."""

    def test_dashboard_url(self):
        url = reverse("security:dashboard")
        self.assertEqual(url, "/security/")

    def test_mfa_setup_url(self):
        url = reverse("security:mfa-setup")
        self.assertEqual(url, "/security/mfa/setup/")

    def test_sessions_url(self):
        url = reverse("security:sessions")
        self.assertEqual(url, "/security/sessions/")


# ─── JWT Tests ───────────────────────────────────────────────────────────────


class JWTServiceTests(TestCase):
    """Test JWT service directly."""

    def setUp(self):
        self.user = User.objects.create_user(
            email="jwt@example.com", 
            password="password123"
        )

    def test_create_access_token(self):
        token = jwt_service.create_access_token(self.user)
        self.assertIsInstance(token, str)
        
        payload = jwt_service.verify_access_token(token)
        self.assertEqual(payload["sub"], str(self.user.id))
        self.assertIn("exp", payload)

    def test_verify_invalid_token(self):
        with self.assertRaises(AuthenticationError):
            jwt_service.verify_access_token("invalid.token.here")

    def test_verify_expired_token(self):
        with override_settings(JWT_ACCESS_EXPIRY_MINUTES=-1):
            token = jwt_service.create_access_token(self.user)
            with self.assertRaises(AuthenticationError) as cm:
                jwt_service.verify_access_token(token)
            self.assertEqual(str(cm.exception), "Token has expired.")

    def test_verify_tampered_token(self):
        token = jwt_service.create_access_token(self.user)
        tampered_token = token[:-1] + ("0" if token[-1] != "0" else "1")
        with self.assertRaises(AuthenticationError):
            jwt_service.verify_access_token(tampered_token)