update:新增限制时间内上传次数，移除IP黑名单列表过期IP，防止炸内存

Author: vastsa

depends.py

@@ -13,25 +13,36 @@ async def admin_required(pwd: Union[str, None] = Header(default=None)):
 
 class IPRateLimit:
     ips = {}
+    count = 0
+    minutes = 0
+
+    def __init__(self, count, minutes):
+        self.count = count
+        self.minutes = minutes
 
     def check_ip(self, ip):
         # 检查ip是否被禁止
         if ip in self.ips:
-            if self.ips[ip]['count'] >= settings.ERROR_COUNT:
-                if self.ips[ip]['time'] + timedelta(minutes=settings.ERROR_MINUTE) > datetime.now():
+            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(cls, ip):
         ip_info = cls.ips.get(ip, {'count': 0, 'time': datetime.now()})
         ip_info['count'] += 1
         cls.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.client.host
         if not self.check_ip(ip):
-            raise HTTPException(status_code=400, detail="错误次数过多，请稍后再试")
+            raise HTTPException(status_code=400, detail=f"请求次数过多，请稍后再试")
         return ip

main.py

@@ -11,58 +11,66 @@
 from sqlalchemy.ext.asyncio.session import AsyncSession
 
 import settings
-from utils import delete_expire_files, storage, get_code
+from utils import delete_expire_files, storage, get_code, error_ip_limit, upload_ip_limit
 from database import get_session, Codes, init_models
-from depends import admin_required, IPRateLimit
+from depends import admin_required
 
+# 实例化FastAPI
 app = FastAPI(debug=settings.DEBUG)
 
+# 数据存储文件夹
 DATA_ROOT = Path(settings.DATA_ROOT)
 if not DATA_ROOT.exists():
     DATA_ROOT.mkdir(parents=True)
-
+# 静态文件夹
 app.mount(settings.STATIC_URL, StaticFiles(directory=DATA_ROOT), name="static")
 
 
 @app.on_event('startup')
 async def startup():
+    # 初始化数据库
     await init_models()
+    # 启动后台任务，不定时删除过期文件
     asyncio.create_task(delete_expire_files())
 
 
+# 首页页面
 index_html = open('templates/index.html', 'r', encoding='utf-8').read() \
     .replace('{{title}}', settings.TITLE) \
     .replace('{{description}}', settings.DESCRIPTION) \
     .replace('{{keywords}}', settings.KEYWORDS)
+# 管理页面
 admin_html = open('templates/admin.html', 'r', encoding='utf-8').read() \
     .replace('{{title}}', settings.TITLE) \
     .replace('{{description}}', settings.DESCRIPTION) \
     .replace('{{keywords}}', settings.KEYWORDS)
 
-ip_limit = IPRateLimit()
-
 
-@app.get(f'/{settings.ADMIN_ADDRESS}')
+@app.get(f'/{settings.ADMIN_ADDRESS}', description='管理页面', response_class=HTMLResponse)
 async def admin():
     return HTMLResponse(admin_html)
 
 
-@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
+@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)], description='查询数据库列表')
 async def admin_post(s: AsyncSession = Depends(get_session)):
-    query = select(Codes)
-    codes = (await s.execute(query)).scalars().all()
+    # 查询数据库列表
+    codes = (await s.execute(select(Codes))).scalars().all()
     return {'detail': '查询成功', 'data': codes}
 
 
-@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)])
+@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)], description='删除数据库记录')
 async def admin_delete(code: str, s: AsyncSession = Depends(get_session)):
+    # 找到相应记录
     query = select(Codes).where(Codes.code == code)
+    # 找到第一条记录
     file = (await s.execute(query)).scalars().first()
-    if file:
-        if file.type != 'text':
-            await storage.delete_file(file.text)
-        await s.delete(file)
-        await s.commit()
+    # 如果记录存在，并且不是文本
+    if file and file.type != 'text':
+        # 删除文件
+        await storage.delete_file(file.text)
+    # 删除数据库记录
+    await s.delete(file)
+    await s.commit()
     return {'detail': '删除成功'}
 
 
@@ -72,25 +80,30 @@ async def index():
 
 
 @app.get('/select')
-async def get_file(code: str, s: AsyncSession = Depends(get_session)):
+async def get_file(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)):
+    # 查出数据库记录
     query = select(Codes).where(Codes.code == code)
     info = (await s.execute(query)).scalars().first()
+    # 如果记录不存在，IP错误次数+1
     if not info:
-        raise HTTPException(status_code=404, detail="口令不存在")
+        error_ip_limit.add_ip(ip)
+        raise HTTPException(status_code=404, detail="口令不存在，次数过多将被禁止访问")
+    # 如果是文本，直接返回
     if info.type == 'text':
         return {'detail': '查询成功', 'data': info.text}
+    # 如果是文件，返回文件
     else:
         filepath = await storage.get_filepath(info.text)
         return FileResponse(filepath, filename=info.name)
 
 
 @app.post('/')
-async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)):
+async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)):
     query = select(Codes).where(Codes.code == code)
     info = (await s.execute(query)).scalars().first()
     if not info:
-        error_count = settings.ERROR_COUNT - ip_limit.add_ip(ip)
-        raise HTTPException(status_code=404, detail=f"取件码错误，错误{error_count}次将被禁止10分钟")
+        error_count = settings.ERROR_COUNT - error_ip_limit.add_ip(ip)
+        raise HTTPException(status_code=404, detail=f"取件码错误，{error_count}次后将被禁止{settings.ERROR_MINUTE}分钟")
     if info.exp_time < datetime.datetime.now() or info.count == 0:
         if info.type != "text":
             await storage.delete_file(info.text)
@@ -109,7 +122,7 @@ async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depend
 
 @app.post('/share')
 async def share(background_tasks: BackgroundTasks, text: str = Form(default=None), style: str = Form(default='2'),
-                value: int = Form(default=1), file: UploadFile = File(default=None),
+                value: int = Form(default=1), file: UploadFile = File(default=None), ip: str = Depends(upload_ip_limit),
                 s: AsyncSession = Depends(get_session)):
     code = await get_code(s)
     if style == '2':
@@ -137,6 +150,7 @@ async def share(background_tasks: BackgroundTasks, text: str = Form(default=None
     info = Codes(code=code, text=_text, size=size, type=_type, name=name, count=exp_count, exp_time=exp_time, key=key)
     s.add(info)
     await s.commit()
+    upload_ip_limit.add_ip(ip)
     return {
         'detail': '分享成功，请点击取件码按钮查看上传列表',
         'data': {'code': code, 'key': key, 'name': name, 'text': _text}

settings.py

@@ -1,31 +1,36 @@
 from starlette.config import Config
 
+# 配置文件.env
 config = Config(".env")
-
+# 是否开启DEBUG模式
 DEBUG = config('DEBUG', cast=bool, default=False)
-
+# 端口
 PORT = config('PORT', cast=int, default=12345)
-
+# Sqlite数据库文件
 DATABASE_URL = config('DATABASE_URL', cast=str, default="sqlite+aiosqlite:///database.db")
-
+# 静态文件夹
 DATA_ROOT = config('DATA_ROOT', cast=str, default="./static")
-
+# 静态文件夹URL
 STATIC_URL = config('STATIC_URL', cast=str, default="/static")
-
+# 错误次数
 ERROR_COUNT = config('ERROR_COUNT', cast=int, default=5)
-
+# 错误限制分钟数
 ERROR_MINUTE = config('ERROR_MINUTE', cast=int, default=10)
-
+# 上传次数
+UPLOAD_COUNT = config('UPLOAD_COUNT', cast=int, default=60)
+# 上传限制分钟数
+UPLOAD_MINUTE = config('UPLOAD_MINUTE', cast=int, default=1)
+# 管理地址
 ADMIN_ADDRESS = config('ADMIN_ADDRESS', cast=str, default="admin")
-
+# 管理密码
 ADMIN_PASSWORD = config('ADMIN_PASSWORD', cast=str, default="admin")
-
+# 文件大小限制，默认10MB
 FILE_SIZE_LIMIT = config('FILE_SIZE_LIMIT', cast=int, default=10) * 1024 * 1024
-
+# 网站标题
 TITLE = config('TITLE', cast=str, default="文件快递柜")
-
+# 网站描述
 DESCRIPTION = config('DESCRIPTION', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件")
-
+# 网站关键词
 KEYWORDS = config('KEYWORDS', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件")
-
+# 存储引擎
 STORAGE_ENGINE = config('STORAGE_ENGINE', cast=str, default="filesystem")

utils.py

@@ -5,14 +5,22 @@
 from sqlalchemy.ext.asyncio.session import AsyncSession
 import settings
 from database import Codes, engine
+from depends import IPRateLimit
 from storage import STORAGE_ENGINE
 
 storage = STORAGE_ENGINE[settings.STORAGE_ENGINE]()
 
+# 错误IP限制器
+error_ip_limit = IPRateLimit(settings.ERROR_COUNT, settings.ERROR_MINUTE)
+# 上传文件限制器
+upload_ip_limit = IPRateLimit(settings.UPLOAD_COUNT, settings.UPLOAD_MINUTE)
+
 
 async def delete_expire_files():
     while True:
         async with AsyncSession(engine, expire_on_commit=False) as s:
+            await error_ip_limit.remove_expired_ip()
+            await upload_ip_limit.remove_expired_ip()
             query = select(Codes).where(or_(Codes.exp_time < datetime.datetime.now(), Codes.count == 0))
             exps = (await s.execute(query)).scalars().all()
             files = []
@@ -25,7 +33,7 @@ async def delete_expire_files():
             query = delete(Codes).where(Codes.id.in_(exps_ids))
             await s.execute(query)
             await s.commit()
-        await asyncio.sleep(random.randint(60, 300))
+        await asyncio.sleep(random.randint(2, 2))
 
 
 async def get_code(s: AsyncSession):