From 4e069486bf33917657849c65b81e7c4992d40883 Mon Sep 17 00:00:00 2001 From: veoco Date: Mon, 12 Dec 2022 11:18:55 +0800 Subject: [PATCH] =?UTF-8?q?=E6=96=B0=E5=A2=9E=E5=AD=98=E5=82=A8=E5=BC=95?= =?UTF-8?q?=E6=93=8E=E7=BB=9F=E4=B8=80=E6=96=87=E4=BB=B6=E8=AF=BB=E5=86=99?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 37 +++++++++++-------------------------- settings.py | 2 ++ storage.py | 45 +++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 58 insertions(+), 26 deletions(-) create mode 100644 storage.py diff --git a/main.py b/main.py index 3214062..bd9471c 100644 --- a/main.py +++ b/main.py @@ -1,7 +1,5 @@ import datetime -import os import uuid -import threading import random import asyncio from pathlib import Path @@ -16,6 +14,7 @@ from sqlalchemy.ext.asyncio.session import AsyncSession import settings from database import get_session, Codes, init_models, engine +from storage import STORAGE_ENGINE app = FastAPI(debug=settings.DEBUG) @@ -26,6 +25,8 @@ if not DATA_ROOT.exists(): STATIC_URL = settings.STATIC_URL app.mount(STATIC_URL, StaticFiles(directory=DATA_ROOT), name="static") +storage = STORAGE_ENGINE[settings.STORAGE_ENGINE]() + @app.on_event('startup') async def startup(): @@ -45,18 +46,13 @@ admin_html = open('templates/admin.html', 'r', encoding='utf-8').read() \ error_ip_count = {} -def delete_file(files): - for file in files: - if file['type'] != 'text': - os.remove(DATA_ROOT / file['text'].lstrip(STATIC_URL + '/')) - - async def delete_expire_files(): while True: async with AsyncSession(engine, expire_on_commit=False) as s: query = select(Codes).where(or_(Codes.exp_time < datetime.datetime.now(), Codes.count == 0)) exps = (await s.execute(query)).scalars().all() - await asyncio.to_thread(delete_file, [{'type': old.type, 'text': old.text} for old in exps]) + files = [{'type': old.type, 'text': old.text} for old in exps] + await storage.delete_files(files) exps_ids = [exp.id for exp in exps] query = delete(Codes).where(Codes.id.in_(exps_ids)) await s.execute(query) @@ -71,18 +67,6 @@ async def get_code(s: AsyncSession): return str(code) -def get_file_name(key, ext, file, file_bytes): - now = datetime.datetime.now() - path = DATA_ROOT / f"upload/{now.year}/{now.month}/{now.day}/" - name = f'{key}.{ext}' - if not path.exists(): - path.mkdir(parents=True) - filepath = path / name - with open(filepath, 'wb') as f: - f.write(file_bytes) - return f"{STATIC_URL}/{filepath.relative_to(DATA_ROOT)}", file.content_type, file.filename - - @app.get(f'/{settings.ADMIN_ADDRESS}') async def admin(): return HTMLResponse(admin_html) @@ -103,7 +87,7 @@ async def admin_delete(request: Request, code: str, s: AsyncSession = Depends(ge if request.headers.get('pwd') == settings.ADMIN_PASSWORD: query = select(Codes).where(Codes.code == code) file = (await s.execute(query)).scalars().first() - await asyncio.to_thread(delete_file, [{'type': file.type, 'text': file.text}]) + await storage.delete_file({'type': file.type, 'text': file.text}) await s.delete(file) await s.commit() return {'code': 200, 'msg': '删除成功'} @@ -142,7 +126,8 @@ async def get_file(code: str, s: AsyncSession = Depends(get_session)): if info.type == 'text': return {'code': code, 'msg': '查询成功', 'data': info.text} else: - return FileResponse(DATA_ROOT / info.text.lstrip(STATIC_URL + '/'), filename=info.name) + filepath = await storage.get_filepath(info.text) + return FileResponse(filepath, filename=info.name) else: return {'code': 404, 'msg': '口令不存在'} @@ -157,7 +142,7 @@ async def index(request: Request, code: str, s: AsyncSession = Depends(get_sessi if not info: return {'code': 404, 'msg': f'取件码错误，错误{settings.ERROR_COUNT - ip_error(ip)}次将被禁止10分钟'} if info.exp_time < datetime.datetime.now() or info.count == 0: - threading.Thread(target=delete_file, args=([{'type': info.type, 'text': info.text}],)).start() + await storage.delete_file({'type': info.type, 'text': info.text}) await s.delete(info) await s.commit() return {'code': 404, 'msg': '取件码已过期，请联系寄件人'} @@ -193,11 +178,11 @@ async def share(text: str = Form(default=None), style: str = Form(default='2'), exp_count = -1 key = uuid.uuid4().hex if file: - file_bytes = file.file.read() + file_bytes = await file.read() size = len(file_bytes) if size > settings.FILE_SIZE_LIMIT: return {'code': 404, 'msg': '文件过大'} - _text, _type, name = get_file_name(key, file.filename.split('.')[-1], file, file_bytes) + _text, _type, name = await storage.save_file(file, file_bytes, key), file.content_type, file.filename else: size, _text, _type, name = len(text), text, 'text', '文本分享' info = Codes( diff --git a/settings.py b/settings.py index ed14dcf..9959f14 100644 --- a/settings.py +++ b/settings.py @@ -25,3 +25,5 @@ TITLE = config('TITLE', cast=str, default="文件快递柜") DESCRIPTION = config('DESCRIPTION', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件") KEYWORDS = config('KEYWORDS', cast=str, default="FileCodeBox，文件快递柜，口令传送箱，匿名口令分享文本，文件等文件") + +STORAGE_ENGINE = config('STORAGE_ENGINE', cast=str, default="filesystem") diff --git a/storage.py b/storage.py new file mode 100644 index 0000000..76c9d69 --- /dev/null +++ b/storage.py @@ -0,0 +1,45 @@ +import os +import asyncio +import datetime +from pathlib import Path + +import settings + + +class FileSystemStorage: + DATA_ROOT = Path(settings.DATA_ROOT) + STATIC_URL = settings.STATIC_URL + NAME = "filesystem" + + async def get_filepath(self, path): + return self.DATA_ROOT / path.lstrip(self.STATIC_URL + '/') + + def _save(self, filepath, file_bytes): + with open(filepath, 'wb') as f: + f.write(file_bytes) + + async def save_file(self, file, file_bytes, key): + now = datetime.datetime.now() + path = self.DATA_ROOT / f"upload/{now.year}/{now.month}/{now.day}/" + ext = file.filename.split('.')[-1] + name = f'{key}.{ext}' + if not path.exists(): + path.mkdir(parents=True) + filepath = path / name + await asyncio.to_thread(self._save, filepath, file_bytes) + text = f"{self.STATIC_URL}/{filepath.relative_to(self.DATA_ROOT)}" + return text + + async def delete_file(self, file): + filepath = self.DATA_ROOT / file['text'].lstrip(self.STATIC_URL + '/') + await asyncio.to_thread(os.remove, filepath) + + async def delete_files(self, files): + for file in files: + if file['type'] != 'text': + await self.delete_file(file) + + +STORAGE_ENGINE = { + "filesystem": FileSystemStorage +} \ No newline at end of file