refact: some codes

This commit is contained in:
Lan
2024-10-06 17:14:16 +08:00
parent fd1eea7779
commit cc3a521d31
12 changed files with 370 additions and 357 deletions
+37
View File
@@ -0,0 +1,37 @@
from typing import Dict, Union
from datetime import datetime, timedelta
from fastapi import HTTPException, Request
class IPRateLimit:
def __init__(self, count: int, minutes: int):
self.ips: Dict[str, Dict[str, Union[int, datetime]]] = {}
self.count = count
self.minutes = minutes
def check_ip(self, ip: str) -> bool:
if ip in self.ips:
ip_info = self.ips[ip]
if ip_info['count'] >= self.count:
if ip_info['time'] + timedelta(minutes=self.minutes) > datetime.now():
return False
self.ips.pop(ip)
return True
def add_ip(self, ip: str) -> int:
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) -> None:
now = datetime.now()
expiration = timedelta(minutes=self.minutes)
self.ips = {ip: info for ip, info in self.ips.items() if info['time'] + expiration >= now}
def __call__(self, request: Request) -> str:
ip = request.headers.get('X-Real-IP') or request.headers.get('X-Forwarded-For') or request.client.host
if not self.check_ip(ip):
raise HTTPException(status_code=423, detail="请求次数过多,请稍后再试")
return ip
-45
View File
@@ -1,45 +0,0 @@
# @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
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 int(self.ips[ip]['count']) >= int(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
+38 -56
View File
@@ -1,85 +1,67 @@
# @Time : 2023/8/14 01:10
# @Author : Lan
# @File : utils.py
# @Software: PyCharm
import datetime
import uuid
import os
from fastapi import UploadFile, HTTPException
from typing import Tuple, Optional
from apps.base.depends import IPRateLimit
from apps.base.dependencies import IPRateLimit
from apps.base.models import FileCodes
from core.settings import settings
from core.utils import get_random_num, get_random_string, max_save_times_desc
async def get_file_path_name(file: UploadFile):
"""
获取文件路径和文件名
:param file:
:return: {
'path': 'share/data/2021/08/13',
'suffix': '.jpg',
'prefix': 'test',
'file_uuid': '44a83bbd70e04c8aa7fd93bfd8c88249',
'uuid_file_name': '44a83bbd70e04c8aa7fd93bfd8c88249.jpg',
'save_path': 'share/data/2021/08/13/44a83bbd70e04c8aa7fd93bfd8c88249.jpg'
}
"""
async def get_file_path_name(file: UploadFile) -> Tuple[str, str, str, str, str]:
"""获取文件路径和文件名"""
today = datetime.datetime.now()
path = f"share/data/{today.strftime('%Y/%m/%d')}"
prefix, suffix = os.path.splitext(file.filename)
file_uuid = f"{uuid.uuid4().hex}"
file_uuid = uuid.uuid4().hex
uuid_file_name = f"{file_uuid}{suffix}"
save_path = f"{path}/{uuid_file_name}"
return path, suffix, prefix, uuid_file_name, save_path
async def get_expire_info(expire_value: int, expire_style: str):
"""
获取过期信息
:param expire_value:
:param expire_style:
:return: expired_at 过期时间, expired_count 可用次数, used_count 已用次数, code 随机码
"""
expired_count, used_count, now, code = -1, 0, datetime.datetime.now(), None
if int(settings.max_save_seconds) > 0:
max_timedelta = datetime.timedelta(seconds=settings.max_save_seconds)
detail = await max_save_times_desc(settings.max_save_seconds)
detail = f'限制最长时间为 {detail[0]},可换用其他方式'
else:
max_timedelta = datetime.timedelta(days=7)
detail = '限制最长时间为 7天,可换用其他方式'
if expire_style == 'day':
if datetime.timedelta(days=expire_value) > max_timedelta:
async def get_expire_info(expire_value: int, expire_style: str) -> Tuple[Optional[datetime.datetime], int, int, str]:
"""获取过期信息"""
expired_count, used_count = -1, 0
now = datetime.datetime.now()
code = None
max_timedelta = datetime.timedelta(seconds=settings.max_save_seconds) if settings.max_save_seconds > 0 else datetime.timedelta(days=7)
detail = await max_save_times_desc(settings.max_save_seconds) if settings.max_save_seconds > 0 else '7天'
detail = f'限制最长时间为 {detail[0]},可换用其他方式'
expire_styles = {
'day': lambda: now + datetime.timedelta(days=expire_value),
'hour': lambda: now + datetime.timedelta(hours=expire_value),
'minute': lambda: now + datetime.timedelta(minutes=expire_value),
'count': lambda: (now + datetime.timedelta(days=1), expire_value),
'forever': lambda: (None, None), # 修改这里
}
if expire_style in expire_styles:
result = expire_styles[expire_style]()
if isinstance(result, tuple):
expired_at, extra = result
if expire_style == 'count':
expired_count = extra
elif expire_style == 'forever':
code = await get_random_code(style='string') # 移动到这里
else:
expired_at = result
if expired_at and expired_at - now > max_timedelta:
raise HTTPException(status_code=403, detail=detail)
expired_at = now + datetime.timedelta(days=expire_value)
elif expire_style == 'hour':
if datetime.timedelta(hours=expire_value) > max_timedelta:
raise HTTPException(status_code=403, detail=detail)
expired_at = now + datetime.timedelta(hours=expire_value)
elif expire_style == 'minute':
if datetime.timedelta(minutes=expire_value) > max_timedelta:
raise HTTPException(status_code=403, detail=detail)
expired_at = now + datetime.timedelta(minutes=expire_value)
elif expire_style == 'count':
expired_at = now + datetime.timedelta(days=1)
expired_count = expire_value
elif expire_style == 'forever':
expired_at = None
code = await get_random_code(style='string')
else:
expired_at = now + datetime.timedelta(days=1)
if not code:
code = await get_random_code()
return expired_at, expired_count, used_count, code
async def get_random_code(style='num'):
"""
获取随机字符串
:return:
"""
async def get_random_code(style='num') -> str:
"""获取随机字符串"""
while True:
code = await get_random_num() if style == 'num' else await get_random_string()
if not await FileCodes.filter(code=code).exists():
+53 -86
View File
@@ -1,38 +1,40 @@
# @Time : 2023/8/14 03:59
# @Author : Lan
# @File : views.py
# @Software: PyCharm
# 导入所需的库和模块
from fastapi import APIRouter, Form, UploadFile, File, Depends, HTTPException
from apps.admin.depends import admin_required
from apps.admin.dependencies import admin_required
from apps.base.models import FileCodes
from apps.base.pydantics import SelectFileModel
from apps.base.schemas import SelectFileModel
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
from core.utils import get_select_token
# 创建一个API路由
share_api = APIRouter(
prefix='/share', # 路由前缀
tags=['分享'], # 标签
)
share_api = APIRouter(prefix='/share', tags=['分享'])
async def validate_file_size(file: UploadFile, max_size: int):
if file.size > max_size:
max_size_mb = max_size / (1024 * 1024)
raise HTTPException(status_code=403, detail=f'大小超过限制,最大为{max_size_mb:.2f} MB')
async def create_file_code(code, **kwargs):
return await FileCodes.create(code=code, **kwargs)
# 分享文本的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(ip_limit['upload'])):
# 获取大小
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'])
):
text_size = len(text.encode('utf-8'))
# 限制 222KB
max_txt_size = 222 * 1024 # 转换为字节
max_txt_size = 222 * 1024
if text_size > max_txt_size:
raise HTTPException(status_code=403, detail=f'内容过多建议采用文件形式')
# 获取过期信息
raise HTTPException(status_code=403, detail='内容过多,建议采用文件形式')
expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style)
# 创建一个新的FileCodes实例
await FileCodes.create(
await create_file_code(
code=code,
text=text,
expired_at=expired_at,
@@ -41,33 +43,29 @@ async def share_text(text: str = Form(...), expire_value: int = Form(default=1,
size=len(text),
prefix='文本分享'
)
# 添加IP到限制列表
ip_limit['upload'].add_ip(ip)
# 返回API响应
return APIResponse(detail={
'code': code,
})
return APIResponse(detail={'code': code})
# 分享文件的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(ip_limit['upload'])):
if file.size > settings.uploadSize:
# 转换为 MB 并格式化输出
max_size_mb = settings.uploadSize / (1024 * 1024)
raise HTTPException(status_code=403, detail=f'大小超过限制,最大为{max_size_mb:.2f} MB')
# 获取过期信息
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'])
):
await validate_file_size(file, settings.uploadSize)
if expire_style not in settings.expireStyle:
raise HTTPException(status_code=400, detail='过期时间类型错误')
expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style)
# 获取文件路径和名称
path, suffix, prefix, uuid_file_name, save_path = await get_file_path_name(file)
# 保存文件
file_storage: FileStorageInterface = storages[settings.file_storage]()
await file_storage.save_file(file, save_path)
# 创建一个新的FileCodes实例
await FileCodes.create(
await create_file_code(
code=code,
prefix=prefix,
suffix=suffix,
@@ -78,69 +76,47 @@ async def share_file(expire_value: int = Form(default=1, gt=0), expire_style: st
expired_count=expired_count,
used_count=used_count,
)
# 添加IP到限制列表
ip_limit['upload'].add_ip(ip)
# 返回API响应
return APIResponse(detail={
'code': code,
'name': file.filename,
})
return APIResponse(detail={'code': code, 'name': file.filename})
# 根据code获取文件
async def get_code_file_by_code(code, check=True):
# 查询文件
file_code = await FileCodes.filter(code=code).first()
# 检查文件是否存在
if not file_code:
return False, '文件不存在'
# 检查文件是否过期
if await file_code.is_expired() and check:
return False, '文件已过期',
return False, '文件已过期'
return True, file_code
# 获取文件的API
async def update_file_usage(file_code):
file_code.used_count += 1
if file_code.expired_count > 0:
file_code.expired_count -= 1
await file_code.save()
@share_api.get('/select/')
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到限制列表
ip_limit['error'].add_ip(ip)
# 返回API响应
return APIResponse(code=404, detail=file_code)
# 更新文件的使用次数和过期次数
file_code.used_count += 1
if file_code.expired_count > 0:
file_code.expired_count -= 1
# 保存文件
await file_code.save()
# 返回文件响应
await update_file_usage(file_code)
return await file_storage.get_file_response(file_code)
# 选择文件的API
@share_api.post('/select/')
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到限制列表
ip_limit['error'].add_ip(ip)
# 返回API响应
return APIResponse(code=404, detail=file_code)
# 更新文件的使用次数和过期次数
file_code.used_count += 1
if file_code.expired_count > 0:
file_code.expired_count -= 1
# 保存文件
await file_code.save()
# 返回API响应
await update_file_usage(file_code)
return APIResponse(detail={
'code': file_code.code,
'name': file_code.prefix + file_code.suffix,
@@ -149,23 +125,14 @@ async def select_file(data: SelectFileModel, ip: str = Depends(ip_limit['error']
})
# 下载文件的API
@share_api.get('/download')
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到限制列表
if await get_select_token(code) != key:
ip_limit['error'].add_ip(ip)
# 获取文件
has, file_code = await get_code_file_by_code(code, False)
# 检查文件是否存在
if not has:
# 返回API响应
return APIResponse(code=404, detail='文件不存在')
# 如果文件是文本,返回文本内容,否则返回文件响应
if file_code.text:
return APIResponse(detail=file_code.text)
else:
return await file_storage.get_file_response(file_code)
return APIResponse(detail=file_code.text) if file_code.text else await file_storage.get_file_response(file_code)