#!/usr/bin/env python3 """ Token Bucket Rate Limiter Demo Server A simple HTTP server demonstrating the token bucket rate limiter. Run this server and visit http://localhost:8080 in your browser to test the rate limiting behavior interactively. Usage: python demo_server.py [--port PORT] [--redis-host HOST] [--redis-port PORT] """ import argparse import json import sys from http.server import HTTPServer, BaseHTTPRequestHandler from urllib.parse import parse_qs, urlparse # Add current directory to path to import token_bucket sys.path.insert(0, '.') try: from token_bucket import TokenBucket import redis except ImportError as e: print(f"Error: {e}") print("Make sure 'redis' package is installed: pip install redis") sys.exit(1) class RateLimiterHandler(BaseHTTPRequestHandler): """HTTP request handler for the rate limiter demo.""" # Class-level variables for configuration limiter = None config = { 'capacity': 10, 'refill_rate': 1, 'refill_interval': 1.0 } def do_GET(self): """Handle GET requests - serve the main page.""" if self.path == '/' or self.path == '/index.html': self.send_response(200) self.send_header('Content-type', 'text/html') self.end_headers() self.wfile.write(self._get_html_page().encode()) else: self.send_error(404) def do_POST(self): """Handle POST requests - test rate limiting or update config.""" parsed_path = urlparse(self.path) if parsed_path.path == '/test': self._handle_test_request() elif parsed_path.path == '/config': self._handle_config_update() else: self.send_error(404) def _handle_test_request(self): """Test the rate limiter and return the result.""" allowed, remaining = self.limiter.allow('demo:request') response = { 'allowed': allowed, 'remaining': remaining, 'config': self.config } self.send_response(200) self.send_header('Content-type', 'application/json') self.end_headers() self.wfile.write(json.dumps(response).encode()) def _handle_config_update(self): """Update the rate limiter configuration.""" content_length = int(self.headers['Content-Length']) post_data = self.rfile.read(content_length) params = parse_qs(post_data.decode()) try: capacity = int(params.get('capacity', [self.config['capacity']])[0]) refill_rate = int(params.get('refill_rate', [self.config['refill_rate']])[0]) refill_interval = float(params.get('refill_interval', [self.config['refill_interval']])[0]) # Update configuration self.config['capacity'] = capacity self.config['refill_rate'] = refill_rate self.config['refill_interval'] = refill_interval # Recreate limiter with new config self.limiter = TokenBucket( redis_client=self.limiter.redis, capacity=capacity, refill_rate=refill_rate, refill_interval=refill_interval ) response = {'success': True, 'config': self.config} except (ValueError, KeyError) as e: response = {'success': False, 'error': str(e)} self.send_response(200) self.send_header('Content-type', 'application/json') self.end_headers() self.wfile.write(json.dumps(response).encode()) def _get_html_page(self): """Generate the HTML page for the demo.""" return f""" Token Bucket Rate Limiter Demo

Token Bucket Rate Limiter Demo

How it works: The token bucket starts with a capacity of tokens. Each request consumes one token. Tokens refill at a constant rate. When the bucket is empty, requests are denied until tokens refill.

Configuration

Current settings:

Capacity: {self.config['capacity']} tokens
Refill Rate: {self.config['refill_rate']} tokens
Refill Interval: {self.config['refill_interval']} seconds

Test Rate Limiting

Click the button below to submit a request and see if it's allowed or denied.

""" def log_message(self, format, *args): """Override to customize logging.""" # Simple logging format print(f"[{self.log_date_time_string()}] {format % args}") def main(): """Main entry point for the demo server.""" parser = argparse.ArgumentParser(description='Token Bucket Rate Limiter Demo Server') parser.add_argument('--port', type=int, default=8080, help='Port to run the server on (default: 8080)') parser.add_argument('--redis-host', default='localhost', help='Redis host (default: localhost)') parser.add_argument('--redis-port', type=int, default=6379, help='Redis port (default: 6379)') args = parser.parse_args() # Initialize Redis connection try: redis_client = redis.Redis( host=args.redis_host, port=args.redis_port, decode_responses=True ) # Test connection redis_client.ping() print(f"✓ Connected to Redis at {args.redis_host}:{args.redis_port}") except redis.ConnectionError as e: print(f"✗ Failed to connect to Redis: {e}") print(f" Make sure Redis is running at {args.redis_host}:{args.redis_port}") sys.exit(1) # Initialize the rate limiter RateLimiterHandler.limiter = TokenBucket( redis_client=redis_client, capacity=RateLimiterHandler.config['capacity'], refill_rate=RateLimiterHandler.config['refill_rate'], refill_interval=RateLimiterHandler.config['refill_interval'] ) # Start the server server = HTTPServer(('', args.port), RateLimiterHandler) print(f"✓ Server started at http://localhost:{args.port}") print(f" Open your browser and visit http://localhost:{args.port}") print(f" Press Ctrl+C to stop the server") try: server.serve_forever() except KeyboardInterrupt: print("\n✓ Server stopped") sys.exit(0) if __name__ == '__main__': main()