add:解决select可跳过次数限制,新增sha256加密token,限制取件码+ip+时间999s内取件

This commit is contained in:
lan
2023-02-08 03:17:22 +08:00
parent a9aac9ed40
commit 745ada59bb
2 changed files with 16 additions and 5 deletions
+7
View File
@@ -1,6 +1,9 @@
import datetime import datetime
import hashlib
import random import random
import asyncio import asyncio
import time
from sqlalchemy import or_, select, delete from sqlalchemy import or_, select, delete
from sqlalchemy.ext.asyncio.session import AsyncSession from sqlalchemy.ext.asyncio.session import AsyncSession
from .database import Codes, engine 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(): while (await s.execute(select(Codes.id).where(Codes.code == code))).scalar():
code = random.randint(10000, 99999) code = random.randint(10000, 99999)
return str(code) 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()
+9 -5
View File
@@ -1,5 +1,6 @@
import asyncio import asyncio
import datetime import datetime
import time
import uuid import uuid
from pathlib import Path from pathlib import Path
from sqlalchemy import select, func, update from sqlalchemy import select, func, update
@@ -7,13 +8,12 @@ from sqlalchemy.ext.asyncio import AsyncSession
from starlette.requests import Request from starlette.requests import Request
from starlette.responses import HTMLResponse, FileResponse from starlette.responses import HTMLResponse, FileResponse
from starlette.staticfiles import StaticFiles from starlette.staticfiles import StaticFiles
import os from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files, get_token
import shutil
from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files
from core.depends import admin_required from core.depends import admin_required
from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException, BackgroundTasks from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException, BackgroundTasks
from core.database import init_models, Options, Codes, get_session from core.database import init_models, Options, Codes, get_session
from settings import settings from settings import settings
import hashlib
# 实例化FastAPI # 实例化FastAPI
app = FastAPI(debug=settings.DEBUG, redoc_url=None, docs_url=None, openapi_url=None) 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.execute(update(Codes).where(Codes.id == info.id).values(count=info.count - 1))
await s.commit() await s.commit()
if info.type != 'text': if info.type != 'text':
info.text = await storage.get_url(info) info.text = f'/select?code={info.code}&token={get_token(code, ip)}'
return { return {
'detail': f'取件成功,请立即下载,避免失效!', 'detail': f'取件成功,请立即下载,避免失效!',
'data': {'type': info.type, 'text': info.text, 'name': info.name, 'code': info.code} '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') @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) query = select(Codes).where(Codes.code == code)
info = (await s.execute(query)).scalars().first() info = (await s.execute(query)).scalars().first()