update:优化分析模式

Author: vastsa

Dockerfile

@@ -5,6 +5,8 @@ LABEL version="6"
 
 
 COPY . /app
+RUN ln -sf /usr/share/zoneinfo/Asia/Shanghai /etc/localtime
+RUN echo 'Asia/Shanghai' >/etc/timezone
 WORKDIR /app
 RUN pip install -r requirements.txt
 EXPOSE 12345

core/base_model.py

@@ -2,6 +2,7 @@
 import hashlib
 import random
 import asyncio
+import string
 import time
 
 from sqlalchemy import or_, select, delete
@@ -39,12 +40,44 @@ async def delete_expire_files():
         await asyncio.sleep(settings.DELETE_EXPIRE_FILES_INTERVAL * 60)
 
 
-async def get_code(s: AsyncSession):
-    code = random.randint(10000, 99999)
+async def get_random_num():
+    return random.randint(10000, 99999)
+
+
+async def get_random_string():
+    r_s = string.ascii_letters + string.digits
+    return ''.join(random.choice(r_s) for _ in range(5)).upper()
+
+
+async def get_code(s: AsyncSession, exp_style):
+    if exp_style == 'forever':
+        generate = get_random_string
+    else:
+        generate = get_random_num
+    code = await generate()
     while (await s.execute(select(Codes.id).where(Codes.code == code))).scalar():
-        code = random.randint(10000, 99999)
+        code = await generate()
     return str(code)
 
 
 async def get_token(ip, code):
     return hashlib.sha256(f"{ip}{code}{int(time.time() / 1000)}000{settings.ADMIN_PASSWORD}".encode()).hexdigest()
+
+
+async def get_expire_info(expire_style, expire_value, s):
+    now = datetime.datetime.now()
+    if expire_value <= 0 or expire_value > 999:
+        return True, None, None
+    code = await get_code(s, expire_style)
+    if expire_style == 'day':
+        return False, now + datetime.timedelta(days=expire_value), -1, code
+    elif expire_style == 'hour':
+        return False, now + datetime.timedelta(hours=expire_value), -1, code
+    elif expire_style == 'minute':
+        return False, now + datetime.timedelta(minutes=expire_value), -1, code
+    elif expire_style == 'forever':
+        return False, None, -1, code
+    elif expire_style == 'count':
+        return False, None, expire_value, code
+    else:
+        return True, None, None

core/utils.py

@@ -2,12 +2,15 @@
 import datetime
 import uuid
 from pathlib import Path
+
+from pydantic import BaseModel
 from sqlalchemy import select, func, update
 from sqlalchemy.ext.asyncio import AsyncSession
 from starlette.requests import Request
 from starlette.responses import HTMLResponse, FileResponse, RedirectResponse
 from starlette.staticfiles import StaticFiles
-from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files, get_token
+from core.utils import error_ip_limit, upload_ip_limit, get_code, storage, delete_expire_files, get_token, \
+    get_expire_info
 from core.depends import admin_required
 from fastapi import FastAPI, Depends, UploadFile, Form, File, HTTPException, BackgroundTasks, Header
 from core.database import init_models, Options, Codes, get_session
@@ -42,16 +45,22 @@ async def startup(s: AsyncSession = Depends(get_session)):
 @app.get('/')
 async def index():
     return HTMLResponse(
-        index_html.replace('{{title}}', settings.TITLE).replace('{{description}}', settings.DESCRIPTION).replace(
-            '{{keywords}}', settings.KEYWORDS).replace("'{{fileSizeLimit}}'", str(settings.FILE_SIZE_LIMIT))
+        index_html
+        .replace('{{title}}', settings.TITLE)
+        .replace('{{description}}', settings.DESCRIPTION)
+        .replace('{{keywords}}', settings.KEYWORDS)
+        .replace("'{{fileSizeLimit}}'", str(settings.FILE_SIZE_LIMIT))
     )
 
 
 @app.get(f'/{settings.ADMIN_ADDRESS}', description='管理页面')
 async def admin():
     return HTMLResponse(
-        admin_html.replace('{{title}}', settings.TITLE).replace('{{description}}', settings.DESCRIPTION).replace(
-            '{{admin_address}}', settings.ADMIN_ADDRESS).replace('{{keywords}}', settings.KEYWORDS)
+        admin_html
+        .replace('{{title}}', settings.TITLE)
+        .replace('{{description}}', settings.DESCRIPTION)
+        .replace('{{admin_address}}', settings.ADMIN_ADDRESS)
+        .replace('{{keywords}}', settings.KEYWORDS)
     )
 
 
@@ -125,11 +134,10 @@ async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession =
     if not info:
         error_count = settings.ERROR_COUNT - error_ip_limit.add_ip(ip)
         raise HTTPException(status_code=404, detail=f"取件码错误，{error_count}次后将被禁止{settings.ERROR_MINUTE}分钟")
-    if info.exp_time < datetime.datetime.now() or info.count == 0:
+    if (info.exp_time and info.exp_time < datetime.datetime.now()) or info.count == 0:
         raise HTTPException(status_code=404, detail="取件码已失效，请联系寄件人")
     await s.execute(update(Codes).where(Codes.id == info.id).values(count=info.count - 1))
     await s.commit()
-    print(info.type)
     if info.type != 'text':
         info.text = f'/select?code={info.code}&token={await get_token(code, ip)}'
     return {
@@ -140,7 +148,6 @@ async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession =
 
 @app.get('/banner')
 async def banner(request: Request):
-    # 数据库查询config
     return {
         'detail': '查询成功',
         'data': settings.BANNERS,
@@ -191,40 +198,40 @@ async def merge_chunks(file_key: str, file_name: str, total_chunks: int):
     return {'code': 200, 'data': await storage.merge_chunks(file_key, file_name, total_chunks)}
 
 
-@app.post('/share', dependencies=[Depends(admin_required)], description='分享文件')
-async def share(text: str = Form(default=None), size: int = Form(default=0), file_key: str = Form(default=None),
-                _type: str = Form(default=None), name: str = Form(default='text'),
-                style: str = Form(default='2'), value: int = Form(default=1), is_file: int = Form(default=True),
-                ip: str = Depends(upload_ip_limit), s: AsyncSession = Depends(get_session)):
-    if style == '2':
-        if value > settings.MAX_DAYS:
-            raise HTTPException(status_code=400, detail=f"最大有效天数为{settings.MAX_DAYS}天")
-        # 如果天数大于0，就设置过期时间，否则就设置为永久，无过期时间
-        if settings.ENABLE_PERMANENT and value < 0:
-            exp_time = None
-        else:
-            exp_time = datetime.datetime.now() + datetime.timedelta(days=value)
-        exp_count = -1
-    elif style == '1':
-        if value < 1:
-            raise HTTPException(status_code=400, detail="最小有效次数为1次")
-        exp_time = datetime.datetime.now() + datetime.timedelta(days=1)
-        exp_count = value
-    else:
-        exp_time = datetime.datetime.now() + datetime.timedelta(days=1)
-        exp_count = -1
-    print(is_file)
-    if is_file:
-        key = file_key
-        # size = await storage.get_size(file)
-        # if size > settings.FILE_SIZE_LIMIT:
-        #     raise HTTPException(status_code=400, detail="文件过大")
-        _text, size, name = text, size, name
-        # background_tasks.add_task(storage.save_file, file, _text)
-    else:
-        key = uuid.uuid4().hex
-        size, _text, _type, name = len(text), text, 'text', '文本分享'
-    code = await get_code(s)
+class ShareDataModel(BaseModel):
+    text: str
+    size: int = 0
+    exp_style: str
+    exp_value: int
+    type: str
+    name: str
+    key: str = uuid.uuid4().hex
+
+
+@app.post('/share/file/', dependencies=[Depends(admin_required)], description='分享文件')
+async def share_file(file_model: ShareDataModel, s: AsyncSession = Depends(get_session),
+                     ip: str = Depends(error_ip_limit)):
+    exp_error, exp_time, exp_count, code = await get_expire_info(file_model.exp_style, file_model.exp_value, s)
+    if exp_error:
+        raise HTTPException(status_code=400, detail='过期值异常')
+    s.add(Codes(code=code, text=file_model.text, size=file_model.size, type=file_model.type, name=file_model.name,
+                count=exp_count, exp_time=exp_time, key=file_model.key))
+    await s.commit()
+    upload_ip_limit.add_ip(ip)
+    return {
+        'detail': '分享成功，请点击我的文件按钮查看上传列表',
+        'data': {'code': code, 'key': file_model.key, 'name': file_model.name}
+    }
+
+
+@app.post('/share/text/', dependencies=[Depends(admin_required)])
+async def share_text(text_model: ShareDataModel, s: AsyncSession = Depends(get_session),
+                     ip: str = Depends(error_ip_limit)):
+    exp_error, exp_time, exp_count, code = await get_expire_info(text_model.exp_style, text_model.exp_value, s, )
+    if exp_error:
+        raise HTTPException(status_code=400, detail='过期值异常')
+    exp_status, exp_time, exp_count, code = await get_expire_info(text_model.exp_style, text_model.exp_value, s)
+    size, _text, _type, name, key = len(text_model.text), text_model.text, 'text', '文本分享', text_model.key
     s.add(Codes(code=code, text=_text, size=size, type=_type, name=name, count=exp_count, exp_time=exp_time, key=key))
     await s.commit()
     upload_ip_limit.add_ip(ip)