Spaces:

techprotrade
/

annator-command-center

Paused

annator-command-center / tests /security /test_jwt_security.py

Deploy ATOM FastAPI command center runtime (part 9)

81e3673 verified 3 months ago

13.3 kB

"""
JWT token security tests (SECU-05).
Tests cover:
- JWT validation
- Token expiration
- Token refresh flow
- Signature verification
- Payload validation
"""
import pytest
import jwt
from datetime import datetime, timedelta
from freezegun import freeze_time
from fastapi.testclient import TestClient
from jose import JWTError
from sqlalchemy.orm import Session
from core.auth import SECRET_KEY, ALGORITHM, create_access_token, decode_token, ACCESS_TOKEN_EXPIRE_MINUTES


class TestJWTValidation:
    """Test JWT token validation."""

    def test_valid_token_accepted(self, client: TestClient, valid_auth_token):
        """Test valid JWT token is accepted."""
        response = client.get("/api/auth/me",
            headers={"Authorization": f"Bearer {valid_auth_token}"}
        )

        # Should succeed or fail gracefully depending on user existence
        assert response.status_code in [200, 401]

    def test_expired_token_rejected(self, client: TestClient, expired_auth_token):
        """Test expired JWT token is rejected."""
        with freeze_time("2026-02-01 12:00:00"):  # After expiration
            response = client.get("/api/auth/me",
                headers={"Authorization": f"Bearer {expired_auth_token}"}
            )

            assert response.status_code == 401

    def test_invalid_token_rejected(self, client: TestClient, invalid_auth_token):
        """Test invalid JWT token format is rejected."""
        response = client.get("/api/auth/me",
            headers={"Authorization": f"Bearer {invalid_auth_token}"}
        )

        assert response.status_code == 401

    def test_tampered_token_rejected(self, client: TestClient, tampered_token):
        """Test tampered JWT token is rejected (signature validation)."""
        response = client.get("/api/auth/me",
            headers={"Authorization": f"Bearer {tampered_token}"}
        )

        assert response.status_code == 401

    def test_missing_token_rejected(self, client: TestClient):
        """Test requests without token are rejected."""
        response = client.get("/api/auth/me")

        assert response.status_code == 401

    def test_token_with_invalid_signature_rejected(self, client: TestClient):
        """Test token signed with wrong secret is rejected."""
        import jose

        # Create token with wrong secret
        payload = {"sub": "user_123", "exp": datetime.utcnow() + timedelta(hours=1)}
        token = jose.jwt.encode(payload, "wrong_secret", algorithm="HS256")

        response = client.get("/api/auth/me",
            headers={"Authorization": f"Bearer {token}"}
        )

        assert response.status_code == 401

    def test_token_without_bearer_prefix_rejected(self, client: TestClient, valid_auth_token):
        """Test token without 'Bearer' prefix is rejected."""
        response = client.get("/api/auth/me",
            headers={"Authorization": valid_auth_token}
        )

        # May or may not be rejected depending on implementation
        # Some implementations are lenient
        assert response.status_code in [200, 401]


class TestTokenExpiration:
    """Test token expiration logic."""

    def test_token_expires_after_configured_time(self, client: TestClient):
        """Test token expires after configured time limit."""
        # Create token with 1 minute expiration
        token = create_access_token(
            data={"sub": "test_user"},
            expires_delta=timedelta(minutes=1)
        )

        # Token should work immediately
        with freeze_time("2026-02-01 10:00:00"):
            response = client.get("/api/auth/me",
                headers={"Authorization": f"Bearer {token}"}
            )
            # May fail if test_user doesn't exist, but should not fail due to expiration
            assert response.status_code in [200, 401]

        # Token should fail after expiration
        with freeze_time("2026-02-01 10:02:00"):  # 2 minutes later
            response = client.get("/api/auth/me",
                headers={"Authorization": f"Bearer {token}"}
            )
            assert response.status_code == 401

    def test_default_expiration_time(self):
        """Test default token expiration is 24 hours."""
        # Default should be 24 hours (1440 minutes)
        assert ACCESS_TOKEN_EXPIRE_MINUTES >= 60 * 23

    def test_custom_expiration_time(self):
        """Test token with custom expiration time."""
        custom_expire = timedelta(hours=2)
        token = create_access_token(
            data={"sub": "test_user"},
            expires_delta=custom_expire
        )

        # Decode and check expiration
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        exp_timestamp = payload["exp"]
        exp_datetime = datetime.utcfromtimestamp(exp_timestamp)

        # Should be approximately 2 hours from now
        # We can't check exact time due to execution time, but can check structure
        assert exp_datetime is not None

    def test_expired_token_raises_error(self):
        """Test decoding expired token raises JWTError."""
        # Create expired token
        with freeze_time("2026-02-01 10:00:00"):
            token = create_access_token(
                data={"sub": "test_user"},
                expires_delta=timedelta(minutes=1)
            )

        # Try to decode after expiration
        with freeze_time("2026-02-01 10:05:00"):
            with pytest.raises(JWTError):
                jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])


class TestTokenPayload:
    """Test JWT token payload structure."""

    def test_token_contains_user_id(self, valid_auth_token):
        """Test token contains user_id in payload."""
        from jose import jwt

        payload = jwt.decode(
            valid_auth_token,
            SECRET_KEY,
            algorithms=[ALGORITHM]
        )

        assert "sub" in payload  # User ID standard claim
        assert payload["sub"] is not None

    def test_token_contains_expiration(self, valid_auth_token):
        """Test token contains expiration claim."""
        from jose import jwt

        payload = jwt.decode(
            valid_auth_token,
            SECRET_KEY,
            algorithms=[ALGORITHM]
        )

        assert "exp" in payload  # Expiration standard claim
        assert payload["exp"] > 0

    def test_decode_token_function(self):
        """Test decode_token helper function."""
        # Create test token
        token = create_access_token(data={"sub": "user_456"})

        # Decode should return payload
        payload = decode_token(token)
        assert payload is not None
        assert payload["sub"] == "user_456"

        # Invalid token should return None
        assert decode_token("invalid") is None

    def test_token_payload_is_immutable(self, valid_auth_token):
        """Test that token payload cannot be modified without invalidating signature."""
        from jose import jwt

        # Decode original token
        original_payload = jwt.decode(
            valid_auth_token,
            SECRET_KEY,
            algorithms=[ALGORITHM]
        )

        # Try to create a new token with modified payload
        modified_payload = original_payload.copy()
        modified_payload["admin"] = True  # Add privilege escalation

        # Encode with wrong secret
        tampered_token = jwt.encode(modified_payload, "wrong_secret", algorithm=ALGORITHM)

        # Should fail to verify
        with pytest.raises(JWTError):
            jwt.decode(tampered_token, SECRET_KEY, algorithms=[ALGORITHM])


class TestTokenRefresh:
    """Test token refresh flow."""

    def test_refresh_token_endpoint_exists(self, client: TestClient):
        """Test refresh token endpoint is available."""
        # This test verifies the endpoint exists
        # Implementation may vary based on actual refresh flow
        response = client.post("/api/auth/refresh", json={
            "refresh_token": "test_refresh_token"
        })

        # Should respond (either success or failure, but endpoint exists)
        assert response.status_code in [200, 401, 400]

    def test_can_refresh_expired_token(self, client: TestClient, refresh_token):
        """Test can refresh token with valid refresh token."""
        # Implementation depends on refresh flow
        # Should test: expired access token + valid refresh token = new access token
        response = client.post("/api/auth/refresh", json={
            "refresh_token": refresh_token
        })

        # Should either succeed (200) or fail with specific error
        assert response.status_code in [200, 401, 400]

    def test_refresh_without_refresh_token_fails(self, client: TestClient):
        """Test refresh fails without valid refresh token."""
        response = client.post("/api/auth/refresh", json={
            "refresh_token": "invalid_refresh_token"
        })

        assert response.status_code in [401, 400]

    def test_refresh_with_empty_token_fails(self, client: TestClient):
        """Test refresh fails with empty refresh token."""
        response = client.post("/api/auth/refresh", json={
            "refresh_token": ""
        })

        assert response.status_code in [400, 422, 401]


class TestTokenSecurity:
    """Test JWT security properties."""

    def test_algorithm_is_hs256(self):
        """Test JWT uses HS256 algorithm (not 'none')."""
        from jose import jwt

        token = create_access_token(data={"sub": "test_user"})
        header = jwt.get_unverified_header(token)

        assert header["alg"] == "HS256"
        # Ensure 'none' algorithm is not used (security vulnerability)

    def test_token_uses_secret_key(self):
        """Test token is signed with configured secret."""
        from jose import jwt

        token = create_access_token(data={"sub": "test_user"})

        # Should decode with correct secret
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        assert payload["sub"] == "test_user"

        # Should fail with wrong secret
        with pytest.raises(JWTError):
            jwt.decode(token, "wrong_secret", algorithms=[ALGORITHM])

    def test_none_algorithm_prevented(self):
        """Test that 'none' algorithm cannot be used (security vulnerability)."""
        from jose import jwt

        # Try to create a token with 'none' algorithm
        payload = {"sub": "test_user", "exp": datetime.utcnow() + timedelta(hours=1)}

        # The library should prevent this, or validation should reject it
        # Our implementation always uses HS256
        token = create_access_token(data={"sub": "test_user"})

        # Verify it's not using 'none'
        header = jwt.get_unverified_header(token)
        assert header["alg"] != "none"

    def test_token_structure_is_valid(self):
        """Test token has valid JWT structure (header.payload.signature)."""
        token = create_access_token(data={"sub": "test_user"})

        parts = token.split(".")
        assert len(parts) == 3

        # Each part should be base64url encoded
        # (non-empty and valid length)
        assert len(parts[0]) > 0  # header
        assert len(parts[1]) > 0  # payload
        assert len(parts[2]) > 0  # signature

    def test_token_with_future_expiration(self):
        """Test token with future expiration is valid."""
        future_expire = timedelta(days=365)
        token = create_access_token(
            data={"sub": "test_user"},
            expires_delta=future_expire
        )

        # Should be valid now
        payload = decode_token(token)
        assert payload is not None
        assert payload["sub"] == "test_user"


class TestTokenEdgeCases:
    """Test edge cases in token handling."""

    def test_token_with_empty_payload(self):
        """Test creating token with minimal payload."""
        token = create_access_token(data={})

        # Should still create valid token structure
        parts = token.split(".")
        assert len(parts) == 3

    def test_token_very_long_expiration(self):
        """Test token with very long expiration."""
        # 10 years
        long_expire = timedelta(days=365 * 10)
        token = create_access_token(
            data={"sub": "test_user"},
            expires_delta=long_expire
        )

        # Should create valid token
        payload = decode_token(token)
        assert payload is not None

    def test_multiple_tokens_for_same_user(self):
        """Test creating multiple tokens for same user."""
        user_id = "test_user_123"

        token1 = create_access_token(data={"sub": user_id})
        token2 = create_access_token(data={"sub": user_id})

        # Tokens should be different (different iat/exp)
        assert token1 != token2

        # But both should be valid
        assert decode_token(token1) is not None
        assert decode_token(token2) is not None

    def test_token_with_special_characters_in_subject(self):
        """Test token with special characters in subject."""
        special_subject = "user@example.com"

        token = create_access_token(data={"sub": special_subject})
        payload = decode_token(token)

        assert payload is not None
        assert payload["sub"] == special_subject

	"""
	JWT token security tests (SECU-05).

	Tests cover:
	- JWT validation
	- Token expiration
	- Token refresh flow
	- Signature verification
	- Payload validation
	"""
	import pytest
	import jwt
	from datetime import datetime, timedelta
	from freezegun import freeze_time
	from fastapi.testclient import TestClient
	from jose import JWTError
	from sqlalchemy.orm import Session
	from core.auth import SECRET_KEY, ALGORITHM, create_access_token, decode_token, ACCESS_TOKEN_EXPIRE_MINUTES


	class TestJWTValidation:
	"""Test JWT token validation."""

	def test_valid_token_accepted(self, client: TestClient, valid_auth_token):
	"""Test valid JWT token is accepted."""
	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {valid_auth_token}"}
	)

	# Should succeed or fail gracefully depending on user existence
	assert response.status_code in [200, 401]

	def test_expired_token_rejected(self, client: TestClient, expired_auth_token):
	"""Test expired JWT token is rejected."""
	with freeze_time("2026-02-01 12:00:00"): # After expiration
	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {expired_auth_token}"}
	)

	assert response.status_code == 401

	def test_invalid_token_rejected(self, client: TestClient, invalid_auth_token):
	"""Test invalid JWT token format is rejected."""
	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {invalid_auth_token}"}
	)

	assert response.status_code == 401

	def test_tampered_token_rejected(self, client: TestClient, tampered_token):
	"""Test tampered JWT token is rejected (signature validation)."""
	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {tampered_token}"}
	)

	assert response.status_code == 401

	def test_missing_token_rejected(self, client: TestClient):
	"""Test requests without token are rejected."""
	response = client.get("/api/auth/me")

	assert response.status_code == 401

	def test_token_with_invalid_signature_rejected(self, client: TestClient):
	"""Test token signed with wrong secret is rejected."""
	import jose

	# Create token with wrong secret
	payload = {"sub": "user_123", "exp": datetime.utcnow() + timedelta(hours=1)}
	token = jose.jwt.encode(payload, "wrong_secret", algorithm="HS256")

	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {token}"}
	)

	assert response.status_code == 401

	def test_token_without_bearer_prefix_rejected(self, client: TestClient, valid_auth_token):
	"""Test token without 'Bearer' prefix is rejected."""
	response = client.get("/api/auth/me",
	headers={"Authorization": valid_auth_token}
	)

	# May or may not be rejected depending on implementation
	# Some implementations are lenient
	assert response.status_code in [200, 401]


	class TestTokenExpiration:
	"""Test token expiration logic."""

	def test_token_expires_after_configured_time(self, client: TestClient):
	"""Test token expires after configured time limit."""
	# Create token with 1 minute expiration
	token = create_access_token(
	data={"sub": "test_user"},
	expires_delta=timedelta(minutes=1)
	)

	# Token should work immediately
	with freeze_time("2026-02-01 10:00:00"):
	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {token}"}
	)
	# May fail if test_user doesn't exist, but should not fail due to expiration
	assert response.status_code in [200, 401]

	# Token should fail after expiration
	with freeze_time("2026-02-01 10:02:00"): # 2 minutes later
	response = client.get("/api/auth/me",
	headers={"Authorization": f"Bearer {token}"}
	)
	assert response.status_code == 401

	def test_default_expiration_time(self):
	"""Test default token expiration is 24 hours."""
	# Default should be 24 hours (1440 minutes)
	assert ACCESS_TOKEN_EXPIRE_MINUTES >= 60 * 23

	def test_custom_expiration_time(self):
	"""Test token with custom expiration time."""
	custom_expire = timedelta(hours=2)
	token = create_access_token(
	data={"sub": "test_user"},
	expires_delta=custom_expire
	)

	# Decode and check expiration
	payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
	exp_timestamp = payload["exp"]
	exp_datetime = datetime.utcfromtimestamp(exp_timestamp)

	# Should be approximately 2 hours from now
	# We can't check exact time due to execution time, but can check structure
	assert exp_datetime is not None

	def test_expired_token_raises_error(self):
	"""Test decoding expired token raises JWTError."""
	# Create expired token
	with freeze_time("2026-02-01 10:00:00"):
	token = create_access_token(
	data={"sub": "test_user"},
	expires_delta=timedelta(minutes=1)
	)

	# Try to decode after expiration
	with freeze_time("2026-02-01 10:05:00"):
	with pytest.raises(JWTError):
	jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])


	class TestTokenPayload:
	"""Test JWT token payload structure."""

	def test_token_contains_user_id(self, valid_auth_token):
	"""Test token contains user_id in payload."""
	from jose import jwt

	payload = jwt.decode(
	valid_auth_token,
	SECRET_KEY,
	algorithms=[ALGORITHM]
	)

	assert "sub" in payload # User ID standard claim
	assert payload["sub"] is not None

	def test_token_contains_expiration(self, valid_auth_token):
	"""Test token contains expiration claim."""
	from jose import jwt

	payload = jwt.decode(
	valid_auth_token,
	SECRET_KEY,
	algorithms=[ALGORITHM]
	)

	assert "exp" in payload # Expiration standard claim
	assert payload["exp"] > 0

	def test_decode_token_function(self):
	"""Test decode_token helper function."""
	# Create test token
	token = create_access_token(data={"sub": "user_456"})

	# Decode should return payload
	payload = decode_token(token)
	assert payload is not None
	assert payload["sub"] == "user_456"

	# Invalid token should return None
	assert decode_token("invalid") is None

	def test_token_payload_is_immutable(self, valid_auth_token):
	"""Test that token payload cannot be modified without invalidating signature."""
	from jose import jwt

	# Decode original token
	original_payload = jwt.decode(
	valid_auth_token,
	SECRET_KEY,
	algorithms=[ALGORITHM]
	)

	# Try to create a new token with modified payload
	modified_payload = original_payload.copy()
	modified_payload["admin"] = True # Add privilege escalation

	# Encode with wrong secret
	tampered_token = jwt.encode(modified_payload, "wrong_secret", algorithm=ALGORITHM)

	# Should fail to verify
	with pytest.raises(JWTError):
	jwt.decode(tampered_token, SECRET_KEY, algorithms=[ALGORITHM])


	class TestTokenRefresh:
	"""Test token refresh flow."""

	def test_refresh_token_endpoint_exists(self, client: TestClient):
	"""Test refresh token endpoint is available."""
	# This test verifies the endpoint exists
	# Implementation may vary based on actual refresh flow
	response = client.post("/api/auth/refresh", json={
	"refresh_token": "test_refresh_token"
	})

	# Should respond (either success or failure, but endpoint exists)
	assert response.status_code in [200, 401, 400]

	def test_can_refresh_expired_token(self, client: TestClient, refresh_token):
	"""Test can refresh token with valid refresh token."""
	# Implementation depends on refresh flow
	# Should test: expired access token + valid refresh token = new access token
	response = client.post("/api/auth/refresh", json={
	"refresh_token": refresh_token
	})

	# Should either succeed (200) or fail with specific error
	assert response.status_code in [200, 401, 400]

	def test_refresh_without_refresh_token_fails(self, client: TestClient):
	"""Test refresh fails without valid refresh token."""
	response = client.post("/api/auth/refresh", json={
	"refresh_token": "invalid_refresh_token"
	})

	assert response.status_code in [401, 400]

	def test_refresh_with_empty_token_fails(self, client: TestClient):
	"""Test refresh fails with empty refresh token."""
	response = client.post("/api/auth/refresh", json={
	"refresh_token": ""
	})

	assert response.status_code in [400, 422, 401]


	class TestTokenSecurity:
	"""Test JWT security properties."""

	def test_algorithm_is_hs256(self):
	"""Test JWT uses HS256 algorithm (not 'none')."""
	from jose import jwt

	token = create_access_token(data={"sub": "test_user"})
	header = jwt.get_unverified_header(token)

	assert header["alg"] == "HS256"
	# Ensure 'none' algorithm is not used (security vulnerability)

	def test_token_uses_secret_key(self):
	"""Test token is signed with configured secret."""
	from jose import jwt

	token = create_access_token(data={"sub": "test_user"})

	# Should decode with correct secret
	payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
	assert payload["sub"] == "test_user"

	# Should fail with wrong secret
	with pytest.raises(JWTError):
	jwt.decode(token, "wrong_secret", algorithms=[ALGORITHM])

	def test_none_algorithm_prevented(self):
	"""Test that 'none' algorithm cannot be used (security vulnerability)."""
	from jose import jwt

	# Try to create a token with 'none' algorithm
	payload = {"sub": "test_user", "exp": datetime.utcnow() + timedelta(hours=1)}

	# The library should prevent this, or validation should reject it
	# Our implementation always uses HS256
	token = create_access_token(data={"sub": "test_user"})

	# Verify it's not using 'none'
	header = jwt.get_unverified_header(token)
	assert header["alg"] != "none"

	def test_token_structure_is_valid(self):
	"""Test token has valid JWT structure (header.payload.signature)."""
	token = create_access_token(data={"sub": "test_user"})

	parts = token.split(".")
	assert len(parts) == 3

	# Each part should be base64url encoded
	# (non-empty and valid length)
	assert len(parts[0]) > 0 # header
	assert len(parts[1]) > 0 # payload
	assert len(parts[2]) > 0 # signature

	def test_token_with_future_expiration(self):
	"""Test token with future expiration is valid."""
	future_expire = timedelta(days=365)
	token = create_access_token(
	data={"sub": "test_user"},
	expires_delta=future_expire
	)

	# Should be valid now
	payload = decode_token(token)
	assert payload is not None
	assert payload["sub"] == "test_user"


	class TestTokenEdgeCases:
	"""Test edge cases in token handling."""

	def test_token_with_empty_payload(self):
	"""Test creating token with minimal payload."""
	token = create_access_token(data={})

	# Should still create valid token structure
	parts = token.split(".")
	assert len(parts) == 3

	def test_token_very_long_expiration(self):
	"""Test token with very long expiration."""
	# 10 years
	long_expire = timedelta(days=365 * 10)
	token = create_access_token(
	data={"sub": "test_user"},
	expires_delta=long_expire
	)

	# Should create valid token
	payload = decode_token(token)
	assert payload is not None

	def test_multiple_tokens_for_same_user(self):
	"""Test creating multiple tokens for same user."""
	user_id = "test_user_123"

	token1 = create_access_token(data={"sub": user_id})
	token2 = create_access_token(data={"sub": user_id})

	# Tokens should be different (different iat/exp)
	assert token1 != token2

	# But both should be valid
	assert decode_token(token1) is not None
	assert decode_token(token2) is not None

	def test_token_with_special_characters_in_subject(self):
	"""Test token with special characters in subject."""
	special_subject = "user@example.com"

	token = create_access_token(data={"sub": special_subject})
	payload = decode_token(token)

	assert payload is not None
	assert payload["sub"] == special_subject