diff --git a/apps/base/views.py b/apps/base/views.py index 8301cdf..f16e4d8 100644 --- a/apps/base/views.py +++ b/apps/base/views.py @@ -3,6 +3,7 @@ # @File : views.py # @Software: PyCharm from fastapi import APIRouter, Form, UploadFile, File, Depends +from starlette.responses import FileResponse from apps.base.models import FileCodes from apps.base.pydantics import SelectFileModel @@ -71,3 +72,20 @@ async def select_file(data: SelectFileModel, ip: str = Depends(error_ip_limit)): 'size': file_code.size, 'text': file_code.text if file_code.text is not None else await file_storage.get_file_url(file_code), }) + + +@share_api.get('/download') +async def download_file(key: str, code: str, ip: str = Depends(error_ip_limit)): + is_valid = await file_storage.get_select_token(code) == key + if not is_valid: + error_ip_limit.add_ip(ip) + file_code = await FileCodes.filter(code=code).first() + if not file_code: + return APIResponse(code=404, detail='文件不存在') + if file_code.text: + return APIResponse(detail=file_code.text) + else: + file_path = file_storage.root_path / await file_code.get_file_path() + if not file_path.exists(): + return APIResponse(code=404, detail='文件已过期删除') + return FileResponse(file_path, filename=file_code.prefix + file_code.suffix) diff --git a/core/storage.py b/core/storage.py index 88027af..426197c 100644 --- a/core/storage.py +++ b/core/storage.py @@ -3,6 +3,9 @@ # @File : storage.py # @Software: PyCharm import asyncio +import hashlib +import time + import aioboto3 from fastapi import UploadFile from pathlib import Path @@ -14,6 +17,7 @@ class SystemFileStorage: def __init__(self): self.chunk_size = 256 * 1024 self.root_path = Path('./data') + self.token = '123456' def _save(self, file, save_path): with open(save_path, 'wb') as f: @@ -33,8 +37,11 @@ class SystemFileStorage: if save_path.exists(): save_path.unlink() + async def get_select_token(self, code): + return hashlib.sha256(f"{code}{int(time.time() / 1000)}000{self.token}".encode()).hexdigest() + async def get_file_url(self, file_code: FileCodes): - return f"/api/select?code={file_code.code}" + return f'/share/download?key={await self.get_select_token(file_code.code)}&code={file_code.code}' class S3FileStorage: diff --git a/filecodebox.db-wal b/filecodebox.db-wal index 5c7411d..e69de29 100644 Binary files a/filecodebox.db-wal and b/filecodebox.db-wal differ