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)

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: