demo-oidc/tests/test_token_validation.py

import pytest
import requests
from unittest.mock import patch, MagicMock
import jwt
from main import (
    verify_token, 
    refresh_token, 
    get_well_known_config, 
    get_jwks,
    get_user_info_from_endpoint,
    get_userinfo_endpoint,
    get_token_endpoint
)

def test_verify_token_invalid_input():
    """Test verify_token with invalid inputs."""
    # Test with None
    with pytest.raises(ValueError, match="Invalid token provided"):
        verify_token(None)
    
    # Test with empty string
    with pytest.raises(ValueError, match="Invalid token provided"):
        verify_token("")
    
    # Test with non-string
    with pytest.raises(ValueError, match="Invalid token provided"):
        verify_token(123)

@patch('main.get_jwks')
@patch('main.jwt.get_unverified_header')
def test_verify_token_missing_kid(mock_get_header, mock_get_jwks):
    """Test verify_token when token header is missing 'kid'."""
    # Mock the token header to not include 'kid'
    mock_get_header.return_value = {}
    
    with pytest.raises(ValueError, match="Token header missing 'kid' field"):
        verify_token("some.token.string")

@patch('main.get_jwks')
@patch('main.jwt.get_unverified_header')
def test_verify_token_no_matching_key(mock_get_header, mock_get_jwks):
    """Test verify_token when no matching key is found in JWKS."""
    # Mock the token header to include 'kid'
    mock_get_header.return_value = {"kid": "nonexistent-key"}
    
    # Mock JWKS with a different key ID
    mock_get_jwks.return_value = {
        "keys": [
            {
                "kid": "different-key",
                "kty": "RSA",
                "n": "modulus",
                "e": "AQAB"
            }
        ]
    }
    
    with pytest.raises(ValueError, match="Unable to find matching key in JWKS"):
        verify_token("some.token.string")


def _mock_jwks_and_key():
    """Helper: return (jwks, header) fixtures for reaching the jwt.decode step."""
    jwks = {
        "keys": [
            {"kid": "test-key-id", "kty": "RSA", "n": "modulus", "e": "AQAB"}
        ]
    }
    header = {"kid": "test-key-id", "alg": "RS256"}
    return jwks, header


@patch('main.jwt.decode')
@patch('main.jwt.algorithms.RSAAlgorithm.from_jwk')
@patch('main.jwt.get_unverified_header')
@patch('main.get_jwks')
def test_verify_token_nonce_match(mock_get_jwks, mock_get_header, mock_from_jwk, mock_decode):
    """verify_token returns payload when expected_nonce matches the token nonce."""
    jwks, header = _mock_jwks_and_key()
    mock_get_jwks.return_value = jwks
    mock_get_header.return_value = header
    mock_from_jwk.return_value = "mock-key"
    mock_decode.return_value = {
        "iss": "https://login.infomaniak.com",
        "aud": "test-client-id",
        "email": "rene.luria@infomaniak.com",
        "nonce": "abc123",
    }

    payload = verify_token("some.token.string", expected_nonce="abc123")
    assert payload["nonce"] == "abc123"


@patch('main.jwt.decode')
@patch('main.jwt.algorithms.RSAAlgorithm.from_jwk')
@patch('main.jwt.get_unverified_header')
@patch('main.get_jwks')
def test_verify_token_nonce_mismatch(mock_get_jwks, mock_get_header, mock_from_jwk, mock_decode):
    """verify_token rejects an id_token whose nonce does not match the expected one."""
    jwks, header = _mock_jwks_and_key()
    mock_get_jwks.return_value = jwks
    mock_get_header.return_value = header
    mock_from_jwk.return_value = "mock-key"
    mock_decode.return_value = {
        "iss": "https://login.infomaniak.com",
        "aud": "test-client-id",
        "email": "attacker@evil.com",
        "nonce": "VERYWRONG_NONCE",
    }

    with pytest.raises(ValueError, match="Nonce mismatch"):
        verify_token("some.token.string", expected_nonce="abc123")


@patch('main.jwt.decode')
@patch('main.jwt.algorithms.RSAAlgorithm.from_jwk')
@patch('main.jwt.get_unverified_header')
@patch('main.get_jwks')
def test_verify_token_nonce_missing_in_token(mock_get_jwks, mock_get_header, mock_from_jwk, mock_decode):
    """verify_token rejects a token lacking a nonce claim when an expected nonce is supplied."""
    jwks, header = _mock_jwks_and_key()
    mock_get_jwks.return_value = jwks
    mock_get_header.return_value = header
    mock_from_jwk.return_value = "mock-key"
    mock_decode.return_value = {
        "iss": "https://login.infomaniak.com",
        "aud": "test-client-id",
        "email": "rene.luria@infomaniak.com",
    }

    with pytest.raises(ValueError, match="Token missing nonce claim"):
        verify_token("some.token.string", expected_nonce="abc123")


@patch('main.jwt.decode')
@patch('main.jwt.algorithms.RSAAlgorithm.from_jwk')
@patch('main.jwt.get_unverified_header')
@patch('main.get_jwks')
def test_verify_token_nonce_not_checked_when_none(mock_get_jwks, mock_get_header, mock_from_jwk, mock_decode):
    """verify_token skips the nonce check when no expected_nonce is provided (returning visits / refresh)."""
    jwks, header = _mock_jwks_and_key()
    mock_get_jwks.return_value = jwks
    mock_get_header.return_value = header
    mock_from_jwk.return_value = "mock-key"
    mock_decode.return_value = {
        "iss": "https://login.infomaniak.com",
        "aud": "test-client-id",
        "email": "rene.luria@infomaniak.com",
    }

    payload = verify_token("some.token.string")
    assert payload["email"] == "rene.luria@infomaniak.com"

@patch('main.requests.get')
def test_get_well_known_config_success(mock_get):
    """Test successful retrieval of well-known configuration."""
    # Mock successful response
    mock_response = MagicMock()
    mock_response.json.return_value = {
        "issuer": "https://login.infomaniak.com",
        "jwks_uri": "https://login.infomaniak.com/.well-known/jwks.json"
    }
    mock_response.text = '{"issuer": "https://login.infomaniak.com"}'
    mock_response.raise_for_status.return_value = None
    mock_get.return_value = mock_response
    
    config = get_well_known_config()
    assert "issuer" in config
    assert config["issuer"] == "https://login.infomaniak.com"

@patch('main.requests.get')
def test_get_well_known_config_failure(mock_get):
    """Test failed retrieval of well-known configuration."""
    # Clear cache first
    get_well_known_config.cache_clear()
    
    # Mock failed response with requests exception
    mock_get.side_effect = requests.exceptions.RequestException("Network error")
    
    with pytest.raises(ValueError, match="Failed to fetch well-known configuration"):
        get_well_known_config()

@patch('main.requests.get')
def test_get_jwks_success(mock_get):
    """Test successful retrieval of JWKS."""
    # Mock successful response for well-known config
    with patch('main.get_well_known_config') as mock_config:
        mock_config.return_value = {
            "jwks_uri": "https://login.infomaniak.com/.well-known/jwks.json"
        }
        
        # Mock successful JWKS response
        mock_response = MagicMock()
        mock_response.json.return_value = {
            "keys": [{"kid": "test-key"}]
        }
        mock_response.text = '{"keys": [{"kid": "test-key"}]}'
        mock_response.raise_for_status.return_value = None
        mock_get.return_value = mock_response
        
        jwks = get_jwks()
        assert "keys" in jwks
        assert len(jwks["keys"]) > 0

@patch('main.requests.get')
def test_get_jwks_failure(mock_get):
    """Test failed retrieval of JWKS."""
    # Clear cache first
    get_jwks.cache_clear()
    
    # Mock successful response for well-known config
    with patch('main.get_well_known_config') as mock_config:
        mock_config.return_value = {
            "jwks_uri": "https://login.infomaniak.com/.well-known/jwks.json"
        }
        
        # Mock failed response with requests exception
        mock_get.side_effect = requests.exceptions.RequestException("Network error")
        
        with pytest.raises(ValueError, match="Failed to fetch JWKS"):
            get_jwks()

@patch('main.requests.post')
def test_refresh_token_success(mock_post):
    """Test successful token refresh."""
    # Mock environment variables
    with patch.dict('main.os.environ', {'CLIENT_ID': 'test-client', 'CLIENT_SECRET': 'test-secret'}):
        # Mock successful response
        mock_response = MagicMock()
        mock_response.ok = True
        mock_response.json.return_value = {
            "access_token": "new-access-token",
            "id_token": "new-id-token",
            "expires_in": 3600,
            "token_type": "Bearer"
        }
        mock_post.return_value = mock_response
        
        with patch('main.get_token_endpoint') as mock_endpoint:
            mock_endpoint.return_value = "https://login.infomaniak.com/token"
            
            result = refresh_token("valid-refresh-token")
            assert "access_token" in result
            assert result["access_token"] == "new-access-token"

@patch('main.requests.post')
def test_refresh_token_failure(mock_post):
    """Test failed token refresh."""
    # Mock environment variables
    with patch.dict('main.os.environ', {'CLIENT_ID': 'test-client', 'CLIENT_SECRET': 'test-secret'}):
        # Mock failed response
        mock_response = MagicMock()
        mock_response.ok = False
        mock_response.status_code = 400
        mock_response.text = "Bad Request"
        mock_post.return_value = mock_response
        
        with patch('main.get_token_endpoint') as mock_endpoint:
            mock_endpoint.return_value = "https://login.infomaniak.com/token"
            
            with pytest.raises(ValueError, match="Token refresh failed"):
                refresh_token("invalid-refresh-token")

def test_refresh_token_invalid_input():
    """Test refresh_token with invalid inputs."""
    # Test with None
    with pytest.raises(ValueError, match="Invalid refresh token provided"):
        refresh_token(None)
    
    # Test with empty string
    with pytest.raises(ValueError, match="Invalid refresh token provided"):
        refresh_token("")
    
    # Test with non-string
    with pytest.raises(ValueError, match="Invalid refresh token provided"):
        refresh_token(123)

@patch('main.get_token_endpoint')
@patch('main.requests.post')
def test_refresh_token_missing_secret(mock_post, mock_get_token_endpoint):
    """Test refresh_token when client secret is not configured."""
    # Mock the token endpoint to avoid that error
    mock_get_token_endpoint.return_value = "https://login.infomaniak.com/token"
    
    # Mock the response to avoid going further in the function
    mock_response = MagicMock()
    mock_response.ok = True
    mock_response.json.return_value = {}
    mock_post.return_value = mock_response
    
    import main
    original_client_secret = main.CLIENT_SECRET
    
    # Temporarily set CLIENT_SECRET to empty string
    main.CLIENT_SECRET = ""
    
    try:
        with pytest.raises(ValueError, match="Client secret not configured"):
            refresh_token("valid-refresh-token")
    finally:
        # Restore original value
        main.CLIENT_SECRET = original_client_secret