From fafc56bc322ea05dae235459743fb1c4fef29a99 Mon Sep 17 00:00:00 2001 From: lan Date: Sun, 28 Jul 2024 19:03:37 +0800 Subject: [PATCH] fix: issue #186 --- apps/base/utils.py | 8 ++++---- apps/base/views.py | 23 ++++++++++++----------- core/tasks.py | 6 +++--- main.py | 6 ++++++ 4 files changed, 25 insertions(+), 18 deletions(-) diff --git a/apps/base/utils.py b/apps/base/utils.py index 370eed7..b3851fd 100644 --- a/apps/base/utils.py +++ b/apps/base/utils.py @@ -86,7 +86,7 @@ async def get_random_code(style='num'): return code -# 错误IP限制器 -error_ip_limit = IPRateLimit(count=settings.errorCount, minutes=settings.errorMinute) -# 上传文件限制器 -upload_ip_limit = IPRateLimit(count=settings.uploadCount, minutes=settings.errorMinute) +ip_limit = { + 'error': IPRateLimit(count=settings.uploadCount, minutes=settings.errorMinute), + 'upload': IPRateLimit(count=settings.errorCount, minutes=settings.errorMinute) +} diff --git a/apps/base/views.py b/apps/base/views.py index 318a391..c838e2b 100644 --- a/apps/base/views.py +++ b/apps/base/views.py @@ -7,7 +7,7 @@ from fastapi import APIRouter, Form, UploadFile, File, Depends, HTTPException from apps.admin.depends import admin_required from apps.base.models import FileCodes from apps.base.pydantics import SelectFileModel -from apps.base.utils import get_expire_info, get_file_path_name, error_ip_limit, upload_ip_limit +from apps.base.utils import get_expire_info, get_file_path_name, ip_limit from core.response import APIResponse from core.settings import settings from core.storage import storages, FileStorageInterface @@ -22,7 +22,7 @@ share_api = APIRouter( # 分享文本的API @share_api.post('/text/', dependencies=[Depends(admin_required)]) -async def share_text(text: str = Form(...), expire_value: int = Form(default=1, gt=0), expire_style: str = Form(default='day'), ip: str = Depends(upload_ip_limit)): +async def share_text(text: str = Form(...), expire_value: int = Form(default=1, gt=0), expire_style: str = Form(default='day'), ip: str = Depends(ip_limit['upload'])): # 获取过期信息 expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style) # 创建一个新的FileCodes实例 @@ -36,7 +36,7 @@ async def share_text(text: str = Form(...), expire_value: int = Form(default=1, prefix='文本分享' ) # 添加IP到限制列表 - upload_ip_limit.add_ip(ip) + ip_limit['upload'].add_ip(ip) # 返回API响应 return APIResponse(detail={ 'code': code, @@ -45,7 +45,8 @@ async def share_text(text: str = Form(...), expire_value: int = Form(default=1, # 分享文件的API @share_api.post('/file/', dependencies=[Depends(admin_required)]) -async def share_file(expire_value: int = Form(default=1, gt=0), expire_style: str = Form(default='day'), file: UploadFile = File(...), ip: str = Depends(upload_ip_limit)): +async def share_file(expire_value: int = Form(default=1, gt=0), expire_style: str = Form(default='day'), file: UploadFile = File(...), + ip: str = Depends(ip_limit['upload'])): # 检查文件大小是否超过限制 if file.size > int(settings.uploadSize): raise HTTPException(status_code=403, detail=f'文件大小超过限制，最大为{settings.uploadSize}字节') @@ -71,7 +72,7 @@ async def share_file(expire_value: int = Form(default=1, gt=0), expire_style: st used_count=used_count, ) # 添加IP到限制列表 - upload_ip_limit.add_ip(ip) + ip_limit['upload'].add_ip(ip) # 返回API响应 return APIResponse(detail={ 'code': code, @@ -94,14 +95,14 @@ async def get_code_file_by_code(code, check=True): # 获取文件的API @share_api.get('/select/') -async def get_code_file(code: str, ip: str = Depends(error_ip_limit)): +async def get_code_file(code: str, ip: str = Depends(ip_limit['error'])): file_storage: FileStorageInterface = storages[settings.file_storage]() # 获取文件 has, file_code = await get_code_file_by_code(code) # 检查文件是否存在 if not has: # 添加IP到限制列表 - error_ip_limit.add_ip(ip) + ip_limit['error'].add_ip(ip) # 返回API响应 return APIResponse(code=404, detail=file_code) # 更新文件的使用次数和过期次数 @@ -116,14 +117,14 @@ async def get_code_file(code: str, ip: str = Depends(error_ip_limit)): # 选择文件的API @share_api.post('/select/') -async def select_file(data: SelectFileModel, ip: str = Depends(error_ip_limit)): +async def select_file(data: SelectFileModel, ip: str = Depends(ip_limit['error'])): file_storage: FileStorageInterface = storages[settings.file_storage]() # 获取文件 has, file_code = await get_code_file_by_code(data.code) # 检查文件是否存在 if not has: # 添加IP到限制列表 - error_ip_limit.add_ip(ip) + ip_limit['error'].add_ip(ip) # 返回API响应 return APIResponse(code=404, detail=file_code) # 更新文件的使用次数和过期次数 @@ -143,13 +144,13 @@ async def select_file(data: SelectFileModel, ip: str = Depends(error_ip_limit)): # 下载文件的API @share_api.get('/download') -async def download_file(key: str, code: str, ip: str = Depends(error_ip_limit)): +async def download_file(key: str, code: str, ip: str = Depends(ip_limit['error'])): file_storage: FileStorageInterface = storages[settings.file_storage]() # 检查token是否有效 is_valid = await get_select_token(code) == key if not is_valid: # 添加IP到限制列表 - error_ip_limit.add_ip(ip) + ip_limit['error'].add_ip(ip) # 获取文件 has, file_code = await get_code_file_by_code(code, False) # 检查文件是否存在 diff --git a/core/tasks.py b/core/tasks.py index fc6ff07..7a55fe3 100644 --- a/core/tasks.py +++ b/core/tasks.py @@ -7,7 +7,7 @@ import asyncio from tortoise.expressions import Q from apps.base.models import FileCodes -from apps.base.utils import error_ip_limit, upload_ip_limit +from apps.base.utils import ip_limit from core.settings import settings from core.storage import FileStorageInterface, storages from core.utils import get_now @@ -17,8 +17,8 @@ async def delete_expire_files(): file_storage: FileStorageInterface = storages[settings.file_storage]() while True: try: - await error_ip_limit.remove_expired_ip() - await upload_ip_limit.remove_expired_ip() + await ip_limit['error'].remove_expired_ip() + await ip_limit['upload'].remove_expired_ip() expire_data = await FileCodes.filter(Q(expired_at__lt=await get_now()) | Q(expired_count=0)).all() for exp in expire_data: await file_storage.delete_file(exp) diff --git a/main.py b/main.py index 5625083..3658500 100644 --- a/main.py +++ b/main.py @@ -10,7 +10,9 @@ from fastapi.responses import HTMLResponse from fastapi.staticfiles import StaticFiles from tortoise.contrib.fastapi import register_tortoise +from apps.base.depends import IPRateLimit from apps.base.models import KeyValue +from apps.base.utils import ip_limit from apps.base.views import share_api from apps.admin.views import admin_api from core.response import APIResponse @@ -59,6 +61,10 @@ async def startup_event(): # 读取用户配置 user_config, created = await KeyValue.get_or_create(key='settings', defaults={'value': DEFAULT_CONFIG}) settings.user_config = user_config.value + ip_limit['error'].minutes = settings.errorMinute + ip_limit['error'].count = settings.errorCount + ip_limit['upload'].minutes = settings.uploadMinute + ip_limit['upload'].count = settings.uploadCount @app.get('/')