diff --git a/agent-framework/prometheus_swarm/utils/nonce.py b/agent-framework/prometheus_swarm/utils/nonce.py new file mode 100644 index 00000000..6bd15b8e --- /dev/null +++ b/agent-framework/prometheus_swarm/utils/nonce.py @@ -0,0 +1,91 @@ +import uuid +import time +import hashlib +from typing import Dict, Any, Optional + +class NonceRequestInterface: + """ + A class to manage nonce (number used once) requests for secure communication. + + The nonce is a unique, one-time use value that helps prevent replay attacks + and ensures the uniqueness of each request. + """ + + def __init__(self, max_age_seconds: int = 300): + """ + Initialize the Nonce Request Interface. + + :param max_age_seconds: Maximum time a nonce is considered valid (default 5 minutes) + """ + self._nonce_store: Dict[str, Dict[str, Any]] = {} + self._max_age_seconds = max_age_seconds + + def generate_nonce(self, context: Optional[str] = None) -> str: + """ + Generate a unique nonce. + + :param context: Optional context for the nonce (e.g., user ID, request type) + :return: A unique nonce string + """ + # Generate a UUID-based nonce + nonce = str(uuid.uuid4()) + + # Store nonce details + self._nonce_store[nonce] = { + 'created_at': time.time(), + 'context': context, + 'used': False + } + + return nonce + + def validate_nonce(self, nonce: str, context: Optional[str] = None) -> bool: + """ + Validate a nonce. + + :param nonce: Nonce to validate + :param context: Optional context to match against + :return: True if nonce is valid, False otherwise + """ + # Check if nonce exists + if nonce not in self._nonce_store: + return False + + # Get nonce details + nonce_details = self._nonce_store[nonce] + + # Check if nonce has already been used + if nonce_details['used']: + return False + + # Check nonce age + current_time = time.time() + if current_time - nonce_details['created_at'] > self._max_age_seconds: + del self._nonce_store[nonce] + return False + + # Check context if provided + if context is not None and nonce_details['context'] != context: + return False + + # Mark nonce as used + nonce_details['used'] = True + + return True + + def clear_expired_nonces(self) -> int: + """ + Clear expired nonces from the store. + + :return: Number of nonces removed + """ + current_time = time.time() + expired_nonces = [ + nonce for nonce, details in self._nonce_store.items() + if details['used'] or current_time - details['created_at'] > self._max_age_seconds + ] + + for nonce in expired_nonces: + del self._nonce_store[nonce] + + return len(expired_nonces) \ No newline at end of file diff --git a/agent-framework/tests/unit/test_nonce_interface.py b/agent-framework/tests/unit/test_nonce_interface.py new file mode 100644 index 00000000..dab662d8 --- /dev/null +++ b/agent-framework/tests/unit/test_nonce_interface.py @@ -0,0 +1,79 @@ +import time +import pytest +from prometheus_swarm.utils.nonce import NonceRequestInterface + +def test_generate_nonce(): + """Test generating a unique nonce""" + nonce_interface = NonceRequestInterface() + nonce1 = nonce_interface.generate_nonce() + nonce2 = nonce_interface.generate_nonce() + + assert nonce1 != nonce2 + assert isinstance(nonce1, str) + assert len(nonce1) > 0 + +def test_validate_nonce(): + """Test validating a generated nonce""" + nonce_interface = NonceRequestInterface() + nonce = nonce_interface.generate_nonce() + + assert nonce_interface.validate_nonce(nonce) is True + assert nonce_interface.validate_nonce(nonce) is False # Cannot reuse nonce + +def test_nonce_context(): + """Test nonce validation with context""" + nonce_interface = NonceRequestInterface() + nonce1 = nonce_interface.generate_nonce(context='user1') + nonce2 = nonce_interface.generate_nonce(context='user2') + + assert nonce_interface.validate_nonce(nonce1, context='user1') is True + assert nonce_interface.validate_nonce(nonce1, context='user2') is False + assert nonce_interface.validate_nonce(nonce2, context='user1') is False + +def test_nonce_expiration(): + """Test nonce expiration""" + nonce_interface = NonceRequestInterface(max_age_seconds=1) + nonce = nonce_interface.generate_nonce() + + time.sleep(2) # Wait for nonce to expire + assert nonce_interface.validate_nonce(nonce) is False + +def test_clear_expired_nonces(): + """Test clearing expired nonces""" + nonce_interface = NonceRequestInterface(max_age_seconds=1) + + # Generate multiple nonces + nonces = [nonce_interface.generate_nonce() for _ in range(5)] + + # Use some nonces + nonce_interface.validate_nonce(nonces[0]) + nonce_interface.validate_nonce(nonces[1]) + + time.sleep(2) # Wait for nonces to expire + + # Clear expired nonces + cleared_count = nonce_interface.clear_expired_nonces() + assert cleared_count > 0 + +def test_invalid_nonce(): + """Test validation of an invalid nonce""" + nonce_interface = NonceRequestInterface() + invalid_nonce = 'not-a-real-nonce' + + assert nonce_interface.validate_nonce(invalid_nonce) is False + +def test_nonce_thread_safety(mocker): + """Test basic thread safety of nonce generation""" + nonce_interface = NonceRequestInterface() + + # Simulate multiple threads generating nonces + mock_threads = [ + mocker.Mock(name=f'thread_{i}') for i in range(10) + ] + + nonces = set() + for thread in mock_threads: + nonce = nonce_interface.generate_nonce() + nonces.add(nonce) + + assert len(nonces) == len(mock_threads) \ No newline at end of file