From d23cffb570f0c1c330c9e4c978c5ce7fcf67f576 Mon Sep 17 00:00:00 2001 From: lan Date: Tue, 13 Dec 2022 19:16:29 +0800 Subject: [PATCH] =?UTF-8?q?update:=E6=96=B0=E5=A2=9E=E9=99=90=E5=88=B6?= =?UTF-8?q?=E6=97=B6=E9=97=B4=E5=86=85=E4=B8=8A=E4=BC=A0=E6=AC=A1=E6=95=B0?= =?UTF-8?q?=EF=BC=8C=E7=A7=BB=E9=99=A4IP=E9=BB=91=E5=90=8D=E5=8D=95?= =?UTF-8?q?=E5=88=97=E8=A1=A8=E8=BF=87=E6=9C=9FIP=EF=BC=8C=E9=98=B2?= =?UTF-8?q?=E6=AD=A2=E7=82=B8=E5=86=85=E5=AD=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- depends.py | 21 +++++++++++++++----- main.py | 56 +++++++++++++++++++++++++++++++++-------------------- settings.py | 33 +++++++++++++++++-------------- utils.py | 10 +++++++++- 4 files changed, 79 insertions(+), 41 deletions(-) diff --git a/depends.py b/depends.py index 72bda9c..0924227 100644 --- a/depends.py +++ b/depends.py @@ -13,25 +13,36 @@ async def admin_required(pwd: Union[str, None] = Header(default=None)): class IPRateLimit: ips = {} + count = 0 + minutes = 0 + + def __init__(self, count, minutes): + self.count = count + self.minutes = minutes def check_ip(self, ip): # 检查ip是否被禁止 if ip in self.ips: - if self.ips[ip]['count'] >= settings.ERROR_COUNT: - if self.ips[ip]['time'] + timedelta(minutes=settings.ERROR_MINUTE) > datetime.now(): + 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(cls, ip): ip_info = cls.ips.get(ip, {'count': 0, 'time': datetime.now()}) ip_info['count'] += 1 cls.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.client.host if not self.check_ip(ip): - raise HTTPException(status_code=400, detail="错误次数过多，请稍后再试") + raise HTTPException(status_code=400, detail=f"请求次数过多，请稍后再试") return ip diff --git a/main.py b/main.py index c6ddd37..32effe6 100644 --- a/main.py +++ b/main.py @@ -11,58 +11,66 @@ from sqlalchemy import select, update from sqlalchemy.ext.asyncio.session import AsyncSession import settings -from utils import delete_expire_files, storage, get_code +from utils import delete_expire_files, storage, get_code, error_ip_limit, upload_ip_limit from database import get_session, Codes, init_models -from depends import admin_required, IPRateLimit +from depends import admin_required +# 实例化FastAPI app = FastAPI(debug=settings.DEBUG) +# 数据存储文件夹 DATA_ROOT = Path(settings.DATA_ROOT) if not DATA_ROOT.exists(): DATA_ROOT.mkdir(parents=True) - +# 静态文件夹 app.mount(settings.STATIC_URL, StaticFiles(directory=DATA_ROOT), name="static") @app.on_event('startup') async def startup(): + # 初始化数据库 await init_models() + # 启动后台任务，不定时删除过期文件 asyncio.create_task(delete_expire_files()) +# 首页页面 index_html = open('templates/index.html', 'r', encoding='utf-8').read() \ .replace('{{title}}', settings.TITLE) \ .replace('{{description}}', settings.DESCRIPTION) \ .replace('{{keywords}}', settings.KEYWORDS) +# 管理页面 admin_html = open('templates/admin.html', 'r', encoding='utf-8').read() \ .replace('{{title}}', settings.TITLE) \ .replace('{{description}}', settings.DESCRIPTION) \ .replace('{{keywords}}', settings.KEYWORDS) -ip_limit = IPRateLimit() - -@app.get(f'/{settings.ADMIN_ADDRESS}') +@app.get(f'/{settings.ADMIN_ADDRESS}', description='管理页面', response_class=HTMLResponse) async def admin(): return HTMLResponse(admin_html) -@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)]) +@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)], description='查询数据库列表') async def admin_post(s: AsyncSession = Depends(get_session)): - query = select(Codes) - codes = (await s.execute(query)).scalars().all() + # 查询数据库列表 + codes = (await s.execute(select(Codes))).scalars().all() return {'detail': '查询成功', 'data': codes} -@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)]) +@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)], description='删除数据库记录') async def admin_delete(code: str, s: AsyncSession = Depends(get_session)): + # 找到相应记录 query = select(Codes).where(Codes.code == code) + # 找到第一条记录 file = (await s.execute(query)).scalars().first() - if file: - if file.type != 'text': - await storage.delete_file(file.text) - await s.delete(file) - await s.commit() + # 如果记录存在，并且不是文本 + if file and file.type != 'text': + # 删除文件 + await storage.delete_file(file.text) + # 删除数据库记录 + await s.delete(file) + await s.commit() return {'detail': '删除成功'} @@ -72,25 +80,30 @@ async def index(): @app.get('/select') -async def get_file(code: str, s: AsyncSession = Depends(get_session)): +async def get_file(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)): + # 查出数据库记录 query = select(Codes).where(Codes.code == code) info = (await s.execute(query)).scalars().first() + # 如果记录不存在，IP错误次数+1 if not info: - raise HTTPException(status_code=404, detail="口令不存在") + error_ip_limit.add_ip(ip) + raise HTTPException(status_code=404, detail="口令不存在，次数过多将被禁止访问") + # 如果是文本，直接返回 if info.type == 'text': return {'detail': '查询成功', 'data': info.text} + # 如果是文件，返回文件 else: filepath = await storage.get_filepath(info.text) return FileResponse(filepath, filename=info.name) @app.post('/') -async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)): +async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)): query = select(Codes).where(Codes.code == code) info = (await s.execute(query)).scalars().first() if not info: - error_count = settings.ERROR_COUNT - ip_limit.add_ip(ip) - raise HTTPException(status_code=404, detail=f"取件码错误，错误{error_count}次将被禁止10分钟") + error_count = settings.ERROR_COUNT - error_ip_limit.add_ip(ip) + raise HTTPException(status_code=404, detail=f"取件码错误，{error_count}次后将被禁止{settings.ERROR_MINUTE}分钟") if info.exp_time < datetime.datetime.now() or info.count == 0: if info.type != "text": await storage.delete_file(info.text) @@ -109,7 +122,7 @@ async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depend @app.post('/share') async def share(background_tasks: BackgroundTasks, text: str = Form(default=None), style: str = Form(default='2'), - value: int = Form(default=1), file: UploadFile = File(default=None), + value: int = Form(default=1), file: UploadFile = File(default=None), ip: str = Depends(upload_ip_limit), s: AsyncSession = Depends(get_session)): code = await get_code(s) if style == '2': @@ -137,6 +150,7 @@ async def share(background_tasks: BackgroundTasks, text: str = Form(default=None info = Codes(code=code, text=_text, size=size, type=_type, name=name, count=exp_count, exp_time=exp_time, key=key) s.add(info) await s.commit() + upload_ip_limit.add_ip(ip) return { 'detail': '分享成功，请点击取件码按钮查看上传列表', 'data': {'code': code, 'key': key, 'name': name, 'text': _text} diff --git a/settings.py b/settings.py index 8405216..e7b7a0a 100644 --- a/settings.py +++ b/settings.py @@ -1,31 +1,36 @@ from starlette.config import Config +# 配置文件.env config = Config(".env") - +# 是否开启DEBUG模式 DEBUG = config('DEBUG', cast=bool, default=False) - +# 端口 PORT = config('PORT', cast=int, default=12345) - +# Sqlite数据库文件 DATABASE_URL = config('DATABASE_URL', cast=str, default="sqlite+aiosqlite:///database.db") - +# 静态文件夹 DATA_ROOT = config('DATA_ROOT', cast=str, default="./static") - +# 静态文件夹URL STATIC_URL = config('STATIC_URL', cast=str, default="/static") - +# 错误次数 ERROR_COUNT = config('ERROR_COUNT', cast=int, default=5) - +# 错误限制分钟数 ERROR_MINUTE = config('ERROR_MINUTE', cast=int, default=10) - +# 上传次数 +UPLOAD_COUNT = config('UPLOAD_COUNT', cast=int, default=60) +# 上传限制分钟数 +UPLOAD_MINUTE = config('UPLOAD_MINUTE', cast=int, default=1) +# 管理地址 ADMIN_ADDRESS = config('ADMIN_ADDRESS', cast=str, default="admin") - +# 管理密码 ADMIN_PASSWORD = config('ADMIN_PASSWORD', cast=str, default="admin") - +# 文件大小限制，默认10MB FILE_SIZE_LIMIT = config('FILE_SIZE_LIMIT', cast=int, default=10) * 1024 * 1024 - +# 网站标题 TITLE = config('TITLE', cast=str, default="文件快递柜") - +# 网站描述 DESCRIPTION = config('DESCRIPTION', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件") - +# 网站关键词 KEYWORDS = config('KEYWORDS', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件") - +# 存储引擎 STORAGE_ENGINE = config('STORAGE_ENGINE', cast=str, default="filesystem") diff --git a/utils.py b/utils.py index 0b297f3..fd648bd 100644 --- a/utils.py +++ b/utils.py @@ -5,14 +5,22 @@ from sqlalchemy import or_, select, delete from sqlalchemy.ext.asyncio.session import AsyncSession import settings from database import Codes, engine +from depends import IPRateLimit from storage import STORAGE_ENGINE storage = STORAGE_ENGINE[settings.STORAGE_ENGINE]() +# 错误IP限制器 +error_ip_limit = IPRateLimit(settings.ERROR_COUNT, settings.ERROR_MINUTE) +# 上传文件限制器 +upload_ip_limit = IPRateLimit(settings.UPLOAD_COUNT, settings.UPLOAD_MINUTE) + async def delete_expire_files(): while True: async with AsyncSession(engine, expire_on_commit=False) as s: + await error_ip_limit.remove_expired_ip() + await upload_ip_limit.remove_expired_ip() query = select(Codes).where(or_(Codes.exp_time < datetime.datetime.now(), Codes.count == 0)) exps = (await s.execute(query)).scalars().all() files = [] @@ -25,7 +33,7 @@ async def delete_expire_files(): query = delete(Codes).where(Codes.id.in_(exps_ids)) await s.execute(query) await s.commit() - await asyncio.sleep(random.randint(60, 300)) + await asyncio.sleep(random.randint(2, 2)) async def get_code(s: AsyncSession):