refact: some codes
This commit is contained in:
@@ -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
|
||||
@@ -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
@@ -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
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user