Trixy7/trixy_core/network/security/rate_limiter.py

# -*- coding: utf-8 -*-
"""
Rate Limiting für Netzwerk-Verbindungen.

Bietet Token-Bucket-basiertes Rate-Limiting pro IP und global.
"""

import asyncio
import time
from dataclasses import dataclass, field
from enum import IntEnum, auto
from typing import Any


class RateLimitResult(IntEnum):
    """Ergebnis einer Rate-Limit-Prüfung."""
    ALLOWED = auto()       # Anfrage erlaubt
    LIMITED = auto()       # Rate-Limit erreicht
    BLOCKED = auto()       # Dauerhaft blockiert (zu viele Überschreitungen)


@dataclass
class RateLimitConfig:
    """
    Konfiguration für Rate-Limiting.

    Attributes:
        requests_per_second: Maximale Anfragen pro Sekunde
        burst_size: Maximale Burst-Größe (Token-Bucket-Kapazität)
        block_threshold: Anzahl Überschreitungen bis zur Blockierung
        block_duration: Blockierungsdauer in Sekunden
        cleanup_interval: Intervall für Speicherbereinigung in Sekunden
    """
    requests_per_second: float = 10.0
    burst_size: int = 20
    block_threshold: int = 100
    block_duration: float = 300.0  # 5 Minuten
    cleanup_interval: float = 60.0  # 1 Minute


@dataclass
class TokenBucket:
    """
    Token-Bucket für Rate-Limiting.

    Der Bucket füllt sich kontinuierlich mit Tokens.
    Jede Anfrage verbraucht ein Token.
    """
    capacity: int
    tokens: float
    refill_rate: float  # Tokens pro Sekunde
    last_update: float = field(default_factory=time.monotonic)
    violations: int = 0
    blocked_until: float | None = None

    def refill(self) -> None:
        """Füllt den Bucket basierend auf verstrichener Zeit."""
        now = time.monotonic()
        elapsed = now - self.last_update
        self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
        self.last_update = now

    def consume(self, tokens: int = 1) -> bool:
        """
        Versucht, Tokens zu verbrauchen.

        Args:
            tokens: Anzahl zu verbrauchender Tokens

        Returns:
            True wenn erfolgreich, False wenn nicht genug Tokens
        """
        self.refill()

        if self.tokens >= tokens:
            self.tokens -= tokens
            return True

        self.violations += 1
        return False

    def is_blocked(self) -> bool:
        """Prüft, ob der Bucket blockiert ist."""
        if self.blocked_until is None:
            return False
        if time.monotonic() > self.blocked_until:
            self.blocked_until = None
            return False
        return True

    def block(self, duration: float) -> None:
        """Blockiert den Bucket für eine bestimmte Zeit."""
        self.blocked_until = time.monotonic() + duration


class IPRateLimiter:
    """
    Rate-Limiter pro IP-Adresse.

    Verfolgt Anfragen pro IP und limitiert übermäßige Anfragen.

    Example:
        limiter = IPRateLimiter(RateLimitConfig(
            requests_per_second=10,
            burst_size=20
        ))

        result = limiter.check("192.168.1.100")
        if result == RateLimitResult.ALLOWED:
            # Anfrage verarbeiten
            pass
        elif result == RateLimitResult.LIMITED:
            # 429 Too Many Requests senden
            pass
        elif result == RateLimitResult.BLOCKED:
            # Verbindung abweisen
            pass
    """

    def __init__(self, config: RateLimitConfig | None = None) -> None:
        """
        Initialisiert den IP-Rate-Limiter.

        Args:
            config: Rate-Limit-Konfiguration
        """
        self._config = config or RateLimitConfig()
        self._buckets: dict[str, TokenBucket] = {}
        self._lock = asyncio.Lock()
        self._cleanup_task: asyncio.Task | None = None

    def _get_or_create_bucket(self, ip: str) -> TokenBucket:
        """Gibt den Bucket für eine IP zurück oder erstellt einen neuen."""
        if ip not in self._buckets:
            self._buckets[ip] = TokenBucket(
                capacity=self._config.burst_size,
                tokens=float(self._config.burst_size),
                refill_rate=self._config.requests_per_second
            )
        return self._buckets[ip]

    def check(self, ip: str, tokens: int = 1) -> RateLimitResult:
        """
        Prüft und verbraucht Rate-Limit für eine IP.

        Args:
            ip: Die IP-Adresse
            tokens: Anzahl zu verbrauchender Tokens

        Returns:
            RateLimitResult
        """
        bucket = self._get_or_create_bucket(ip)

        # Prüfe Blockierung
        if bucket.is_blocked():
            return RateLimitResult.BLOCKED

        # Prüfe Rate-Limit
        if not bucket.consume(tokens):
            # Prüfe ob Blockierung nötig
            if bucket.violations >= self._config.block_threshold:
                bucket.block(self._config.block_duration)
                return RateLimitResult.BLOCKED

            return RateLimitResult.LIMITED

        return RateLimitResult.ALLOWED

    async def check_async(self, ip: str, tokens: int = 1) -> RateLimitResult:
        """Thread-sichere Prüfung mit Lock."""
        async with self._lock:
            return self.check(ip, tokens)

    def reset(self, ip: str) -> None:
        """Setzt das Rate-Limit für eine IP zurück."""
        if ip in self._buckets:
            del self._buckets[ip]

    def unblock(self, ip: str) -> bool:
        """
        Hebt die Blockierung einer IP auf.

        Setzt auch die Tokens auf volle Kapazität zurück,
        damit die IP wieder Anfragen stellen kann.

        Args:
            ip: Die IP-Adresse

        Returns:
            True wenn IP blockiert war und entsperrt wurde
        """
        if ip in self._buckets:
            bucket = self._buckets[ip]
            if bucket.is_blocked():
                bucket.blocked_until = None
                bucket.violations = 0
                # Tokens auf volle Kapazität zurücksetzen
                bucket.tokens = float(bucket.capacity)
                return True
        return False

    def get_status(self, ip: str) -> dict[str, Any]:
        """
        Gibt den Status für eine IP zurück.

        Args:
            ip: Die IP-Adresse

        Returns:
            Status-Dictionary
        """
        if ip not in self._buckets:
            return {
                "exists": False,
                "tokens": self._config.burst_size,
                "blocked": False,
                "violations": 0
            }

        bucket = self._buckets[ip]
        bucket.refill()

        return {
            "exists": True,
            "tokens": bucket.tokens,
            "capacity": bucket.capacity,
            "blocked": bucket.is_blocked(),
            "blocked_until": bucket.blocked_until,
            "violations": bucket.violations
        }

    def get_blocked_ips(self) -> list[str]:
        """Gibt alle blockierten IPs zurück."""
        return [
            ip for ip, bucket in self._buckets.items()
            if bucket.is_blocked()
        ]

    def cleanup(self) -> int:
        """
        Entfernt inaktive Buckets.

        Returns:
            Anzahl entfernter Buckets
        """
        now = time.monotonic()
        stale_threshold = 300.0  # 5 Minuten Inaktivität

        to_remove = []
        for ip, bucket in self._buckets.items():
            # Entferne nur wenn nicht blockiert und lange inaktiv
            if not bucket.is_blocked():
                if now - bucket.last_update > stale_threshold:
                    to_remove.append(ip)

        for ip in to_remove:
            del self._buckets[ip]

        return len(to_remove)

    async def start_cleanup_task(self) -> None:
        """Startet den automatischen Cleanup-Task."""
        if self._cleanup_task is not None:
            return

        async def cleanup_loop():
            while True:
                await asyncio.sleep(self._config.cleanup_interval)
                self.cleanup()

        self._cleanup_task = asyncio.create_task(cleanup_loop())

    async def stop_cleanup_task(self) -> None:
        """Stoppt den Cleanup-Task."""
        if self._cleanup_task is not None:
            self._cleanup_task.cancel()
            try:
                await self._cleanup_task
            except asyncio.CancelledError:
                pass
            self._cleanup_task = None

    def get_statistics(self) -> dict[str, Any]:
        """Gibt Statistiken zurück."""
        blocked = sum(1 for b in self._buckets.values() if b.is_blocked())
        total_violations = sum(b.violations for b in self._buckets.values())

        return {
            "tracked_ips": len(self._buckets),
            "blocked_ips": blocked,
            "total_violations": total_violations,
            "config": {
                "requests_per_second": self._config.requests_per_second,
                "burst_size": self._config.burst_size,
                "block_threshold": self._config.block_threshold,
                "block_duration": self._config.block_duration
            }
        }


class GlobalRateLimiter:
    """
    Globales Rate-Limit für alle Verbindungen.

    Schützt den Server vor Überlastung unabhängig von der Quelle.

    Example:
        limiter = GlobalRateLimiter(
            max_requests_per_second=1000,
            burst_size=2000
        )

        if limiter.check() == RateLimitResult.ALLOWED:
            # Anfrage verarbeiten
            pass
    """

    def __init__(
        self,
        max_requests_per_second: float = 1000.0,
        burst_size: int = 2000
    ) -> None:
        """
        Initialisiert den globalen Rate-Limiter.

        Args:
            max_requests_per_second: Maximale Anfragen pro Sekunde
            burst_size: Maximale Burst-Größe
        """
        self._bucket = TokenBucket(
            capacity=burst_size,
            tokens=float(burst_size),
            refill_rate=max_requests_per_second
        )
        self._lock = asyncio.Lock()
        self._total_requests: int = 0
        self._limited_requests: int = 0

    def check(self, tokens: int = 1) -> RateLimitResult:
        """
        Prüft das globale Rate-Limit.

        Args:
            tokens: Anzahl zu verbrauchender Tokens

        Returns:
            RateLimitResult
        """
        self._total_requests += 1

        if self._bucket.consume(tokens):
            return RateLimitResult.ALLOWED

        self._limited_requests += 1
        return RateLimitResult.LIMITED

    async def check_async(self, tokens: int = 1) -> RateLimitResult:
        """Thread-sichere Prüfung mit Lock."""
        async with self._lock:
            return self.check(tokens)

    def get_status(self) -> dict[str, Any]:
        """Gibt den aktuellen Status zurück."""
        self._bucket.refill()

        return {
            "tokens_available": self._bucket.tokens,
            "capacity": self._bucket.capacity,
            "refill_rate": self._bucket.refill_rate,
            "total_requests": self._total_requests,
            "limited_requests": self._limited_requests,
            "limit_ratio": self._limited_requests / max(1, self._total_requests)
        }


class ConnectionRateLimiter:
    """
    Kombinierter Rate-Limiter mit IP- und Global-Limits.

    Kombiniert IP-spezifische und globale Rate-Limits.

    Example:
        limiter = ConnectionRateLimiter(
            ip_config=RateLimitConfig(requests_per_second=10),
            global_rps=1000
        )

        result = await limiter.check("192.168.1.100")
        if result != RateLimitResult.ALLOWED:
            # Anfrage ablehnen
            return
    """

    def __init__(
        self,
        ip_config: RateLimitConfig | None = None,
        global_rps: float = 1000.0,
        global_burst: int = 2000
    ) -> None:
        """
        Initialisiert den kombinierten Rate-Limiter.

        Args:
            ip_config: Konfiguration für IP-basiertes Limiting
            global_rps: Globale Anfragen pro Sekunde
            global_burst: Globale Burst-Größe
        """
        self._ip_limiter = IPRateLimiter(ip_config)
        self._global_limiter = GlobalRateLimiter(global_rps, global_burst)

    @property
    def ip_limiter(self) -> IPRateLimiter:
        """IP-basierter Rate-Limiter."""
        return self._ip_limiter

    @property
    def global_limiter(self) -> GlobalRateLimiter:
        """Globaler Rate-Limiter."""
        return self._global_limiter

    def check(self, ip: str, tokens: int = 1) -> RateLimitResult:
        """
        Prüft beide Rate-Limits.

        Args:
            ip: Die IP-Adresse
            tokens: Anzahl Tokens

        Returns:
            Das restriktivere Ergebnis
        """
        # Globales Limit zuerst
        global_result = self._global_limiter.check(tokens)
        if global_result != RateLimitResult.ALLOWED:
            return global_result

        # Dann IP-Limit
        return self._ip_limiter.check(ip, tokens)

    async def check_async(self, ip: str, tokens: int = 1) -> RateLimitResult:
        """Thread-sichere Prüfung."""
        global_result = await self._global_limiter.check_async(tokens)
        if global_result != RateLimitResult.ALLOWED:
            return global_result

        return await self._ip_limiter.check_async(ip, tokens)

    async def start(self) -> None:
        """Startet Hintergrund-Tasks."""
        await self._ip_limiter.start_cleanup_task()

    async def stop(self) -> None:
        """Stoppt Hintergrund-Tasks."""
        await self._ip_limiter.stop_cleanup_task()

    def get_statistics(self) -> dict[str, Any]:
        """Gibt kombinierte Statistiken zurück."""
        return {
            "ip": self._ip_limiter.get_statistics(),
            "global": self._global_limiter.get_status()
        }