| 66 | |
| 67 | |
| 68 | class IPRateLimit: |
| 69 | def __init__(self, count: int, minutes: int): |
| 70 | self.ips: Dict[str, Dict[str, Union[int, datetime]]] = {} |
| 71 | self.count = count |
| 72 | self.minutes = minutes |
| 73 | |
| 74 | def check_ip(self, ip: str) -> bool: |
| 75 | if ip in self.ips: |
| 76 | ip_info = self.ips[ip] |
| 77 | if ip_info["count"] >= self.count: |
| 78 | if ip_info["time"] + timedelta(minutes=self.minutes) > datetime.now(): |
| 79 | return False |
| 80 | self.ips.pop(ip) |
| 81 | return True |
| 82 | |
| 83 | def add_ip(self, ip: str) -> int: |
| 84 | ip_info = self.ips.get(ip, {"count": 0, "time": datetime.now()}) |
| 85 | ip_info["count"] += 1 |
| 86 | ip_info["time"] = datetime.now() |
| 87 | self.ips[ip] = ip_info |
| 88 | return ip_info["count"] |
| 89 | |
| 90 | async def remove_expired_ip(self) -> None: |
| 91 | now = datetime.now() |
| 92 | expiration = timedelta(minutes=self.minutes) |
| 93 | self.ips = { |
| 94 | ip: info |
| 95 | for ip, info in self.ips.items() |
| 96 | if info["time"] + expiration >= now |
| 97 | } |
| 98 | |
| 99 | def __call__(self, request: Request) -> str: |
| 100 | ip = get_client_ip(request) |
| 101 | if not self.check_ip(ip): |
| 102 | raise HTTPException(status_code=423, detail="请求次数过多，请稍后再试") |
| 103 | return ip |