Crypto Agent - Stage 3: Secure Implementation¶

# authentication.py (client)
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
import json

class SecureClient:
    def __init__(self, private_key_path):
        # Load client's private key
        with open(private_key_path, 'rb') as f:
            self.private_key = Ed25519PrivateKey.from_private_bytes(f.read())

    def sign_request(self, request_data):
        """Sign request with private key"""
        # 1. Create canonical message
        message = json.dumps(request_data, sort_keys=True)

        # 2. Sign with private key
        signature = self.private_key.sign(message.encode())

        # 3. Return base64-encoded signature
        return base64.b64encode(signature).decode()

# authentication.py (server)
from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PublicKey

class AuthenticationManager:
    def __init__(self):
        # Load authorized client public keys
        self.authorized_keys = self.load_authorized_keys()

    def verify_signature(self, request_data, signature, client_id):
        """Verify request signature"""
        # 1. Get client's public key
        public_key = self.authorized_keys.get(client_id)
        if not public_key:
            return False, "Unknown client"

        # 2. Recreate canonical message
        message = json.dumps(request_data, sort_keys=True)

        # 3. Verify signature
        try:
            public_key.verify(
                base64.b64decode(signature),
                message.encode()
            )
            return True, "Signature valid"
        except Exception as e:
            return False, f"Invalid signature: {e}"

# Keys are stored in separate files
authorized_clients/
├── client-123_public.key    # Client 123's public key
├── client-456_public.key    # Client 456's public key
└── client-789_public.key    # Client 789's public key

# Each client has their own private key (never shared!)
client-123/
└── private.key              # Client 123's private key (secret!)

# replay_protection.py
import redis
import uuid

class ReplayProtection:
    def __init__(self, redis_client):
        self.redis = redis_client
        self.nonce_ttl = 300  # 5 minutes

    def check_nonce(self, nonce):
        """Check if nonce has been used before"""
        # Redis key for this nonce
        key = f"nonce:{nonce}"

        # Try to set the nonce (atomic operation)
        # NX = only set if doesn't exist
        # EX = expire after ttl seconds
        was_set = self.redis.set(
            key, 
            "used", 
            nx=True,  # Only if not exists
            ex=self.nonce_ttl  # Expire after 5 min
        )

        if not was_set:
            # Nonce already exists = replay attack!
            return False, "Nonce already used (replay attack detected)"

        return True, "Nonce is fresh"

    def generate_nonce(self):
        """Generate cryptographically random nonce"""
        return str(uuid.uuid4())

def check_timestamp(self, timestamp):
    """Ensure request is recent"""
    import time

    now = time.time()
    request_time = float(timestamp)

    # Check age
    age = abs(now - request_time)

    # Reject if older than 5 minutes
    if age > 300:
        return False, f"Request too old ({age}s > 300s)"

    # Reject if in future (clock skew tolerance: 60s)
    if request_time > now + 60:
        return False, "Request timestamp in future"

    return True, "Timestamp valid"

def validate_request(self, request):
    # 1. Verify signature
    valid, msg = self.auth.verify_signature(
        request, 
        request['signature'],
        request['client_id']
    )
    if not valid:
        return False, msg

    # 2. Check timestamp
    valid, msg = self.check_timestamp(request['timestamp'])
    if not valid:
        return False, msg

    # 3. Check nonce (LAST - only if everything else is valid)
    valid, msg = self.check_nonce(request['nonce'])
    if not valid:
        return False, msg

    return True, "Request valid"

# Attacker captures valid request
captured = {
    "client_id": "client-123",
    "query": "Transfer $100",
    "timestamp": 1703001045.234,
    "nonce": "a3d2c9f1-...",
    "signature": "valid_sig..."
}

# First replay (immediately)
server.handle(captured)  # ❌ Nonce already used!

# Second replay (10 minutes later)
server.handle(captured)  # ❌ Timestamp too old!

# Modified replay (new nonce)
captured['nonce'] = "new-nonce"
server.handle(captured)  # ❌ Signature doesn't match!

# validation.py
class ComprehensiveValidator:

    # Layer 1: Type Validation
    def validate_types(self, request):
        """Ensure correct data types"""
        if not isinstance(request.get('query'), str):
            return False, "Query must be string"
        if not isinstance(request.get('timestamp'), (int, float)):
            return False, "Timestamp must be numeric"
        return True, "Types valid"

    # Layer 2: Size Validation
    def validate_sizes(self, request):
        """Check size limits"""
        query = request.get('query', '')
        if len(query) > 500:
            return False, "Query too long (max 500)"
        if len(query) < 1:
            return False, "Query too short (min 1)"

        # Check total request size
        import sys
        size = sys.getsizeof(str(request))
        if size > 10240:  # 10KB
            return False, f"Request too large ({size} > 10KB)"

        return True, "Sizes valid"

    # Layer 3: Format Validation
    def validate_format(self, request):
        """Check required fields and format"""
        required = ['client_id', 'query', 'timestamp', 'nonce']
        for field in required:
            if field not in request:
                return False, f"Missing required field: {field}"

        # Validate UUID format for nonce
        import uuid
        try:
            uuid.UUID(request['nonce'])
        except ValueError:
            return False, "Invalid nonce format"

        return True, "Format valid"

    # Layer 4: Character Validation
    def validate_characters(self, query):
        """Check for dangerous characters"""
        # Whitelist approach
        import re
        allowed_pattern = r'^[a-zA-Z0-9\s\?\.\,\!\-]+$'
        if not re.match(allowed_pattern, query):
            return False, "Query contains invalid characters"
        return True, "Characters valid"

    # Layer 5: Content Validation
    def validate_content(self, query):
        """Check for known attack patterns"""
        # SQL injection patterns
        sql_patterns = ['DROP', 'DELETE', 'UPDATE', 'INSERT', '--', ';']
        query_upper = query.upper()
        for pattern in sql_patterns:
            if pattern in query_upper:
                return False, f"Suspicious pattern detected: {pattern}"

        # Command injection patterns
        cmd_patterns = ['|', '&', '$', '`', ';', '\n']
        for pattern in cmd_patterns:
            if pattern in query:
                return False, f"Dangerous character: {pattern}"

        return True, "Content valid"

    # Layer 6: Semantic Validation
    def validate_semantics(self, query):
        """Validate query makes sense"""
        # Must mention a coin
        coins = ['BTC', 'ETH', 'SOL', 'ADA', 'DOGE', 'BITCOIN', 'ETHEREUM']
        query_upper = query.upper()
        has_coin = any(coin in query_upper for coin in coins)

        if not has_coin:
            return False, "Query must mention a cryptocurrency"

        # Must have price-related words
        price_words = ['PRICE', 'COST', 'VALUE', 'WORTH', 'TRADING']
        has_price_word = any(word in query_upper for word in price_words)

        if not has_price_word:
            return False, "Query must ask about price"

        return True, "Semantics valid"

    # Layer 7: Business Logic Validation
    def validate_business_logic(self, query, client_id):
        """Check business rules"""
        # Check if client has permission
        if not self.client_has_permission(client_id, 'price_query'):
            return False, "Client not authorized for price queries"

        # Check if requested coin is supported
        mentioned_coins = self.extract_coins(query)
        for coin in mentioned_coins:
            if coin not in self.SUPPORTED_COINS:
                return False, f"Unsupported coin: {coin}"

        return True, "Business logic valid"

    # Layer 8: Rate Limit Check
    def validate_rate_limit(self, client_id):
        """Check if client has exceeded rate limit"""
        current_rate = self.rate_limiter.get_current_rate(client_id)
        limit = self.rate_limiter.get_limit(client_id)

        if current_rate >= limit:
            return False, f"Rate limit exceeded ({current_rate}/{limit})"

        return True, "Rate limit OK"

# Sophisticated attack that bypasses some layers
malicious = {
    "query": "What is the price of BTC" + " " * 400,  # ✓ Valid chars
    "timestamp": time.time(),                          # ✓ Valid timestamp
    "nonce": str(uuid.uuid4()),                        # ✓ Valid nonce
}

# Layer 1: ✓ Types correct
# Layer 2: ✓ Size OK (500 chars)
# Layer 3: ✓ Format OK
# Layer 4: ✓ Characters OK
# Layer 5: ✓ No SQL/cmd injection
# Layer 6: ❌ BLOCKED! Semantics invalid (excessive spaces)

# rate_limit.py
import time
import redis

class TokenBucketRateLimiter:
    def __init__(self, redis_client):
        self.redis = redis_client

    def check_rate_limit(self, client_id, max_tokens=100, refill_rate=100):
        """
        Token bucket algorithm

        max_tokens: Bucket capacity (burst allowance)
        refill_rate: Tokens per minute
        """
        key = f"rate_limit:{client_id}"
        now = time.time()

        # Get current bucket state
        pipe = self.redis.pipeline()
        pipe.hmget(key, 'tokens', 'last_refill')
        pipe.expire(key, 3600)  # Expire after 1 hour of inactivity
        result = pipe.execute()[0]

        if result[0] is None:
            # First request - initialize bucket
            tokens = max_tokens - 1  # Use one token
            last_refill = now
        else:
            # Calculate refill
            tokens = float(result[0])
            last_refill = float(result[1])

            # Time since last refill
            time_passed = now - last_refill

            # Add tokens (refill_rate per minute)
            tokens_to_add = (time_passed / 60.0) * refill_rate
            tokens = min(max_tokens, tokens + tokens_to_add)

            # Use one token
            tokens -= 1

        # Check if request allowed
        if tokens < 0:
            # No tokens left - rate limited!
            return False, 0, max_tokens

        # Update bucket state
        self.redis.hmset(key, {
            'tokens': tokens,
            'last_refill': now
        })

        return True, int(tokens), max_tokens

# In request handler
allowed, remaining, limit = rate_limiter.check_rate_limit(
    client_id="client-123",
    max_tokens=100,
    refill_rate=100
)

if not allowed:
    return {
        "error": "Rate limit exceeded",
        "retry_after": 60  # seconds
    }, 429

# Add rate limit headers
response.headers['X-RateLimit-Limit'] = str(limit)
response.headers['X-RateLimit-Remaining'] = str(remaining)

# error_handling.py
import logging
import uuid

class SecureErrorHandler:
    def __init__(self):
        self.logger = logging.getLogger('security')

    def handle_error(self, error, context):
        """
        Log detailed error internally,
        return generic error to client
        """
        # Generate error ID for correlation
        error_id = str(uuid.uuid4())

        # Log detailed error internally
        self.logger.error(
            f"Error ID: {error_id}",
            extra={
                'error_id': error_id,
                'error_type': type(error).__name__,
                'error_message': str(error),
                'context': context,
                'traceback': traceback.format_exc()
            }
        )

        # Return generic error to client
        return {
            'error': self.get_generic_message(error),
            'error_id': error_id,  # For support requests
            'timestamp': time.time()
        }

    def get_generic_message(self, error):
        """Map internal errors to generic messages"""
        if isinstance(error, ValidationError):
            return "Invalid request format"
        elif isinstance(error, AuthenticationError):
            return "Authentication failed"
        elif isinstance(error, RateLimitError):
            return "Rate limit exceeded"
        else:
            return "Internal server error"

# audit.py
class AuditLogger:
    def log_request(self, request, response, duration):
        """Log every request"""
        self.logger.info('request', extra={
            # Who
            'client_id': request.client_id,
            'ip_address': request.ip,

            # What
            'endpoint': request.path,
            'method': request.method,
            'query': self.sanitize_for_log(request.query),

            # When
            'timestamp': request.timestamp,

            # How
            'signature_valid': request.signature_valid,
            'nonce': request.nonce,

            # Result
            'status_code': response.status_code,
            'response_time_ms': duration * 1000,
            'rate_limit_remaining': response.rate_limit_remaining,

            # Context
            'user_agent': request.user_agent,
            'request_id': request.id
        })

    def log_security_event(self, event_type, details):
        """Log security events"""
        self.logger.warning('security_event', extra={
            'event_type': event_type,
            'severity': self.get_severity(event_type),
            'details': details,
            'timestamp': time.time()
        })

    def sanitize_for_log(self, data):
        """Remove sensitive data before logging"""
        # Never log:
        # - Private keys
        # - Signatures (only log "present/absent")
        # - Full tokens
        # - Passwords
        # - Credit card numbers
        # - SSNs
        pass

# secure_crypto_server.py
from security.authentication import AuthenticationManager
from security.replay_protection import ReplayProtection
from security.validation import ComprehensiveValidator
from security.rate_limit import TokenBucketRateLimiter
from security.audit import AuditLogger
from security.error_handling import SecureErrorHandler

class SecureCryptoAgent:
    def __init__(self):
        # Initialize security modules
        self.auth = AuthenticationManager()
        self.replay = ReplayProtection(redis_client)
        self.validator = ComprehensiveValidator()
        self.rate_limiter = TokenBucketRateLimiter(redis_client)
        self.audit = AuditLogger()
        self.error_handler = SecureErrorHandler()

    def handle_request(self, request):
        start_time = time.time()

        try:
            # Security checks (in order)
            # 1. Authenticate
            valid, msg = self.auth.verify_signature(request)
            if not valid:
                self.audit.log_security_event('auth_failed', {'reason': msg})
                return {'error': msg}, 401

            # 2. Check replay
            valid, msg = self.replay.check_request(request)
            if not valid:
                self.audit.log_security_event('replay_detected', {'reason': msg})
                return {'error': msg}, 403

            # 3. Validate input
            valid, msg = self.validator.validate_all(request)
            if not valid:
                return {'error': msg}, 400

            # 4. Check rate limit
            allowed, remaining, limit = self.rate_limiter.check_rate_limit(
                request['client_id']
            )
            if not allowed:
                return {'error': 'Rate limit exceeded'}, 429

            # All security checks passed - process request
            response = self.process_query(request['query'])

            # Add security headers
            response['rate_limit'] = {
                'remaining': remaining,
                'limit': limit
            }

            # Log successful request
            duration = time.time() - start_time
            self.audit.log_request(request, response, duration)

            return response, 200

        except Exception as e:
            # Handle errors securely
            return self.error_handler.handle_error(e, request), 500

# Test replay protection
import requests
import time

# Make legitimate request
response1 = client.make_request("What's the price of BTC?")
captured_request = response1.request_data

# Attempt replays
print("Replay 1: Immediate (same nonce)")
response2 = requests.post(url, json=captured_request)
# Should see: "Nonce already used"

print("Replay 2: New nonce")
captured_request['nonce'] = generate_nonce()
response3 = requests.post(url, json=captured_request)
# Should see: "Invalid signature"

print("Replay 3: Old timestamp")
time.sleep(600)  # Wait 10 minutes
response4 = requests.post(url, json=captured_request)
# Should see: "Request too old"

# rate_limit_test.py
import time

# Send requests rapidly
for i in range(150):  # More than limit (100)
    response = client.make_request("What's the price of BTC?")
    print(f"Request {i+1}: {response.status_code} - "
          f"Remaining: {response.headers['X-RateLimit-Remaining']}")

    if response.status_code == 429:
        print(f"Rate limited after {i+1} requests!")
        print(f"Retry-After: {response.json()['retry_after']}s")
        break

# Wait and try again
print("Waiting 60 seconds for token refill...")
time.sleep(60)

response = client.make_request("What's the price of BTC?")
print(f"After wait: {response.status_code}")

test_cases = [
    # Layer 1: Type validation
    {"query": 123},  # Should fail - wrong type

    # Layer 2: Size validation
    {"query": "A" * 1000},  # Should fail - too long

    # Layer 3: Format validation
    {"query": "BTC price?"},  # Should fail - missing nonce

    # Layer 4: Character validation
    {"query": "BTC price<script>"},  # Should fail - invalid chars

    # Layer 5: Content validation
    {"query": "BTC; DROP TABLE"},  # Should fail - SQL injection

    # Layer 6: Semantic validation
    {"query": "Hello there"},  # Should fail - no coin mentioned

    # Layer 7: Business logic validation
    {"query": "Price of INVALIDCOIN"},  # Should fail - unsupported coin

    # Layer 8: Rate limit
    # (Already tested in Exercise 3)
]

for i, test in enumerate(test_cases):
    response = client.make_request(test['query'])
    print(f"Test {i+1}: {response.status_code} - {response.json()['error']}")