Rename project to ddns-service

src/ddns_service/ratelimit.py

@@ -0,0 +1,129 @@
+"""Rate limiting with sliding window."""
+
+import threading
+import time
+from collections import defaultdict
+
+
+class RateLimiter:
+    """Sliding window rate limiter with separate good/bad request tracking."""
+
+    def __init__(self, config):
+        """
+        Initialize rate limiter from config.
+
+        Args:
+            config: Full configuration dictionary.
+        """
+        rl_config = config.get("rate_limit", {})
+        self.enabled = rl_config.get("enabled", False)
+        self.good_window = rl_config.get("good_window_seconds", 60)
+        self.good_max = rl_config.get("good_max_requests", 30)
+        self.bad_window = rl_config.get("bad_window_seconds", 60)
+        self.bad_max = rl_config.get("bad_max_requests", 5)
+
+        self.bad_requests = defaultdict(list)
+        self.good_requests = defaultdict(list)
+        self.bad_lock = threading.Lock()
+        self.good_lock = threading.Lock()
+
+    def _cleanup_old(self, timestamps, window):
+        """Remove timestamps older than window."""
+        cutoff = time.time() - window
+        return [t for t in timestamps if t > cutoff]
+
+    def is_blocked_bad(self, ip):
+        """
+        Check if IP is blocked by the bad request rate limiter with recording
+        when IP is already limited.
+
+        Args:
+            ip: Client IP address.
+
+        Returns:
+            Tuple of (blocked, retry_after_seconds).
+        """
+        if not self.enabled:
+            return False, 0
+
+        now = time.time()
+
+        with self.bad_lock:
+            # Check bad requests
+            self.bad_requests[ip] = self._cleanup_old(
+                self.bad_requests[ip], self.bad_window
+            )
+            if len(self.bad_requests[ip]) >= self.bad_max:
+                self.bad_requests[ip].append(time.time())
+                oldest = min(self.bad_requests[ip][-self.bad_max:])
+                retry_after = int(oldest + self.bad_window - now) + 1
+                return True, max(1, retry_after)
+
+        return False, 0
+
+    def is_blocked_good(self, ip):
+        """
+        Check if IP is blocked by the good request rate limiter without recording.
+
+        Args:
+            ip: Client IP address.
+
+        Returns:
+            Tuple of (blocked, retry_after_seconds).
+        """
+        if not self.enabled:
+            return False, 0
+
+        now = time.time()
+
+        with self.good_lock:
+            # Check good requests
+            self.good_requests[ip] = self._cleanup_old(
+                self.good_requests[ip], self.good_window
+            )
+            if len(self.good_requests[ip]) >= self.good_max:
+                oldest = min(self.good_requests[ip])
+                retry_after = int(oldest + self.good_window - now) + 1
+                return True, max(1, retry_after)
+
+        return False, 0
+
+    def record_bad(self, ip):
+        """Record a bad request (without checking)."""
+        if not self.enabled:
+            return
+
+        with self.bad_lock:
+            self.bad_requests[ip] = self._cleanup_old(
+                self.bad_requests[ip], self.bad_window
+            )
+            self.bad_requests[ip].append(time.time())
+
+    def record_good(self, ip):
+        """Record a good request (without checking)."""
+        if not self.enabled:
+            return
+
+        with self.good_lock:
+            self.good_requests[ip] = self._cleanup_old(
+                self.good_requests[ip], self.good_window
+            )
+            self.good_requests[ip].append(time.time())
+
+    def cleanup(self):
+        """Remove stale entries to prevent memory leak."""
+        with self.good_lock:
+            for ip in list(self.good_requests.keys()):
+                self.good_requests[ip] = self._cleanup_old(
+                    self.good_requests[ip], self.good_window
+                )
+                if not self.good_requests[ip]:
+                    del self.good_requests[ip]
+
+        with self.bad_lock:
+            for ip in list(self.bad_requests.keys()):
+                self.bad_requests[ip] = self._cleanup_old(
+                    self.bad_requests[ip], self.bad_window
+                )
+                if not self.bad_requests[ip]:
+                    del self.bad_requests[ip]