787 lines
30 KiB
Python
787 lines
30 KiB
Python
|
"""
|
||
|
|
Security tests for OAuth state handling and token storage (V-006 Fix).
|
||
|
|
|
||
|
|
Tests verify that:
|
||
|
|
1. JSON serialization is used instead of pickle
|
||
|
|
2. HMAC signatures are properly verified for both state and tokens
|
||
|
|
3. State structure is validated
|
||
|
|
4. Token schema is validated
|
||
|
|
5. Tampering is detected
|
||
|
|
6. Replay attacks are prevented
|
||
|
|
7. Timing attacks are mitigated via constant-time comparison
|
||
|
|
"""
|
||
|
|
|
||
|
|
import base64
|
||
|
|
import hashlib
|
||
|
|
import hmac
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import secrets
|
||
|
|
import sys
|
||
|
|
import tempfile
|
||
|
|
import threading
|
||
|
|
import time
|
||
|
|
from pathlib import Path
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
# Ensure tools directory is in path
|
||
|
|
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||
|
|
|
||
|
|
from tools.mcp_oauth import (
|
||
|
|
OAuthStateError,
|
||
|
|
OAuthStateManager,
|
||
|
|
SecureOAuthState,
|
||
|
|
HermesTokenStorage,
|
||
|
|
_validate_token_schema,
|
||
|
|
_OAUTH_TOKEN_SCHEMA,
|
||
|
|
_OAUTH_CLIENT_SCHEMA,
|
||
|
|
_sign_token_data,
|
||
|
|
_verify_token_signature,
|
||
|
|
_get_token_storage_key,
|
||
|
|
_state_manager,
|
||
|
|
get_state_manager,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# SecureOAuthState Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestSecureOAuthState:
|
||
|
|
"""Tests for the SecureOAuthState class."""
|
||
|
|
|
||
|
|
def test_generate_creates_valid_state(self):
|
||
|
|
"""Test that generated state has all required fields."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
|
||
|
|
assert state.token is not None
|
||
|
|
assert len(state.token) >= 16
|
||
|
|
assert state.timestamp is not None
|
||
|
|
assert isinstance(state.timestamp, float)
|
||
|
|
assert state.nonce is not None
|
||
|
|
assert len(state.nonce) >= 8
|
||
|
|
assert isinstance(state.data, dict)
|
||
|
|
|
||
|
|
def test_generate_unique_tokens(self):
|
||
|
|
"""Test that generated tokens are unique."""
|
||
|
|
tokens = {SecureOAuthState._generate_token() for _ in range(100)}
|
||
|
|
assert len(tokens) == 100
|
||
|
|
|
||
|
|
def test_serialization_format(self):
|
||
|
|
"""Test that serialized state has correct format."""
|
||
|
|
state = SecureOAuthState(data={"test": "value"})
|
||
|
|
serialized = state.serialize()
|
||
|
|
|
||
|
|
# Should have format: data.signature
|
||
|
|
parts = serialized.split(".")
|
||
|
|
assert len(parts) == 2
|
||
|
|
|
||
|
|
# Both parts should be URL-safe base64
|
||
|
|
data_part, sig_part = parts
|
||
|
|
assert all(c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_="
|
||
|
|
for c in data_part)
|
||
|
|
assert all(c in "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_="
|
||
|
|
for c in sig_part)
|
||
|
|
|
||
|
|
def test_serialize_deserialize_roundtrip(self):
|
||
|
|
"""Test that serialize/deserialize preserves state."""
|
||
|
|
original = SecureOAuthState(data={"server": "test123", "user": "alice"})
|
||
|
|
serialized = original.serialize()
|
||
|
|
deserialized = SecureOAuthState.deserialize(serialized)
|
||
|
|
|
||
|
|
assert deserialized.token == original.token
|
||
|
|
assert deserialized.timestamp == original.timestamp
|
||
|
|
assert deserialized.nonce == original.nonce
|
||
|
|
assert deserialized.data == original.data
|
||
|
|
|
||
|
|
def test_deserialize_empty_raises_error(self):
|
||
|
|
"""Test that deserializing empty state raises OAuthStateError."""
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize("")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "empty or wrong type" in str(e)
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(None)
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "empty or wrong type" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_missing_signature_raises_error(self):
|
||
|
|
"""Test that missing signature is detected."""
|
||
|
|
data = json.dumps({"test": "data"})
|
||
|
|
encoded = base64.urlsafe_b64encode(data.encode()).decode()
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(encoded)
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "missing signature" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_invalid_base64_raises_error(self):
|
||
|
|
"""Test that invalid data is rejected (base64 or signature)."""
|
||
|
|
# Invalid characters may be accepted by Python's base64 decoder
|
||
|
|
# but signature verification should fail
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize("!!!invalid!!!.!!!data!!!")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
# Error could be from encoding or signature verification
|
||
|
|
assert "Invalid state" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_tampered_signature_detected(self):
|
||
|
|
"""Test that tampered signature is detected."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
serialized = state.serialize()
|
||
|
|
|
||
|
|
# Tamper with the signature
|
||
|
|
data_part, sig_part = serialized.split(".")
|
||
|
|
tampered_sig = base64.urlsafe_b64encode(b"tampered").decode().rstrip("=")
|
||
|
|
tampered = f"{data_part}.{tampered_sig}"
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(tampered)
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "tampering detected" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_tampered_data_detected(self):
|
||
|
|
"""Test that tampered data is detected via HMAC verification."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
serialized = state.serialize()
|
||
|
|
|
||
|
|
# Tamper with the data but keep signature
|
||
|
|
data_part, sig_part = serialized.split(".")
|
||
|
|
tampered_data = json.dumps({"hacked": True})
|
||
|
|
tampered_encoded = base64.urlsafe_b64encode(tampered_data.encode()).decode().rstrip("=")
|
||
|
|
tampered = f"{tampered_encoded}.{sig_part}"
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(tampered)
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "tampering detected" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_expired_state_raises_error(self):
|
||
|
|
"""Test that expired states are rejected."""
|
||
|
|
# Create a state with old timestamp
|
||
|
|
old_state = SecureOAuthState()
|
||
|
|
old_state.timestamp = time.time() - 1000 # 1000 seconds ago
|
||
|
|
|
||
|
|
serialized = old_state.serialize()
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(serialized)
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "expired" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_invalid_json_raises_error(self):
|
||
|
|
"""Test that invalid JSON raises OAuthStateError."""
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
bad_data = b"not valid json {{{"
|
||
|
|
sig = base64.urlsafe_b64encode(hmac.new(key, bad_data, hashlib.sha256).digest())
|
||
|
|
encoded_data = base64.urlsafe_b64encode(bad_data).decode().rstrip("=")
|
||
|
|
encoded_sig = sig.decode().rstrip("=")
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(f"{encoded_data}.{encoded_sig}")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "Invalid state JSON" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_missing_fields_raises_error(self):
|
||
|
|
"""Test that missing required fields are detected."""
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
bad_data = json.dumps({"token": "test"}).encode() # missing timestamp, nonce
|
||
|
|
sig = base64.urlsafe_b64encode(hmac.new(key, bad_data, hashlib.sha256).digest())
|
||
|
|
encoded_data = base64.urlsafe_b64encode(bad_data).decode().rstrip("=")
|
||
|
|
encoded_sig = sig.decode().rstrip("=")
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(f"{encoded_data}.{encoded_sig}")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "missing fields" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_invalid_token_type_raises_error(self):
|
||
|
|
"""Test that non-string tokens are rejected."""
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
bad_data = json.dumps({
|
||
|
|
"token": 12345, # should be string
|
||
|
|
"timestamp": time.time(),
|
||
|
|
"nonce": "abc123"
|
||
|
|
}).encode()
|
||
|
|
sig = base64.urlsafe_b64encode(hmac.new(key, bad_data, hashlib.sha256).digest())
|
||
|
|
encoded_data = base64.urlsafe_b64encode(bad_data).decode().rstrip("=")
|
||
|
|
encoded_sig = sig.decode().rstrip("=")
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(f"{encoded_data}.{encoded_sig}")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "token must be a string" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_short_token_raises_error(self):
|
||
|
|
"""Test that short tokens are rejected."""
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
bad_data = json.dumps({
|
||
|
|
"token": "short", # too short
|
||
|
|
"timestamp": time.time(),
|
||
|
|
"nonce": "abc123"
|
||
|
|
}).encode()
|
||
|
|
sig = base64.urlsafe_b64encode(hmac.new(key, bad_data, hashlib.sha256).digest())
|
||
|
|
encoded_data = base64.urlsafe_b64encode(bad_data).decode().rstrip("=")
|
||
|
|
encoded_sig = sig.decode().rstrip("=")
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(f"{encoded_data}.{encoded_sig}")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "token must be a string" in str(e)
|
||
|
|
|
||
|
|
def test_deserialize_invalid_timestamp_raises_error(self):
|
||
|
|
"""Test that non-numeric timestamps are rejected."""
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
bad_data = json.dumps({
|
||
|
|
"token": "a" * 32,
|
||
|
|
"timestamp": "not a number",
|
||
|
|
"nonce": "abc123"
|
||
|
|
}).encode()
|
||
|
|
sig = base64.urlsafe_b64encode(hmac.new(key, bad_data, hashlib.sha256).digest())
|
||
|
|
encoded_data = base64.urlsafe_b64encode(bad_data).decode().rstrip("=")
|
||
|
|
encoded_sig = sig.decode().rstrip("=")
|
||
|
|
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(f"{encoded_data}.{encoded_sig}")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "timestamp must be numeric" in str(e)
|
||
|
|
|
||
|
|
def test_validate_against_correct_token(self):
|
||
|
|
"""Test token validation with matching token."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
assert state.validate_against(state.token) is True
|
||
|
|
|
||
|
|
def test_validate_against_wrong_token(self):
|
||
|
|
"""Test token validation with non-matching token."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
assert state.validate_against("wrong-token") is False
|
||
|
|
|
||
|
|
def test_validate_against_non_string(self):
|
||
|
|
"""Test token validation with non-string input."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
assert state.validate_against(None) is False
|
||
|
|
assert state.validate_against(12345) is False
|
||
|
|
|
||
|
|
def test_validate_uses_constant_time_comparison(self):
|
||
|
|
"""Test that validate_against uses constant-time comparison."""
|
||
|
|
state = SecureOAuthState(token="test-token-for-comparison")
|
||
|
|
|
||
|
|
# This test verifies no early return on mismatch
|
||
|
|
# In practice, secrets.compare_digest is used
|
||
|
|
result1 = state.validate_against("wrong-token-for-comparison")
|
||
|
|
result2 = state.validate_against("another-wrong-token-here")
|
||
|
|
|
||
|
|
assert result1 is False
|
||
|
|
assert result2 is False
|
||
|
|
|
||
|
|
def test_to_dict_format(self):
|
||
|
|
"""Test that to_dict returns correct format."""
|
||
|
|
state = SecureOAuthState(data={"custom": "data"})
|
||
|
|
d = state.to_dict()
|
||
|
|
|
||
|
|
assert set(d.keys()) == {"token", "timestamp", "nonce", "data"}
|
||
|
|
assert d["token"] == state.token
|
||
|
|
assert d["timestamp"] == state.timestamp
|
||
|
|
assert d["nonce"] == state.nonce
|
||
|
|
assert d["data"] == {"custom": "data"}
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# OAuthStateManager Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestOAuthStateManager:
|
||
|
|
"""Tests for the OAuthStateManager class."""
|
||
|
|
|
||
|
|
def setup_method(self):
|
||
|
|
"""Reset state manager before each test."""
|
||
|
|
global _state_manager
|
||
|
|
_state_manager.invalidate()
|
||
|
|
_state_manager._used_nonces.clear()
|
||
|
|
|
||
|
|
def test_generate_state_returns_serialized(self):
|
||
|
|
"""Test that generate_state returns a serialized state string."""
|
||
|
|
state_str = _state_manager.generate_state()
|
||
|
|
|
||
|
|
# Should be a string with format: data.signature
|
||
|
|
assert isinstance(state_str, str)
|
||
|
|
assert "." in state_str
|
||
|
|
parts = state_str.split(".")
|
||
|
|
assert len(parts) == 2
|
||
|
|
|
||
|
|
def test_generate_state_with_data(self):
|
||
|
|
"""Test that extra data is included in state."""
|
||
|
|
extra = {"server_name": "test-server", "user_id": "123"}
|
||
|
|
state_str = _state_manager.generate_state(extra_data=extra)
|
||
|
|
|
||
|
|
# Validate and extract
|
||
|
|
is_valid, data = _state_manager.validate_and_extract(state_str)
|
||
|
|
assert is_valid is True
|
||
|
|
assert data == extra
|
||
|
|
|
||
|
|
def test_validate_and_extract_valid_state(self):
|
||
|
|
"""Test validation with a valid state."""
|
||
|
|
extra = {"test": "data"}
|
||
|
|
state_str = _state_manager.generate_state(extra_data=extra)
|
||
|
|
|
||
|
|
is_valid, data = _state_manager.validate_and_extract(state_str)
|
||
|
|
|
||
|
|
assert is_valid is True
|
||
|
|
assert data == extra
|
||
|
|
|
||
|
|
def test_validate_and_extract_none_state(self):
|
||
|
|
"""Test validation with None state."""
|
||
|
|
is_valid, data = _state_manager.validate_and_extract(None)
|
||
|
|
|
||
|
|
assert is_valid is False
|
||
|
|
assert data is None
|
||
|
|
|
||
|
|
def test_validate_and_extract_invalid_state(self):
|
||
|
|
"""Test validation with an invalid state."""
|
||
|
|
is_valid, data = _state_manager.validate_and_extract("invalid.state.here")
|
||
|
|
|
||
|
|
assert is_valid is False
|
||
|
|
assert data is None
|
||
|
|
|
||
|
|
def test_state_cleared_after_validation(self):
|
||
|
|
"""Test that state is cleared after successful validation."""
|
||
|
|
state_str = _state_manager.generate_state()
|
||
|
|
|
||
|
|
# First validation should succeed
|
||
|
|
is_valid1, _ = _state_manager.validate_and_extract(state_str)
|
||
|
|
assert is_valid1 is True
|
||
|
|
|
||
|
|
# Second validation should fail (replay)
|
||
|
|
is_valid2, _ = _state_manager.validate_and_extract(state_str)
|
||
|
|
assert is_valid2 is False
|
||
|
|
|
||
|
|
def test_nonce_tracking_prevents_replay(self):
|
||
|
|
"""Test that nonce tracking prevents replay attacks."""
|
||
|
|
state = SecureOAuthState()
|
||
|
|
serialized = state.serialize()
|
||
|
|
|
||
|
|
# Manually add to used nonces
|
||
|
|
with _state_manager._lock:
|
||
|
|
_state_manager._used_nonces.add(state.nonce)
|
||
|
|
|
||
|
|
# Validation should fail due to nonce replay
|
||
|
|
is_valid, _ = _state_manager.validate_and_extract(serialized)
|
||
|
|
assert is_valid is False
|
||
|
|
|
||
|
|
def test_invalidate_clears_state(self):
|
||
|
|
"""Test that invalidate clears the stored state."""
|
||
|
|
_state_manager.generate_state()
|
||
|
|
assert _state_manager._state is not None
|
||
|
|
|
||
|
|
_state_manager.invalidate()
|
||
|
|
assert _state_manager._state is None
|
||
|
|
|
||
|
|
def test_thread_safety(self):
|
||
|
|
"""Test thread safety of state manager."""
|
||
|
|
results = []
|
||
|
|
|
||
|
|
def generate():
|
||
|
|
state_str = _state_manager.generate_state(extra_data={"thread": threading.current_thread().name})
|
||
|
|
results.append(state_str)
|
||
|
|
|
||
|
|
threads = [threading.Thread(target=generate) for _ in range(10)]
|
||
|
|
for t in threads:
|
||
|
|
t.start()
|
||
|
|
for t in threads:
|
||
|
|
t.join()
|
||
|
|
|
||
|
|
# All states should be unique
|
||
|
|
assert len(set(results)) == 10
|
||
|
|
|
||
|
|
def test_max_nonce_limit(self):
|
||
|
|
"""Test that nonce set is limited to prevent memory growth."""
|
||
|
|
manager = OAuthStateManager()
|
||
|
|
manager._max_used_nonces = 5
|
||
|
|
|
||
|
|
# Generate more nonces than the limit
|
||
|
|
for _ in range(10):
|
||
|
|
state = SecureOAuthState()
|
||
|
|
manager._used_nonces.add(state.nonce)
|
||
|
|
|
||
|
|
# Set should have been cleared at some point
|
||
|
|
# (implementation clears when limit is exceeded)
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# Schema Validation Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestSchemaValidation:
|
||
|
|
"""Tests for JSON schema validation (V-006)."""
|
||
|
|
|
||
|
|
def test_valid_token_schema_accepted(self):
|
||
|
|
"""Test that valid token data passes schema validation."""
|
||
|
|
valid_token = {
|
||
|
|
"access_token": "secret_token_123",
|
||
|
|
"token_type": "Bearer",
|
||
|
|
"refresh_token": "refresh_456",
|
||
|
|
"expires_in": 3600,
|
||
|
|
"expires_at": 1234567890.0,
|
||
|
|
"scope": "read write",
|
||
|
|
"id_token": "id_token_789",
|
||
|
|
}
|
||
|
|
# Should not raise
|
||
|
|
_validate_token_schema(valid_token, _OAUTH_TOKEN_SCHEMA, "token")
|
||
|
|
|
||
|
|
def test_minimal_valid_token_schema(self):
|
||
|
|
"""Test that minimal valid token (only required fields) passes."""
|
||
|
|
minimal_token = {
|
||
|
|
"access_token": "secret",
|
||
|
|
"token_type": "Bearer",
|
||
|
|
}
|
||
|
|
_validate_token_schema(minimal_token, _OAUTH_TOKEN_SCHEMA, "token")
|
||
|
|
|
||
|
|
def test_missing_required_field_rejected(self):
|
||
|
|
"""Test that missing required fields are detected."""
|
||
|
|
invalid_token = {"token_type": "Bearer"} # missing access_token
|
||
|
|
try:
|
||
|
|
_validate_token_schema(invalid_token, _OAUTH_TOKEN_SCHEMA, "token")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "missing required fields" in str(e)
|
||
|
|
assert "access_token" in str(e)
|
||
|
|
|
||
|
|
def test_wrong_type_rejected(self):
|
||
|
|
"""Test that fields with wrong types are rejected."""
|
||
|
|
invalid_token = {
|
||
|
|
"access_token": 12345, # should be string
|
||
|
|
"token_type": "Bearer",
|
||
|
|
}
|
||
|
|
try:
|
||
|
|
_validate_token_schema(invalid_token, _OAUTH_TOKEN_SCHEMA, "token")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "has wrong type" in str(e)
|
||
|
|
|
||
|
|
def test_null_values_accepted(self):
|
||
|
|
"""Test that null values for optional fields are accepted."""
|
||
|
|
token_with_nulls = {
|
||
|
|
"access_token": "secret",
|
||
|
|
"token_type": "Bearer",
|
||
|
|
"refresh_token": None,
|
||
|
|
"expires_in": None,
|
||
|
|
}
|
||
|
|
_validate_token_schema(token_with_nulls, _OAUTH_TOKEN_SCHEMA, "token")
|
||
|
|
|
||
|
|
def test_non_dict_data_rejected(self):
|
||
|
|
"""Test that non-dictionary data is rejected."""
|
||
|
|
try:
|
||
|
|
_validate_token_schema("not a dict", _OAUTH_TOKEN_SCHEMA, "token")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "must be a dictionary" in str(e)
|
||
|
|
|
||
|
|
def test_valid_client_schema(self):
|
||
|
|
"""Test that valid client info passes schema validation."""
|
||
|
|
valid_client = {
|
||
|
|
"client_id": "client_123",
|
||
|
|
"client_secret": "secret_456",
|
||
|
|
"client_name": "Test Client",
|
||
|
|
"redirect_uris": ["http://localhost/callback"],
|
||
|
|
}
|
||
|
|
_validate_token_schema(valid_client, _OAUTH_CLIENT_SCHEMA, "client")
|
||
|
|
|
||
|
|
def test_client_missing_required_rejected(self):
|
||
|
|
"""Test that client info missing client_id is rejected."""
|
||
|
|
invalid_client = {"client_name": "Test"}
|
||
|
|
try:
|
||
|
|
_validate_token_schema(invalid_client, _OAUTH_CLIENT_SCHEMA, "client")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "missing required fields" in str(e)
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# Token Storage Security Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestTokenStorageSecurity:
|
||
|
|
"""Tests for token storage signing and validation (V-006)."""
|
||
|
|
|
||
|
|
def test_sign_and_verify_token_data(self):
|
||
|
|
"""Test that token data can be signed and verified."""
|
||
|
|
data = {"access_token": "test123", "token_type": "Bearer"}
|
||
|
|
sig = _sign_token_data(data)
|
||
|
|
|
||
|
|
assert sig is not None
|
||
|
|
assert len(sig) > 0
|
||
|
|
assert _verify_token_signature(data, sig) is True
|
||
|
|
|
||
|
|
def test_tampered_token_data_rejected(self):
|
||
|
|
"""Test that tampered token data fails verification."""
|
||
|
|
data = {"access_token": "test123", "token_type": "Bearer"}
|
||
|
|
sig = _sign_token_data(data)
|
||
|
|
|
||
|
|
# Modify the data
|
||
|
|
tampered_data = {"access_token": "hacked", "token_type": "Bearer"}
|
||
|
|
assert _verify_token_signature(tampered_data, sig) is False
|
||
|
|
|
||
|
|
def test_empty_signature_rejected(self):
|
||
|
|
"""Test that empty signature is rejected."""
|
||
|
|
data = {"access_token": "test", "token_type": "Bearer"}
|
||
|
|
assert _verify_token_signature(data, "") is False
|
||
|
|
|
||
|
|
def test_invalid_signature_rejected(self):
|
||
|
|
"""Test that invalid signature is rejected."""
|
||
|
|
data = {"access_token": "test", "token_type": "Bearer"}
|
||
|
|
assert _verify_token_signature(data, "invalid") is False
|
||
|
|
|
||
|
|
def test_signature_deterministic(self):
|
||
|
|
"""Test that signing the same data produces the same signature."""
|
||
|
|
data = {"access_token": "test123", "token_type": "Bearer"}
|
||
|
|
sig1 = _sign_token_data(data)
|
||
|
|
sig2 = _sign_token_data(data)
|
||
|
|
assert sig1 == sig2
|
||
|
|
|
||
|
|
def test_different_data_different_signatures(self):
|
||
|
|
"""Test that different data produces different signatures."""
|
||
|
|
data1 = {"access_token": "test1", "token_type": "Bearer"}
|
||
|
|
data2 = {"access_token": "test2", "token_type": "Bearer"}
|
||
|
|
sig1 = _sign_token_data(data1)
|
||
|
|
sig2 = _sign_token_data(data2)
|
||
|
|
assert sig1 != sig2
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# Pickle Security Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestNoPickleUsage:
|
||
|
|
"""Tests to verify pickle is NOT used (V-006 regression tests)."""
|
||
|
|
|
||
|
|
def test_serialization_does_not_use_pickle(self):
|
||
|
|
"""Verify that state serialization uses JSON, not pickle."""
|
||
|
|
state = SecureOAuthState(data={"malicious": "__import__('os').system('rm -rf /')"})
|
||
|
|
serialized = state.serialize()
|
||
|
|
|
||
|
|
# Decode the data part
|
||
|
|
data_part, _ = serialized.split(".")
|
||
|
|
padding = 4 - (len(data_part) % 4) if len(data_part) % 4 else 0
|
||
|
|
decoded = base64.urlsafe_b64decode(data_part + ("=" * padding))
|
||
|
|
|
||
|
|
# Should be valid JSON, not pickle
|
||
|
|
parsed = json.loads(decoded.decode('utf-8'))
|
||
|
|
assert parsed["data"]["malicious"] == "__import__('os').system('rm -rf /')"
|
||
|
|
|
||
|
|
# Should NOT start with pickle protocol markers
|
||
|
|
assert not decoded.startswith(b'\x80') # Pickle protocol marker
|
||
|
|
assert b'cos\n' not in decoded # Pickle module load pattern
|
||
|
|
|
||
|
|
def test_deserialize_rejects_pickle_payload(self):
|
||
|
|
"""Test that pickle payloads are rejected during deserialization."""
|
||
|
|
import pickle
|
||
|
|
|
||
|
|
# Create a pickle payload that would execute code
|
||
|
|
malicious = pickle.dumps({"cmd": "whoami"})
|
||
|
|
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
sig = base64.urlsafe_b64encode(
|
||
|
|
hmac.new(key, malicious, hashlib.sha256).digest()
|
||
|
|
).decode().rstrip("=")
|
||
|
|
data = base64.urlsafe_b64encode(malicious).decode().rstrip("=")
|
||
|
|
|
||
|
|
# Should fail because it's not valid JSON
|
||
|
|
try:
|
||
|
|
SecureOAuthState.deserialize(f"{data}.{sig}")
|
||
|
|
assert False, "Should have raised OAuthStateError"
|
||
|
|
except OAuthStateError as e:
|
||
|
|
assert "Invalid state JSON" in str(e)
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# Key Management Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestSecretKeyManagement:
|
||
|
|
"""Tests for HMAC secret key management."""
|
||
|
|
|
||
|
|
def test_get_secret_key_from_env(self):
|
||
|
|
"""Test that HERMES_OAUTH_SECRET environment variable is used."""
|
||
|
|
with patch.dict(os.environ, {"HERMES_OAUTH_SECRET": "test-secret-key-32bytes-long!!"}):
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
assert key == b"test-secret-key-32bytes-long!!"
|
||
|
|
|
||
|
|
def test_get_token_storage_key_from_env(self):
|
||
|
|
"""Test that HERMES_TOKEN_STORAGE_SECRET environment variable is used."""
|
||
|
|
with patch.dict(os.environ, {"HERMES_TOKEN_STORAGE_SECRET": "storage-secret-key-32bytes!!"}):
|
||
|
|
key = _get_token_storage_key()
|
||
|
|
assert key == b"storage-secret-key-32bytes!!"
|
||
|
|
|
||
|
|
def test_get_secret_key_creates_file(self):
|
||
|
|
"""Test that secret key file is created if it doesn't exist."""
|
||
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
||
|
|
home = Path(tmpdir)
|
||
|
|
with patch('pathlib.Path.home', return_value=home):
|
||
|
|
with patch.dict(os.environ, {}, clear=True):
|
||
|
|
key = SecureOAuthState._get_secret_key()
|
||
|
|
assert len(key) == 64
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# Integration Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestOAuthFlowIntegration:
|
||
|
|
"""Integration tests for the OAuth flow with secure state."""
|
||
|
|
|
||
|
|
def setup_method(self):
|
||
|
|
"""Reset state manager before each test."""
|
||
|
|
global _state_manager
|
||
|
|
_state_manager.invalidate()
|
||
|
|
_state_manager._used_nonces.clear()
|
||
|
|
|
||
|
|
def test_full_oauth_state_flow(self):
|
||
|
|
"""Test the full OAuth state generation and validation flow."""
|
||
|
|
# Step 1: Generate state for OAuth request
|
||
|
|
server_name = "test-mcp-server"
|
||
|
|
state = _state_manager.generate_state(extra_data={"server_name": server_name})
|
||
|
|
|
||
|
|
# Step 2: Simulate OAuth callback with state
|
||
|
|
# (In real flow, this comes back from OAuth provider)
|
||
|
|
is_valid, data = _state_manager.validate_and_extract(state)
|
||
|
|
|
||
|
|
# Step 3: Verify validation succeeded
|
||
|
|
assert is_valid is True
|
||
|
|
assert data["server_name"] == server_name
|
||
|
|
|
||
|
|
# Step 4: Verify state cannot be replayed
|
||
|
|
is_valid_replay, _ = _state_manager.validate_and_extract(state)
|
||
|
|
assert is_valid_replay is False
|
||
|
|
|
||
|
|
def test_csrf_attack_prevention(self):
|
||
|
|
"""Test that CSRF attacks using different states are detected."""
|
||
|
|
# Attacker generates their own state
|
||
|
|
attacker_state = _state_manager.generate_state(extra_data={"malicious": True})
|
||
|
|
|
||
|
|
# Victim generates their state
|
||
|
|
victim_manager = OAuthStateManager()
|
||
|
|
victim_state = victim_manager.generate_state(extra_data={"legitimate": True})
|
||
|
|
|
||
|
|
# Attacker tries to use their state with victim's session
|
||
|
|
# This would fail because the tokens don't match
|
||
|
|
is_valid, _ = victim_manager.validate_and_extract(attacker_state)
|
||
|
|
assert is_valid is False
|
||
|
|
|
||
|
|
def test_mitm_attack_detection(self):
|
||
|
|
"""Test that tampered states from MITM attacks are detected."""
|
||
|
|
# Generate legitimate state
|
||
|
|
state = _state_manager.generate_state()
|
||
|
|
|
||
|
|
# Modify the state (simulating MITM tampering)
|
||
|
|
parts = state.split(".")
|
||
|
|
tampered_state = parts[0] + ".tampered-signature-here"
|
||
|
|
|
||
|
|
# Validation should fail
|
||
|
|
is_valid, _ = _state_manager.validate_and_extract(tampered_state)
|
||
|
|
assert is_valid is False
|
||
|
|
|
||
|
|
|
||
|
|
# =============================================================================
|
||
|
|
# Performance Tests
|
||
|
|
# =============================================================================
|
||
|
|
|
||
|
|
class TestPerformance:
|
||
|
|
"""Performance tests for state operations."""
|
||
|
|
|
||
|
|
def test_serialize_performance(self):
|
||
|
|
"""Test that serialization is fast."""
|
||
|
|
state = SecureOAuthState(data={"key": "value" * 100})
|
||
|
|
|
||
|
|
start = time.time()
|
||
|
|
for _ in range(1000):
|
||
|
|
state.serialize()
|
||
|
|
elapsed = time.time() - start
|
||
|
|
|
||
|
|
# Should complete 1000 serializations in under 1 second
|
||
|
|
assert elapsed < 1.0
|
||
|
|
|
||
|
|
def test_deserialize_performance(self):
|
||
|
|
"""Test that deserialization is fast."""
|
||
|
|
state = SecureOAuthState(data={"key": "value" * 100})
|
||
|
|
serialized = state.serialize()
|
||
|
|
|
||
|
|
start = time.time()
|
||
|
|
for _ in range(1000):
|
||
|
|
SecureOAuthState.deserialize(serialized)
|
||
|
|
elapsed = time.time() - start
|
||
|
|
|
||
|
|
# Should complete 1000 deserializations in under 1 second
|
||
|
|
assert elapsed < 1.0
|
||
|
|
|
||
|
|
|
||
|
|
def run_tests():
|
||
|
|
"""Run all tests."""
|
||
|
|
import inspect
|
||
|
|
|
||
|
|
test_classes = [
|
||
|
|
TestSecureOAuthState,
|
||
|
|
TestOAuthStateManager,
|
||
|
|
TestSchemaValidation,
|
||
|
|
TestTokenStorageSecurity,
|
||
|
|
TestNoPickleUsage,
|
||
|
|
TestSecretKeyManagement,
|
||
|
|
TestOAuthFlowIntegration,
|
||
|
|
TestPerformance,
|
||
|
|
]
|
||
|
|
|
||
|
|
total_tests = 0
|
||
|
|
passed_tests = 0
|
||
|
|
failed_tests = []
|
||
|
|
|
||
|
|
for cls in test_classes:
|
||
|
|
print(f"\n{'='*60}")
|
||
|
|
print(f"Running {cls.__name__}")
|
||
|
|
print('='*60)
|
||
|
|
|
||
|
|
instance = cls()
|
||
|
|
|
||
|
|
# Run setup if exists
|
||
|
|
if hasattr(instance, 'setup_method'):
|
||
|
|
instance.setup_method()
|
||
|
|
|
||
|
|
for name, method in inspect.getmembers(cls, predicate=inspect.isfunction):
|
||
|
|
if name.startswith('test_'):
|
||
|
|
total_tests += 1
|
||
|
|
try:
|
||
|
|
method(instance)
|
||
|
|
print(f" ✓ {name}")
|
||
|
|
passed_tests += 1
|
||
|
|
except Exception as e:
|
||
|
|
print(f" ✗ {name}: {e}")
|
||
|
|
failed_tests.append((cls.__name__, name, str(e)))
|
||
|
|
|
||
|
|
print(f"\n{'='*60}")
|
||
|
|
print(f"Results: {passed_tests}/{total_tests} tests passed")
|
||
|
|
print('='*60)
|
||
|
|
|
||
|
|
if failed_tests:
|
||
|
|
print("\nFailed tests:")
|
||
|
|
for cls_name, test_name, error in failed_tests:
|
||
|
|
print(f" - {cls_name}.{test_name}: {error}")
|
||
|
|
return 1
|
||
|
|
else:
|
||
|
|
print("\nAll tests passed!")
|
||
|
|
return 0
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.exit(run_tests())
|