add:fronted

apps/base/depends.py

@@ -0,0 +1,49 @@
+# @Time    : 2023/8/14 12:20
+# @Author  : Lan
+# @File    : depends.py
+# @Software: PyCharm
+from typing import Union
+from datetime import datetime, timedelta
+
+from fastapi import Header, HTTPException, Request
+
+from core.response import APIResponse
+
+
+async def admin_required(pwd: Union[str, None] = Header(default=None), request: Request = None):
+    return False
+
+
+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
+        ip_info['time'] = datetime.now()
+        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=423, detail=f"请求次数过多，请稍后再试")
+        return ip

apps/base/models.py

@@ -28,7 +28,7 @@ class FileCodes(Model):
     created_at: Optional[datetime] = fields.DatetimeField(auto_now_add=True, description='创建时间')
 
     async def is_expired(self):
-        if self.expired_at:
+        if self.expired_at and (self.expired_count == -1 or self.used_count < self.expired_count):
             return self.expired_at < await get_now()
         else:
             return self.expired_count != -1 and self.used_count >= self.expired_count

apps/base/pydantics.py

@@ -0,0 +1,5 @@
+from pydantic import BaseModel
+
+
+class SelectFileModel(BaseModel):
+    code: str

apps/base/utils.py

@@ -7,6 +7,7 @@ import uuid
 import os
 from fastapi import UploadFile
 
+from apps.base.depends import IPRateLimit
 from apps.base.models import FileCodes
 from core.utils import get_random_num, get_random_string
 
@@ -41,7 +42,6 @@ async def get_expire_info(expire_value: int, expire_style: str):
     :return: expired_at 过期时间, expired_count 可用次数, used_count 已用次数, code 随机码
     """
     expired_count, used_count, now, code = -1, 0, datetime.datetime.now(), None
-
     if expire_style == 'day':
         expired_at = now + datetime.timedelta(days=expire_value)
     elif expire_style == 'hour':
@@ -58,6 +58,7 @@ async def get_expire_info(expire_value: int, expire_style: str):
         expired_at = now + datetime.timedelta(days=1)
     if not code:
         code = await get_random_code()
+    print(expire_style, expire_value, expired_at, expired_count, used_count, code)
     return expired_at, expired_count, used_count, code
 
 
@@ -70,3 +71,9 @@ async def get_random_code(style='num'):
         code = await get_random_num() if style == 'num' else await get_random_string()
         if not await FileCodes.filter(code=code).exists():
             return code
+
+
+# 错误IP限制器
+error_ip_limit = IPRateLimit(1, 1)
+# 上传文件限制器
+upload_ip_limit = IPRateLimit(10, 1)

apps/base/views.py

@@ -2,11 +2,12 @@
 # @Author  : Lan
 # @File    : views.py
 # @Software: PyCharm
-from fastapi import APIRouter, Form, UploadFile, File
-from pydantic import BaseModel
+from fastapi import APIRouter, Form, UploadFile, File, Depends
 
 from apps.base.models import FileCodes
-from apps.base.utils import get_expire_info, get_file_path_name
+from apps.base.pydantics import SelectFileModel
+from apps.base.utils import get_expire_info, get_file_path_name, error_ip_limit
+from core.response import APIResponse
 from core.storage import file_storage
 
 share_api = APIRouter(
@@ -27,13 +28,9 @@ async def share_text(text: str = Form(...), expire_value: int = Form(default=1,
         size=len(text),
         prefix='文本分享'
     )
-    return {
-        'code': 200,
-        'msg': 'success',
-        'data': {
-            'code': code,
-        }
-    }
+    return APIResponse(detail={
+        'code': code,
+    })
 
 
 @share_api.post('/file/')
@@ -52,40 +49,25 @@ async def share_file(expire_value: int = Form(default=1, gt=0), expire_style: st
         expired_count=expired_count,
         used_count=used_count,
     )
-    return {
-        'code': 200,
-        'msg': 'success',
-        'data': {
-            'code': code,
-            'name': file.filename,
-        }
-    }
-
-
-class SelectFileModel(BaseModel):
-    code: str
+    return APIResponse(detail={
+        'code': code,
+        'name': file.filename,
+    })
 
 
 @share_api.post('/select/')
-async def select_file(data: SelectFileModel):
+async def select_file(data: SelectFileModel, ip: str = Depends(error_ip_limit)):
     file_code = await FileCodes.filter(code=data.code).first()
     if not file_code:
-        return {
-            'code': 404,
-            'msg': '文件不存在',
-        }
+        error_ip_limit.add_ip(ip)
+        return APIResponse(code=404, detail='文件不存在')
     if await file_code.is_expired():
-        return {
-            'code': 403,
-            'msg': '文件已过期',
-        }
-    return {
-        'code': 200,
-        'msg': 'success',
-        'data': {
-            'code': file_code.code,
-            'name': file_code.prefix + file_code.suffix,
-            'size': file_code.size,
-            'text': await file_storage.get_file_url(file_code),
-        }
-    }
+        return APIResponse(code=403, detail='文件已过期')
+    file_code.used_count += 1
+    await file_code.save()
+    return APIResponse(detail={
+        'code': file_code.code,
+        'name': file_code.prefix + file_code.suffix,
+        'size': file_code.size,
+        'text': await file_storage.get_file_url(file_code),
+    })