|
| 1 | +import datetime |
| 2 | +import uuid |
| 3 | +import random |
| 4 | +import asyncio |
| 5 | +from pathlib import Path |
| 6 | + |
| 7 | +from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException |
| 8 | +from starlette.responses import HTMLResponse, FileResponse |
| 9 | +from starlette.staticfiles import StaticFiles |
| 10 | + |
| 11 | +from sqlalchemy import or_, select, update, delete |
| 12 | +from sqlalchemy.ext.asyncio.session import AsyncSession |
| 13 | + |
| 14 | +import settings |
| 15 | +from database import get_session, Codes, init_models, engine |
| 16 | +from storage import STORAGE_ENGINE |
| 17 | +from depends import admin_required, IPRateLimit |
| 18 | + |
| 19 | +app = FastAPI(debug=settings.DEBUG) |
| 20 | + |
| 21 | +DATA_ROOT = Path(settings.DATA_ROOT) |
| 22 | +if not DATA_ROOT.exists(): |
| 23 | + DATA_ROOT.mkdir(parents=True) |
| 24 | + |
| 25 | +STATIC_URL = settings.STATIC_URL |
| 26 | +app.mount(STATIC_URL, StaticFiles(directory=DATA_ROOT), name="static") |
| 27 | + |
| 28 | +storage = STORAGE_ENGINE[settings.STORAGE_ENGINE]() |
| 29 | + |
| 30 | + |
| 31 | +@app.on_event('startup') |
| 32 | +async def startup(): |
| 33 | + await init_models() |
| 34 | + asyncio.create_task(delete_expire_files()) |
| 35 | + |
| 36 | + |
| 37 | +index_html = open('templates/index.html', 'r', encoding='utf-8').read() \ |
| 38 | + .replace('{{title}}', settings.TITLE) \ |
| 39 | + .replace('{{description}}', settings.DESCRIPTION) \ |
| 40 | + .replace('{{keywords}}', settings.KEYWORDS) |
| 41 | +admin_html = open('templates/admin.html', 'r', encoding='utf-8').read() \ |
| 42 | + .replace('{{title}}', settings.TITLE) \ |
| 43 | + .replace('{{description}}', settings.DESCRIPTION) \ |
| 44 | + .replace('{{keywords}}', settings.KEYWORDS) |
| 45 | + |
| 46 | +ip_limit = IPRateLimit() |
| 47 | + |
| 48 | + |
| 49 | +async def delete_expire_files(): |
| 50 | + while True: |
| 51 | + async with AsyncSession(engine, expire_on_commit=False) as s: |
| 52 | + query = select(Codes).where(or_(Codes.exp_time < datetime.datetime.now(), Codes.count == 0)) |
| 53 | + exps = (await s.execute(query)).scalars().all() |
| 54 | + files = [{'type': old.type, 'text': old.text} for old in exps] |
| 55 | + await storage.delete_files(files) |
| 56 | + exps_ids = [exp.id for exp in exps] |
| 57 | + query = delete(Codes).where(Codes.id.in_(exps_ids)) |
| 58 | + await s.execute(query) |
| 59 | + await s.commit() |
| 60 | + await asyncio.sleep(random.randint(60, 300)) |
| 61 | + |
| 62 | + |
| 63 | +async def get_code(s: AsyncSession): |
| 64 | + code = random.randint(10000, 99999) |
| 65 | + while (await s.execute(select(Codes.id).where(Codes.code == code))).scalar(): |
| 66 | + code = random.randint(10000, 99999) |
| 67 | + return str(code) |
| 68 | + |
| 69 | + |
| 70 | +@app.get(f'/{settings.ADMIN_ADDRESS}') |
| 71 | +async def admin(): |
| 72 | + return HTMLResponse(admin_html) |
| 73 | + |
| 74 | + |
| 75 | +@app.post(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)]) |
| 76 | +async def admin_post(s: AsyncSession = Depends(get_session)): |
| 77 | + query = select(Codes) |
| 78 | + codes = (await s.execute(query)).scalars().all() |
| 79 | + return {'detail': '查询成功', 'data': codes} |
| 80 | + |
| 81 | + |
| 82 | +@app.delete(f'/{settings.ADMIN_ADDRESS}', dependencies=[Depends(admin_required)]) |
| 83 | +async def admin_delete(code: str, s: AsyncSession = Depends(get_session)): |
| 84 | + query = select(Codes).where(Codes.code == code) |
| 85 | + file = (await s.execute(query)).scalars().first() |
| 86 | + await storage.delete_file({'type': file.type, 'text': file.text}) |
| 87 | + await s.delete(file) |
| 88 | + await s.commit() |
| 89 | + return {'detail': '删除成功'} |
| 90 | + |
| 91 | + |
| 92 | +@app.get('/') |
| 93 | +async def index(): |
| 94 | + return HTMLResponse(index_html) |
| 95 | + |
| 96 | + |
| 97 | +@app.get('/select') |
| 98 | +async def get_file(code: str, s: AsyncSession = Depends(get_session)): |
| 99 | + query = select(Codes).where(Codes.code == code) |
| 100 | + info = (await s.execute(query)).scalars().first() |
| 101 | + if not info: |
| 102 | + raise HTTPException(status_code=404, detail="口令不存在") |
| 103 | + if info.type == 'text': |
| 104 | + return {'detail': '查询成功', 'data': info.text} |
| 105 | + else: |
| 106 | + filepath = await storage.get_filepath(info.text) |
| 107 | + return FileResponse(filepath, filename=info.name) |
| 108 | + |
| 109 | + |
| 110 | +@app.post('/') |
| 111 | +async def index(code: str, ip: str = Depends(ip_limit), s: AsyncSession = Depends(get_session)): |
| 112 | + query = select(Codes).where(Codes.code == code) |
| 113 | + info = (await s.execute(query)).scalars().first() |
| 114 | + if not info: |
| 115 | + error_count = settings.ERROR_COUNT - ip_limit.add_ip(ip) |
| 116 | + raise HTTPException(status_code=404, detail=f"取件码错误，错误{error_count}次将被禁止10分钟") |
| 117 | + if info.exp_time < datetime.datetime.now() or info.count == 0: |
| 118 | + await storage.delete_file({'type': info.type, 'text': info.text}) |
| 119 | + await s.delete(info) |
| 120 | + await s.commit() |
| 121 | + raise HTTPException(status_code=404, detail="取件码已过期，请联系寄件人") |
| 122 | + await s.execute(update(Codes).where(Codes.id == info.id).values(count=info.count - 1)) |
| 123 | + await s.commit() |
| 124 | + if info.type != 'text': |
| 125 | + info.text = f'/select?code={code}' |
| 126 | + return { |
| 127 | + 'detail': '取件成功，请点击"取"查看', |
| 128 | + 'data': {'type': info.type, 'text': info.text, 'name': info.name, 'code': info.code} |
| 129 | + } |
| 130 | + |
| 131 | + |
| 132 | +@app.post('/share') |
| 133 | +async def share(text: str = Form(default=None), style: str = Form(default='2'), value: int = Form(default=1), |
| 134 | + file: UploadFile = File(default=None), s: AsyncSession = Depends(get_session)): |
| 135 | + code = await get_code(s) |
| 136 | + if style == '2': |
| 137 | + if value > 7: |
| 138 | + raise HTTPException(status_code=400, detail="最大有效天数为7天") |
| 139 | + exp_time = datetime.datetime.now() + datetime.timedelta(days=value) |
| 140 | + exp_count = -1 |
| 141 | + elif style == '1': |
| 142 | + if value < 1: |
| 143 | + raise HTTPException(status_code=400, detail="最小有效次数为1次") |
| 144 | + exp_time = datetime.datetime.now() + datetime.timedelta(days=1) |
| 145 | + exp_count = value |
| 146 | + else: |
| 147 | + exp_time = datetime.datetime.now() + datetime.timedelta(days=1) |
| 148 | + exp_count = -1 |
| 149 | + key = uuid.uuid4().hex |
| 150 | + if file: |
| 151 | + file_bytes = await file.read() |
| 152 | + size = len(file_bytes) |
| 153 | + if size > settings.FILE_SIZE_LIMIT: |
| 154 | + raise HTTPException(status_code=400, detail="文件过大") |
| 155 | + _text, _type, name = await storage.save_file(file, file_bytes, key), file.content_type, file.filename |
| 156 | + else: |
| 157 | + size, _text, _type, name = len(text), text, 'text', '文本分享' |
| 158 | + info = Codes( |
| 159 | + code=code, |
| 160 | + text=_text, |
| 161 | + size=size, |
| 162 | + type=_type, |
| 163 | + name=name, |
| 164 | + count=exp_count, |
| 165 | + exp_time=exp_time, |
| 166 | + key=key |
| 167 | + ) |
| 168 | + s.add(info) |
| 169 | + await s.commit() |
| 170 | + return { |
| 171 | + 'detail': '分享成功，请点击文件箱查看取件码', |
| 172 | + 'data': {'code': code, 'key': key, 'name': name, 'text': _text} |
| 173 | + } |
| 174 | + |
| 175 | + |
| 176 | +if __name__ == '__main__': |
| 177 | + import uvicorn |
| 178 | + |
| 179 | + uvicorn.run('main:app', host='0.0.0.0', port=settings.PORT, debug=settings.DEBUG) |
0 commit comments