From b9f9ce6f8c911d82e6ddaac524db9923474a1717 Mon Sep 17 00:00:00 2001
From: lan <hhxk666@gmail.com>
Date: Tue, 15 Aug 2023 21:33:58 +0800
Subject: [PATCH] =?UTF-8?q?update:=E9=87=8D=E6=9E=84=E5=9F=BA=E6=9C=AC?=
 =?UTF-8?q?=E5=AE=8C=E6=88=90?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 apps/base/depends.py | 1 +
 apps/base/utils.py   | 5 +++--
 apps/base/views.py   | 8 +++++---
 core/settings.py     | 1 -
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/apps/base/depends.py b/apps/base/depends.py
index 324bf0f..b2cb926 100644
--- a/apps/base/depends.py
+++ b/apps/base/depends.py
@@ -12,6 +12,7 @@ from core.response import APIResponse
 
 class IPRateLimit:
     def __init__(self, count, minutes):
+        print(count, minutes)
         self.ips = {}
         self.count = count
         self.minutes = minutes
diff --git a/apps/base/utils.py b/apps/base/utils.py
index 64efb15..72155cb 100644
--- a/apps/base/utils.py
+++ b/apps/base/utils.py
@@ -9,6 +9,7 @@ from fastapi import UploadFile
 
 from apps.base.depends import IPRateLimit
 from apps.base.models import FileCodes
+from core.settings import settings
 from core.utils import get_random_num, get_random_string
 
 
@@ -73,6 +74,6 @@ async def get_random_code(style='num'):
 
 
 # 错误IP限制器
-error_ip_limit = IPRateLimit(1, 1)
+error_ip_limit = IPRateLimit(count=settings.errorCount, minutes=settings.errorMinute)
 # 上传文件限制器
-upload_ip_limit = IPRateLimit(10, 1)
+upload_ip_limit = IPRateLimit(count=settings.uploadCount, minutes=settings.errorMinute)
diff --git a/apps/base/views.py b/apps/base/views.py
index 789a528..6addb4e 100644
--- a/apps/base/views.py
+++ b/apps/base/views.py
@@ -8,7 +8,7 @@ from starlette.responses import FileResponse
 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
+from apps.base.utils import get_expire_info, get_file_path_name, error_ip_limit, upload_ip_limit
 from core.response import APIResponse
 from core.settings import settings
 from core.storage import file_storage
@@ -20,7 +20,7 @@ share_api = APIRouter(
 
 
 @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')):
+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)):
     expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style)
     await FileCodes.create(
         code=code,
@@ -31,13 +31,14 @@ async def share_text(text: str = Form(...), expire_value: int = Form(default=1,
         size=len(text),
         prefix='文本分享'
     )
+    upload_ip_limit.add_ip(ip)
     return APIResponse(detail={
         'code': code,
     })
 
 
 @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(...)):
+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)):
     if file.size > settings.uploadSize:
         raise HTTPException(status_code=403, detail=f'文件大小超过限制，最大为{settings.uploadSize}字节')
     expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style)
@@ -54,6 +55,7 @@ async def share_file(expire_value: int = Form(default=1, gt=0), expire_style: st
         expired_count=expired_count,
         used_count=used_count,
     )
+    upload_ip_limit.add_ip(ip)
     return APIResponse(detail={
         'code': code,
         'name': file.filename,
diff --git a/core/settings.py b/core/settings.py
index f309897..df8b7c5 100644
--- a/core/settings.py
+++ b/core/settings.py
@@ -44,7 +44,6 @@ class Settings:
                 if value.isdigit():
                     value = int(value)
                 default_value[key] = value
-
         # 更新self
         for key, value in default_value.items():
             self.__setattr__(key, value)