From 745ada59bb18d2b480fcfb28eda356aec08e851b Mon Sep 17 00:00:00 2001
From: lan <hhxk666@gmail.com>
Date: Wed, 8 Feb 2023 03:17:22 +0800
Subject: [PATCH] =?UTF-8?q?add:=E8=A7=A3=E5=86=B3select=E5=8F=AF=E8=B7=B3?=
 =?UTF-8?q?=E8=BF=87=E6=AC=A1=E6=95=B0=E9=99=90=E5=88=B6=EF=BC=8C=E6=96=B0?=
 =?UTF-8?q?=E5=A2=9Esha256=E5=8A=A0=E5=AF=86token=EF=BC=8C=E9=99=90?=
 =?UTF-8?q?=E5=88=B6=E5=8F=96=E4=BB=B6=E7=A0=81+ip+=E6=97=B6=E9=97=B4999s?=
 =?UTF-8?q?=E5=86=85=E5=8F=96=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 core/utils.py |  7 +++++++
 main.py       | 14 +++++++++-----
 2 files changed, 16 insertions(+), 5 deletions(-)

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