refact: some codes

This commit is contained in:
Lan
2024-10-06 17:14:16 +08:00
parent ef830b9f03
commit 2ca874e0e0
12 changed files with 370 additions and 357 deletions
@@ -7,9 +7,10 @@ from typing import Union
from fastapi import Header, HTTPException from fastapi import Header, HTTPException
from fastapi.requests import Request from fastapi.requests import Request
from core.settings import settings from core.settings import settings
from apps.admin.services import FileService, ConfigService, LocalFileService
async def admin_required(authorization: Union[str, None] = Header(default=None), request: Request = None): async def admin_required(authorization: str = Header(default=None), request: Request = None):
is_admin = authorization == str(settings.admin_token) is_admin = authorization == str(settings.admin_token)
if request.url.path.startswith('/share/'): if request.url.path.startswith('/share/'):
if not settings.openUpload and not is_admin: if not settings.openUpload and not is_admin:
@@ -17,3 +18,16 @@ async def admin_required(authorization: Union[str, None] = Header(default=None),
else: else:
if not is_admin: if not is_admin:
raise HTTPException(status_code=401, detail='未授权或授权校验失败') raise HTTPException(status_code=401, detail='未授权或授权校验失败')
return is_admin
async def get_file_service():
return FileService()
async def get_config_service():
return ConfigService()
async def get_local_file_service():
return LocalFileService()
-5
View File
@@ -1,5 +0,0 @@
from pydantic import BaseModel
class IDData(BaseModel):
id: int
+19
View File
@@ -0,0 +1,19 @@
from pydantic import BaseModel
class IDData(BaseModel):
id: int
class ConfigUpdateData(BaseModel):
admin_token: str
class ShareItem(BaseModel):
expire_value: int
expire_style: str = 'day'
filename: str
class DeleteItem(BaseModel):
filename: str
+123
View File
@@ -0,0 +1,123 @@
import os
import time
from core.response import APIResponse
from core.storage import FileStorageInterface, storages
from core.settings import settings
from apps.base.models import FileCodes, KeyValue
from apps.base.utils import get_expire_info, get_file_path_name
from fastapi import HTTPException
from core.settings import data_root
class FileService:
def __init__(self):
self.file_storage: FileStorageInterface = storages[settings.file_storage]()
async def delete_file(self, file_id: int):
file_code = await FileCodes.get(id=file_id)
await self.file_storage.delete_file(file_code)
await file_code.delete()
async def list_files(self, page: int, size: int):
offset = (page - 1) * size
files = await FileCodes.all().limit(size).offset(offset)
total = await FileCodes.all().count()
return files, total
async def download_file(self, file_id: int):
file_code = await FileCodes.filter(id=file_id).first()
if not file_code:
raise HTTPException(status_code=404, detail='文件不存在')
if file_code.text:
return APIResponse(detail=file_code.text)
else:
return await self.file_storage.get_file_response(file_code)
async def share_local_file(self, item):
local_file = LocalFileClass(item.filename)
if not await local_file.exists():
raise HTTPException(status_code=404, detail='文件不存在')
text = await local_file.read()
expired_at, expired_count, used_count, code = await get_expire_info(item.expire_value, item.expire_style)
path, suffix, prefix, uuid_file_name, save_path = await get_file_path_name(item)
await self.file_storage.save_file(text, save_path)
await FileCodes.create(
code=code,
prefix=prefix,
suffix=suffix,
uuid_file_name=uuid_file_name,
file_path=path,
size=local_file.size,
expired_at=expired_at,
expired_count=expired_count,
used_count=used_count,
)
return {
'code': code,
'name': local_file.file,
}
class ConfigService:
def get_config(self):
return settings.items()
async def update_config(self, data: dict):
admin_token = data.get('admin_token')
if admin_token is None or admin_token == '':
raise HTTPException(status_code=400, detail='管理员密码不能为空')
for key, value in data.items():
if key not in settings.default_config:
continue
if key in ['errorCount', 'errorMinute', 'max_save_seconds', 'onedrive_proxy', 'openUpload', 'port', 's3_proxy', 'uploadCount', 'uploadMinute', 'uploadSize']:
data[key] = int(value)
elif key in ['opacity']:
data[key] = float(value)
else:
data[key] = value
await KeyValue.filter(key='settings').update(value=data)
for k, v in data.items():
settings.__setattr__(k, v)
class LocalFileService:
async def list_files(self):
files = []
for file in os.listdir(data_root / 'local'):
files.append(LocalFileClass(file))
return files
async def delete_file(self, filename: str):
file = LocalFileClass(filename)
if await file.exists():
await file.delete()
return '删除成功'
raise HTTPException(status_code=404, detail='文件不存在')
class LocalFileClass:
def __init__(self, file):
self.file = file
self.path = data_root / 'local' / file
self.ctime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(os.path.getctime(self.path)))
self.size = os.path.getsize(self.path)
async def read(self):
return open(self.path, 'rb')
async def write(self, data):
with open(self.path, 'w') as f:
f.write(data)
async def delete(self):
os.remove(self.path)
async def exists(self):
return os.path.exists(self.path)
+67 -148
View File
@@ -2,180 +2,99 @@
# @Author : Lan # @Author : Lan
# @File : views.py # @File : views.py
# @Software: PyCharm # @Software: PyCharm
import math
import os
import time
from fastapi import APIRouter, Depends, Form from fastapi import APIRouter, Depends
from pydantic import BaseModel from apps.admin.services import FileService, ConfigService, LocalFileService
from apps.admin.dependencies import admin_required, get_file_service, get_config_service, get_local_file_service
from apps.admin.depends import admin_required from apps.admin.schemas import IDData, ConfigUpdateData, ShareItem, DeleteItem
from apps.admin.pydantics import IDData
from apps.base.models import FileCodes, KeyValue
from apps.base.utils import get_expire_info, get_file_path_name
from core.response import APIResponse from core.response import APIResponse
from core.settings import settings, data_root
from core.storage import FileStorageInterface, storages
admin_api = APIRouter( admin_api = APIRouter(prefix='/admin', tags=['管理'])
prefix='/admin',
tags=['管理'],
)
@admin_api.post('/login', dependencies=[Depends(admin_required)]) @admin_api.post('/login')
async def login(): async def login(admin: bool = Depends(admin_required)):
return APIResponse() return APIResponse()
@admin_api.delete('/file/delete', dependencies=[Depends(admin_required)]) @admin_api.delete('/file/delete')
async def file_delete(data: IDData): async def file_delete(
file_storage: FileStorageInterface = storages[settings.file_storage]() data: IDData,
file_code = await FileCodes.get(id=data.id) file_service: FileService = Depends(get_file_service),
await file_storage.delete_file(file_code) admin: bool = Depends(admin_required)
await file_code.delete() ):
await file_service.delete_file(data.id)
return APIResponse() return APIResponse()
@admin_api.get('/file/list', dependencies=[Depends(admin_required)]) @admin_api.get('/file/list')
async def file_list(page: float = 1, size: int = 10): async def file_list(
page: int = 1,
size: int = 10,
file_service: FileService = Depends(get_file_service),
admin: bool = Depends(admin_required)
):
files, total = await file_service.list_files(page, size)
return APIResponse(detail={ return APIResponse(detail={
'page': page, 'page': page,
'size': size, 'size': size,
'data': await FileCodes.all().limit(size).offset((math.ceil(page) - 1) * size), 'data': files,
'total': await FileCodes.all().count(), 'total': total,
}) })
@admin_api.get('/config/get', dependencies=[Depends(admin_required)]) @admin_api.get('/config/get')
async def get_config(): async def get_config(
return APIResponse(detail=settings.items()) config_service: ConfigService = Depends(get_config_service),
admin: bool = Depends(admin_required)
):
return APIResponse(detail=config_service.get_config())
@admin_api.patch('/config/update', dependencies=[Depends(admin_required)]) @admin_api.patch('/config/update')
async def update_config(data: dict): async def update_config(
admin_token = data.get('admin_token') data: ConfigUpdateData,
for key, value in data.items(): config_service: ConfigService = Depends(get_config_service),
if key not in settings.default_config: admin: bool = Depends(admin_required)
continue ):
if key in ['errorCount', 'errorMinute', 'max_save_seconds', 'onedrive_proxy', 'openUpload', 'port', 's3_proxy', 'uploadCount', 'uploadMinute', 'uploadSize']: await config_service.update_config(data)
data[key] = int(value)
elif key in ['opacity']:
data[key] = float(value)
else:
data[key] = value
if admin_token is None or admin_token == '':
return APIResponse(code=400, detail='管理员密码不能为空')
await KeyValue.filter(key='settings').update(value=data)
for k, v in data.items():
settings.__setattr__(k, v)
return APIResponse() return APIResponse()
# 根据code获取文件 @admin_api.get('/file/download')
async def get_file_by_id(id): async def file_download(
# 查询文件 id: int,
file_code = await FileCodes.filter(id=id).first() file_service: FileService = Depends(get_file_service),
# 检查文件是否存在 admin: bool = Depends(admin_required)
if not file_code: ):
return False, '文件不存在' file_content = await file_service.download_file(id)
return True, file_code return file_content
@admin_api.get('/file/download', dependencies=[Depends(admin_required)]) @admin_api.get('/local/lists')
async def file_download(id: int): async def get_local_lists(
file_storage: FileStorageInterface = storages[settings.file_storage]() local_file_service: LocalFileService = Depends(get_local_file_service),
has, file_code = await get_file_by_id(id) admin: bool = Depends(admin_required)
# 检查文件是否存在 ):
if not has: files = await local_file_service.list_files()
# 返回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)
class LocalFileClass:
def __init__(self, file):
self.file = file
self.path = data_root / 'local' / file
self.ctime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(os.path.getctime(self.path)))
self.size = os.path.getsize(self.path)
async def read(self):
return open(self.path, 'rb')
async def write(self, data):
with open(self.path, 'w') as f:
f.write(data)
async def delete(self):
os.remove(self.path)
async def exists(self):
return os.path.exists(self.path)
@admin_api.get('/local/lists', dependencies=[Depends(admin_required)])
async def get_local_lists():
files = []
for file in os.listdir(data_root / 'local'):
files.append(LocalFileClass(file))
return APIResponse(detail=files) return APIResponse(detail=files)
class DeleteItem(BaseModel): @admin_api.delete('/local/delete')
filename: str async def delete_local_file(
item: DeleteItem,
local_file_service: LocalFileService = Depends(get_local_file_service),
admin: bool = Depends(admin_required)
):
result = await local_file_service.delete_file(item.filename)
return APIResponse(detail=result)
@admin_api.delete('/local/delete', dependencies=[Depends(admin_required)]) @admin_api.post('/local/share')
async def delete_local_file(item: DeleteItem): async def share_local_file(
file = LocalFileClass(item.filename) item: ShareItem,
if await file.exists(): file_service: FileService = Depends(get_file_service),
await file.delete() admin: bool = Depends(admin_required)
return APIResponse(detail='删除成功') ):
return APIResponse(code=404, detail='文件不存在') share_info = await file_service.share_local_file(item)
return APIResponse(detail=share_info)
class ShareItem(BaseModel):
expire_value: int
expire_style: str = 'day'
filename: str
class File:
def __init__(self, file):
self.file = file
@admin_api.post('/local/share', dependencies=[Depends(admin_required)])
async def share_local_file(item: ShareItem):
file = LocalFileClass(item.filename)
if not await file.exists():
return APIResponse(code=404, detail='文件不存在')
text = File(await file.read())
expired_at, expired_count, used_count, code = await get_expire_info(item.expire_value, item.expire_style)
# 获取文件路径和名称
path, suffix, prefix, uuid_file_name, save_path = await get_file_path_name(item)
# 保存文件
file_storage: FileStorageInterface = storages[settings.file_storage]()
await file_storage.save_file(text, save_path)
# 创建一个新的FileCodes实例
await FileCodes.create(
code=code,
prefix=prefix,
suffix=suffix,
uuid_file_name=uuid_file_name,
file_path=path,
size=file.size,
expired_at=expired_at,
expired_count=expired_count,
used_count=used_count,
)
# 返回API响应
return APIResponse(detail={
'code': code,
'name': file.file,
})
+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 datetime
import uuid import uuid
import os import os
from fastapi import UploadFile, HTTPException 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 apps.base.models import FileCodes
from core.settings import settings from core.settings import settings
from core.utils import get_random_num, get_random_string, max_save_times_desc from core.utils import get_random_num, get_random_string, max_save_times_desc
async def get_file_path_name(file: UploadFile): async def get_file_path_name(file: UploadFile) -> Tuple[str, str, str, str, str]:
""" """获取文件路径和文件名"""
获取文件路径和文件名
: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'
}
"""
today = datetime.datetime.now() today = datetime.datetime.now()
path = f"share/data/{today.strftime('%Y/%m/%d')}" path = f"share/data/{today.strftime('%Y/%m/%d')}"
prefix, suffix = os.path.splitext(file.filename) 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}" uuid_file_name = f"{file_uuid}{suffix}"
save_path = f"{path}/{uuid_file_name}" save_path = f"{path}/{uuid_file_name}"
return path, suffix, prefix, uuid_file_name, save_path return path, suffix, prefix, uuid_file_name, save_path
async def get_expire_info(expire_value: int, expire_style: str): async def get_expire_info(expire_value: int, expire_style: str) -> Tuple[Optional[datetime.datetime], int, int, str]:
""" """获取过期信息"""
获取过期信息 expired_count, used_count = -1, 0
:param expire_value: now = datetime.datetime.now()
:param expire_style: code = None
:return: expired_at 过期时间, expired_count 可用次数, used_count 已用次数, code 随机码
""" max_timedelta = datetime.timedelta(seconds=settings.max_save_seconds) if settings.max_save_seconds > 0 else datetime.timedelta(days=7)
expired_count, used_count, now, code = -1, 0, datetime.datetime.now(), None detail = await max_save_times_desc(settings.max_save_seconds) if settings.max_save_seconds > 0 else '7天'
if int(settings.max_save_seconds) > 0: detail = f'限制最长时间为 {detail[0]},可换用其他方式'
max_timedelta = datetime.timedelta(seconds=settings.max_save_seconds)
detail = await max_save_times_desc(settings.max_save_seconds) expire_styles = {
detail = f'限制最长时间为 {detail[0]},可换用其他方式' 'day': lambda: now + datetime.timedelta(days=expire_value),
else: 'hour': lambda: now + datetime.timedelta(hours=expire_value),
max_timedelta = datetime.timedelta(days=7) 'minute': lambda: now + datetime.timedelta(minutes=expire_value),
detail = '限制最长时间为 7天,可换用其他方式' 'count': lambda: (now + datetime.timedelta(days=1), expire_value),
if expire_style == 'day': 'forever': lambda: (None, None), # 修改这里
if datetime.timedelta(days=expire_value) > max_timedelta: }
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) 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: else:
expired_at = now + datetime.timedelta(days=1) expired_at = now + datetime.timedelta(days=1)
if not code: if not code:
code = await get_random_code() code = await get_random_code()
return expired_at, expired_count, used_count, code return expired_at, expired_count, used_count, code
async def get_random_code(style='num'): async def get_random_code(style='num') -> str:
""" """获取随机字符串"""
获取随机字符串
:return:
"""
while True: while True:
code = await get_random_num() if style == 'num' else await get_random_string() code = await get_random_num() if style == 'num' else await get_random_string()
if not await FileCodes.filter(code=code).exists(): 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 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.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 apps.base.utils import get_expire_info, get_file_path_name, ip_limit
from core.response import APIResponse from core.response import APIResponse
from core.settings import settings from core.settings import settings
from core.storage import storages, FileStorageInterface from core.storage import storages, FileStorageInterface
from core.utils import get_select_token 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)]) @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')) text_size = len(text.encode('utf-8'))
# 限制 222KB max_txt_size = 222 * 1024
max_txt_size = 222 * 1024 # 转换为字节
if text_size > max_txt_size: 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) expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style)
# 创建一个新的FileCodes实例 await create_file_code(
await FileCodes.create(
code=code, code=code,
text=text, text=text,
expired_at=expired_at, expired_at=expired_at,
@@ -41,33 +43,29 @@ async def share_text(text: str = Form(...), expire_value: int = Form(default=1,
size=len(text), size=len(text),
prefix='文本分享' prefix='文本分享'
) )
# 添加IP到限制列表
ip_limit['upload'].add_ip(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)]) @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(
ip: str = Depends(ip_limit['upload'])): expire_value: int = Form(default=1, gt=0),
if file.size > settings.uploadSize: expire_style: str = Form(default='day'),
# 转换为 MB 并格式化输出 file: UploadFile = File(...),
max_size_mb = settings.uploadSize / (1024 * 1024) ip: str = Depends(ip_limit['upload'])
raise HTTPException(status_code=403, detail=f'大小超过限制,最大为{max_size_mb:.2f} MB') ):
# 获取过期信息 await validate_file_size(file, settings.uploadSize)
if expire_style not in settings.expireStyle: if expire_style not in settings.expireStyle:
raise HTTPException(status_code=400, detail='过期时间类型错误') raise HTTPException(status_code=400, detail='过期时间类型错误')
expired_at, expired_count, used_count, code = await get_expire_info(expire_value, expire_style) 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) path, suffix, prefix, uuid_file_name, save_path = await get_file_path_name(file)
# 保存文件
file_storage: FileStorageInterface = storages[settings.file_storage]() file_storage: FileStorageInterface = storages[settings.file_storage]()
await file_storage.save_file(file, save_path) await file_storage.save_file(file, save_path)
# 创建一个新的FileCodes实例
await FileCodes.create( await create_file_code(
code=code, code=code,
prefix=prefix, prefix=prefix,
suffix=suffix, 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, expired_count=expired_count,
used_count=used_count, used_count=used_count,
) )
# 添加IP到限制列表
ip_limit['upload'].add_ip(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): async def get_code_file_by_code(code, check=True):
# 查询文件
file_code = await FileCodes.filter(code=code).first() file_code = await FileCodes.filter(code=code).first()
# 检查文件是否存在
if not file_code: if not file_code:
return False, '文件不存在' return False, '文件不存在'
# 检查文件是否过期
if await file_code.is_expired() and check: if await file_code.is_expired() and check:
return False, '文件已过期', return False, '文件已过期'
return True, file_code 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/') @share_api.get('/select/')
async def get_code_file(code: str, ip: str = Depends(ip_limit['error'])): async def get_code_file(code: str, ip: str = Depends(ip_limit['error'])):
file_storage: FileStorageInterface = storages[settings.file_storage]() file_storage: FileStorageInterface = storages[settings.file_storage]()
# 获取文件
has, file_code = await get_code_file_by_code(code) has, file_code = await get_code_file_by_code(code)
# 检查文件是否存在
if not has: if not has:
# 添加IP到限制列表
ip_limit['error'].add_ip(ip) ip_limit['error'].add_ip(ip)
# 返回API响应
return APIResponse(code=404, detail=file_code) return APIResponse(code=404, detail=file_code)
# 更新文件的使用次数和过期次数
file_code.used_count += 1 await update_file_usage(file_code)
if file_code.expired_count > 0:
file_code.expired_count -= 1
# 保存文件
await file_code.save()
# 返回文件响应
return await file_storage.get_file_response(file_code) return await file_storage.get_file_response(file_code)
# 选择文件的API
@share_api.post('/select/') @share_api.post('/select/')
async def select_file(data: SelectFileModel, ip: str = Depends(ip_limit['error'])): async def select_file(data: SelectFileModel, ip: str = Depends(ip_limit['error'])):
file_storage: FileStorageInterface = storages[settings.file_storage]() file_storage: FileStorageInterface = storages[settings.file_storage]()
# 获取文件
has, file_code = await get_code_file_by_code(data.code) has, file_code = await get_code_file_by_code(data.code)
# 检查文件是否存在
if not has: if not has:
# 添加IP到限制列表
ip_limit['error'].add_ip(ip) ip_limit['error'].add_ip(ip)
# 返回API响应
return APIResponse(code=404, detail=file_code) return APIResponse(code=404, detail=file_code)
# 更新文件的使用次数和过期次数
file_code.used_count += 1 await update_file_usage(file_code)
if file_code.expired_count > 0:
file_code.expired_count -= 1
# 保存文件
await file_code.save()
# 返回API响应
return APIResponse(detail={ return APIResponse(detail={
'code': file_code.code, 'code': file_code.code,
'name': file_code.prefix + file_code.suffix, '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') @share_api.get('/download')
async def download_file(key: str, code: str, ip: str = Depends(ip_limit['error'])): async def download_file(key: str, code: str, ip: str = Depends(ip_limit['error'])):
file_storage: FileStorageInterface = storages[settings.file_storage]() file_storage: FileStorageInterface = storages[settings.file_storage]()
# 检查token是否有效 if await get_select_token(code) != key:
is_valid = await get_select_token(code) == key
if not is_valid:
# 添加IP到限制列表
ip_limit['error'].add_ip(ip) ip_limit['error'].add_ip(ip)
# 获取文件
has, file_code = await get_code_file_by_code(code, False) has, file_code = await get_code_file_by_code(code, False)
# 检查文件是否存在
if not has: if not has:
# 返回API响应
return APIResponse(code=404, detail='文件不存在') return APIResponse(code=404, detail='文件不存在')
# 如果文件是文本,返回文本内容,否则返回文件响应
if file_code.text: return APIResponse(detail=file_code.text) if file_code.text else await file_storage.get_file_response(file_code)
return APIResponse(detail=file_code.text)
else:
return await file_storage.get_file_response(file_code)
+1 -1
View File
@@ -8,7 +8,7 @@ import random
import string import string
import time import time
from apps.base.depends import IPRateLimit from apps.base.dependencies import IPRateLimit
async def get_random_num(): async def get_random_num():
+17 -15
View File
@@ -38,28 +38,30 @@ async def lifespan(app: FastAPI):
# 初始化数据库 # 初始化数据库
await init_db() await init_db()
# 启动后台任务,不定时删除过期文件 # 启动后台任务
task = asyncio.create_task(delete_expire_files()) task = asyncio.create_task(delete_expire_files())
# 读取用户配置
user_config, created = await KeyValue.get_or_create(key='settings', defaults={'value': DEFAULT_CONFIG}) # 加载配置
await load_config()
try:
yield
finally:
# 清理操作
task.cancel()
await asyncio.gather(task, return_exceptions=True)
await Tortoise.close_connections()
async def load_config():
user_config, _ = await KeyValue.get_or_create(key='settings', defaults={'value': DEFAULT_CONFIG})
settings.user_config = user_config.value settings.user_config = user_config.value
# 更新 ip_limit 配置
ip_limit['error'].minutes = settings.errorMinute ip_limit['error'].minutes = settings.errorMinute
ip_limit['error'].count = settings.errorCount ip_limit['error'].count = settings.errorCount
ip_limit['upload'].minutes = settings.uploadMinute ip_limit['upload'].minutes = settings.uploadMinute
ip_limit['upload'].count = settings.uploadCount ip_limit['upload'].count = settings.uploadCount
yield
# 清理操作
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
# 关闭数据库连接
await Tortoise.close_connections()
app = FastAPI(lifespan=lifespan) app = FastAPI(lifespan=lifespan)