forked from vastsa/FileCodeBox
-
Notifications
You must be signed in to change notification settings - Fork 0
/
depends.py
49 lines (39 loc) · 1.75 KB
/
depends.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
from typing import Union
from datetime import datetime, timedelta
from fastapi import Header, HTTPException, Request
import settings
async def admin_required(pwd: Union[str, None] = Header(default=None), request: Request = None):
if 'share' in request.url.path:
if pwd != settings.ADMIN_PASSWORD and not settings.ENABLE_UPLOAD:
raise HTTPException(status_code=403, detail='本站上传功能已关闭,仅管理员可用')
else:
if not pwd or pwd != settings.ADMIN_PASSWORD:
raise HTTPException(status_code=401, detail="密码错误,请重新登录")
class IPRateLimit:
def __init__(self, count, minutes):
self.ips = {}
self.count = count
self.minutes = minutes
def check_ip(self, ip):
# 检查ip是否被禁止
if ip in self.ips:
if self.ips[ip]['count'] >= self.count:
if self.ips[ip]['time'] + timedelta(minutes=self.minutes) > datetime.now():
return False
else:
self.ips.pop(ip)
return True
def add_ip(self, ip):
ip_info = self.ips.get(ip, {'count': 0, 'time': datetime.now()})
ip_info['count'] += 1
self.ips[ip] = ip_info
return ip_info['count']
async def remove_expired_ip(self):
for ip in list(self.ips.keys()):
if self.ips[ip]['time'] + timedelta(minutes=self.minutes) < datetime.now():
self.ips.pop(ip)
def __call__(self, request: Request):
ip = request.headers.get('X-Real-IP', request.headers.get('X-Forwarded-For', request.client.host))
if not self.check_ip(ip):
raise HTTPException(status_code=400, detail=f"请求次数过多,请稍后再试")
return ip