diff --git a/core/utils.py b/core/utils.py index 8b0f949..6d3fda3 100644 --- a/core/utils.py +++ b/core/utils.py @@ -1,6 +1,9 @@ import datetime +import hashlib import random import asyncio +import time + from sqlalchemy import or_, select, delete from sqlalchemy.ext.asyncio.session import AsyncSession from .database import Codes, engine @@ -41,3 +44,7 @@ async def get_code(s: AsyncSession): while (await s.execute(select(Codes.id).where(Codes.code == code))).scalar(): code = random.randint(10000, 99999) return str(code) + + +async def get_token(ip, code): + return hashlib.sha256(f"{ip}{code}{int(time.time()) / 1000}000{settings.SECRET_KEY}".encode()).hexdigest() diff --git a/main.py b/main.py index 4ff8c46..f1c953a 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,6 @@ import asyncio import datetime +import time import uuid from pathlib import Path from sqlalchemy import select, func, update @@ -7,13 +8,12 @@ from sqlalchemy.ext.asyncio import AsyncSession from starlette.requests import Request from starlette.responses import HTMLResponse, FileResponse from starlette.staticfiles import StaticFiles -import os -import shutil -from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files +from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files, get_token from core.depends import admin_required from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException, BackgroundTasks from core.database import init_models, Options, Codes, get_session from settings import settings +import hashlib # 实例化FastAPI app = FastAPI(debug=settings.DEBUG, redoc_url=None, docs_url=None, openapi_url=None) @@ -132,7 +132,7 @@ async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = await s.execute(update(Codes).where(Codes.id == info.id).values(count=info.count - 1)) await s.commit() if info.type != 'text': - info.text = await storage.get_url(info) + info.text = f'/select?code={info.code}&token={get_token(code, ip)}' return { 'detail': f'取件成功,请立即下载,避免失效!', 'data': {'type': info.type, 'text': info.text, 'name': info.name, 'code': info.code} @@ -150,7 +150,11 @@ async def banner(request: Request): @app.get('/select') -async def get_file(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)): +async def get_file(code: str, token: str, ip: str = Depends(error_ip_limit), s: AsyncSession = Depends(get_session)): + # 验证token + if token != get_token(code, ip): + error_ip_limit.add_ip(ip) + raise HTTPException(status_code=403, detail="口令错误,或已过期,次数过多将被禁止访问") # 查出数据库记录 query = select(Codes).where(Codes.code == code) info = (await s.execute(query)).scalars().first()