invoke-ai--invokeai/tests/app/services/auth/test_token_service.py

"""Unit tests for JWT token service."""

import time
from datetime import timedelta

import pytest

from invokeai.app.services.auth.token_service import TokenData, create_access_token, set_jwt_secret, verify_token


@pytest.fixture(scope="module", autouse=True)
def setup_jwt_secret():
    """Set up JWT secret for all tests in this module."""
    # Use a test secret key
    set_jwt_secret("test-secret-key-for-unit-tests-only-do-not-use-in-production")


# Minimum token length to safely modify middle characters for testing
# JWT tokens have format header.payload.signature and are typically >180 characters
MIN_TOKEN_LENGTH_FOR_MODIFICATION = 50

# Minimum signature length to safely modify middle characters for testing
# JWT signatures are typically 43 characters (base64-encoded HMAC-SHA256)
MIN_SIGNATURE_LENGTH_FOR_MODIFICATION = 10


class TestTokenCreation:
    """Tests for JWT token creation."""

    def test_create_access_token_basic(self):
        """Test creating a basic access token."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)

        assert token is not None
        assert isinstance(token, str)
        assert len(token) > 0

    def test_create_access_token_with_expiration(self):
        """Test creating token with custom expiration."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data, expires_delta=timedelta(hours=1))

        assert token is not None
        # Verify token is valid
        verified_data = verify_token(token)
        assert verified_data is not None
        assert verified_data.user_id == "user123"

    def test_create_access_token_admin_user(self):
        """Test creating token for admin user."""
        token_data = TokenData(
            user_id="admin123",
            email="admin@example.com",
            is_admin=True,
        )

        token = create_access_token(token_data)
        verified_data = verify_token(token)

        assert verified_data is not None
        assert verified_data.is_admin is True

    def test_create_access_token_preserves_all_data(self):
        """Test that all token data is preserved."""
        token_data = TokenData(
            user_id="user_with_complex_id_12345",
            email="complex.email+tag@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)
        verified_data = verify_token(token)

        assert verified_data is not None
        assert verified_data.user_id == token_data.user_id
        assert verified_data.email == token_data.email
        assert verified_data.is_admin == token_data.is_admin

    def test_create_access_token_different_each_time(self):
        """Test that creating token with same data produces different tokens (due to timestamps)."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        # Create tokens with different expiration times to ensure uniqueness
        token1 = create_access_token(token_data, expires_delta=timedelta(hours=1))
        token2 = create_access_token(token_data, expires_delta=timedelta(hours=2))

        # Tokens should be different due to different exp timestamps
        assert token1 != token2


class TestTokenVerification:
    """Tests for JWT token verification."""

    def test_verify_valid_token(self):
        """Test verifying a valid token."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)
        verified_data = verify_token(token)

        assert verified_data is not None
        assert verified_data.user_id == "user123"
        assert verified_data.email == "test@example.com"
        assert verified_data.is_admin is False

    def test_verify_invalid_token(self):
        """Test verifying an invalid token."""
        verified_data = verify_token("invalid_token_string")

        assert verified_data is None

    def test_verify_malformed_token(self):
        """Test verifying malformed tokens."""
        malformed_tokens = [
            "",
            "not.a.token",
            "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.invalid",
            "header.payload",  # Missing signature
        ]

        for token in malformed_tokens:
            verified_data = verify_token(token)
            assert verified_data is None, f"Should reject malformed token: {token}"

    def test_verify_expired_token(self):
        """Test verifying an expired token."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        # Create token that expires in 100 milliseconds (0.1 seconds)
        token = create_access_token(token_data, expires_delta=timedelta(milliseconds=100))

        # Wait for token to expire (wait longer than expiration - 200ms to be safe)
        time.sleep(0.2)

        # Token should be invalid now
        verified_data = verify_token(token)
        assert verified_data is None

    def test_verify_token_with_modified_payload(self):
        """Test that tokens with modified payloads are rejected."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)

        # Try to modify the token by changing a character in the middle
        # JWT tokens are base64 encoded, so changing any character should invalidate the signature
        # Note: We change a character in the middle to avoid Base64 padding issues where
        # the last character might not affect the decoded value
        if len(token) > MIN_TOKEN_LENGTH_FOR_MODIFICATION:
            mid = len(token) // 2
            modified_token = token[:mid] + ("X" if token[mid] != "X" else "Y") + token[mid + 1 :]
            verified_data = verify_token(modified_token)
            assert verified_data is None

    def test_verify_token_preserves_admin_status(self):
        """Test that admin status is correctly preserved through token lifecycle."""
        # Test with regular user
        token_data = TokenData(
            user_id="user123",
            email="user@example.com",
            is_admin=False,
        )
        token = create_access_token(token_data)
        verified = verify_token(token)
        assert verified is not None
        assert verified.is_admin is False

        # Test with admin user
        admin_token_data = TokenData(
            user_id="admin123",
            email="admin@example.com",
            is_admin=True,
        )
        admin_token = create_access_token(admin_token_data)
        admin_verified = verify_token(admin_token)
        assert admin_verified is not None
        assert admin_verified.is_admin is True


class TestTokenExpiration:
    """Tests for token expiration handling."""

    def test_token_not_expired_immediately(self):
        """Test that freshly created token is not expired."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data, expires_delta=timedelta(hours=1))
        verified_data = verify_token(token)

        assert verified_data is not None

    def test_token_with_long_expiration(self):
        """Test token with long expiration time."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        # Create token that expires in 7 days
        token = create_access_token(token_data, expires_delta=timedelta(days=7))
        verified_data = verify_token(token)

        assert verified_data is not None
        assert verified_data.user_id == "user123"

    def test_token_with_short_expiration_not_expired(self):
        """Test token with short but not yet expired time."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        # Create token that expires in 1 second
        token = create_access_token(token_data, expires_delta=timedelta(seconds=1))

        # Immediately verify - should still be valid
        verified_data = verify_token(token)
        assert verified_data is not None


class TestTokenDataModel:
    """Tests for TokenData model."""

    def test_token_data_creation(self):
        """Test creating TokenData instance."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        assert token_data.user_id == "user123"
        assert token_data.email == "test@example.com"
        assert token_data.is_admin is False

    def test_token_data_with_admin(self):
        """Test TokenData for admin user."""
        token_data = TokenData(
            user_id="admin123",
            email="admin@example.com",
            is_admin=True,
        )

        assert token_data.is_admin is True

    def test_token_data_model_dump(self):
        """Test that TokenData can be serialized."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        data_dict = token_data.model_dump()

        assert isinstance(data_dict, dict)
        assert data_dict["user_id"] == "user123"
        assert data_dict["email"] == "test@example.com"
        assert data_dict["is_admin"] is False


class TestTokenSecurity:
    """Tests for token security properties."""

    def test_token_signature_verification(self):
        """Test that token signature is verified."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)

        # Token should verify correctly
        assert verify_token(token) is not None

        # Modified token should fail verification
        if len(token) > MIN_TOKEN_LENGTH_FOR_MODIFICATION:
            # Change a character in the signature part (last part of JWT)
            parts = token.split(".")
            if len(parts) == 3 and len(parts[2]) > MIN_SIGNATURE_LENGTH_FOR_MODIFICATION:
                # Modify a character in the middle of the signature to avoid Base64 padding issues
                # where the last few characters might not affect the decoded value
                mid = len(parts[2]) // 2
                modified_signature = parts[2][:mid] + ("X" if parts[2][mid] != "X" else "Y") + parts[2][mid + 1 :]
                modified_token = f"{parts[0]}.{parts[1]}.{modified_signature}"
                assert verify_token(modified_token) is None

    def test_cannot_forge_admin_token(self):
        """Test that admin status cannot be forged by modifying token."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)

        # Any modification to the token should invalidate it
        # This prevents attackers from changing is_admin=false to is_admin=true
        parts = token.split(".")
        if len(parts) == 3:
            # Try to modify the payload
            modified_payload = parts[1][:-1] + ("X" if parts[1][-1] != "X" else "Y")
            modified_token = f"{parts[0]}.{modified_payload}.{parts[2]}"

            verified_data = verify_token(modified_token)
            # Modified token should be rejected
            assert verified_data is None

    def test_token_uses_strong_algorithm(self):
        """Test that token uses secure algorithm (HS256)."""
        token_data = TokenData(
            user_id="user123",
            email="test@example.com",
            is_admin=False,
        )

        token = create_access_token(token_data)

        # JWT tokens have format: header.payload.signature
        # Header contains algorithm information
        import base64
        import json

        parts = token.split(".")
        if len(parts) >= 1:
            # Decode header (add padding if necessary)
            header_b64 = parts[0]
            # Add padding if necessary
            padding = 4 - len(header_b64) % 4
            if padding != 4:
                header_b64 += "=" * padding

            header = json.loads(base64.urlsafe_b64decode(header_b64))
            # Should use HS256 algorithm
            assert header.get("alg") == "HS256"