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()

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()