Apply PR aliasrobotics#388: add rate limiting to web logs endpoints

Author: giveen

tools/logs.py

@@ -19,22 +19,86 @@
 import matplotlib
 matplotlib.use('Agg')
 
-from flask import Flask, render_template
+from flask import Flask, render_template, request
+from functools import wraps
 import pandas as pd
 import matplotlib.pyplot as plt
 import io
 import base64
-from datetime import datetime
+from datetime import datetime, timedelta
 import os
 import folium
 import requests
 import argparse
 from typing import Dict, Optional
 import numpy as np
 import re
+from collections import defaultdict
+import time
 
 app = Flask(__name__)
 
+# Rate limiting configuration
+class RateLimiter:
+    """In-memory rate limiter to prevent DoS attacks on API endpoints."""
+
+    def __init__(self, requests_per_minute: int = 10, requests_per_hour: int = 50, requests_per_day: int = 200):
+        self.requests_per_minute = requests_per_minute
+        self.requests_per_hour = requests_per_hour
+        self.requests_per_day = requests_per_day
+        self.request_history: Dict[str, list] = defaultdict(list)
+
+    def is_rate_limited(self, client_id: str) -> bool:
+        """Return True if the client has exceeded configured rate limits."""
+        current_time = time.time()
+
+        # Clean up requests older than 24 hours
+        self.request_history[client_id] = [
+            req_time
+            for req_time in self.request_history[client_id]
+            if current_time - req_time < 86400
+        ]
+
+        # Minute window
+        minute_requests = [r for r in self.request_history[client_id] if current_time - r < 60]
+        if len(minute_requests) >= self.requests_per_minute:
+            return True
+
+        # Hour window
+        hour_requests = [r for r in self.request_history[client_id] if current_time - r < 3600]
+        if len(hour_requests) >= self.requests_per_hour:
+            return True
+
+        # Day window
+        day_requests = [r for r in self.request_history[client_id] if current_time - r < 86400]
+        if len(day_requests) >= self.requests_per_day:
+            return True
+
+        # Record this request and allow
+        self.request_history[client_id].append(current_time)
+        return False
+
+
+# Initialize the global rate limiter
+rate_limiter = RateLimiter(requests_per_minute=10, requests_per_hour=50, requests_per_day=200)
+
+
+def apply_rate_limit():
+    """Check request and return a Flask-style response tuple on rate limit, else None."""
+    try:
+        # Prefer X-Forwarded-For when behind proxies
+        forwarded = request.headers.get('X-Forwarded-For', None)
+        if forwarded:
+            client_ip = forwarded.split(',')[0].strip()
+        else:
+            client_ip = request.remote_addr or 'unknown'
+    except Exception:
+        client_ip = 'unknown'
+
+    if rate_limiter.is_rate_limited(client_ip):
+        return "Rate limit exceeded. Please try again later.", 429
+    return None
+
 # Configuration for enabled visualizations
 class Config:
     def __init__(self):
@@ -508,6 +572,10 @@ def create_pypi_plot():
 
 @app.route('/')
 def index():
+    # Apply rate limiting
+    rate_limit_response = apply_rate_limit()
+    if rate_limit_response:
+        return rate_limit_response
     # Get log file path from app config
     log_file = app.config['LOG_FILE']
 
@@ -544,6 +612,10 @@ def index():
 
 @app.route('/pypi-stats')
 def pypi_stats():
+    # Apply rate limiting
+    rate_limit_response = apply_rate_limit()
+    if rate_limit_response:
+        return rate_limit_response
     # Generate PyPI plot
     pypi_plot, stats = create_pypi_plot()