From 68175f958349c44bdb0cb870099f7e2079742ef1 Mon Sep 17 00:00:00 2001
From: lan <hhxk666@gmail.com>
Date: Wed, 14 Dec 2022 09:55:04 +0800
Subject: [PATCH] =?UTF-8?q?update:=E9=99=90=E5=88=B6=E7=AE=A1=E7=90=86?=
 =?UTF-8?q?=E5=91=98=E4=B8=8A=E4=BC=A0=EF=BC=8C=E6=9C=80=E9=95=BF=E5=A4=A9?=
 =?UTF-8?q?=E6=95=B0=E9=85=8D=E7=BD=AE?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 depends.py  | 10 +++++++---
 main.py     | 10 ++++------
 settings.py |  2 ++
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/depends.py b/depends.py
index 0924227..0ca7b15 100644
--- a/depends.py
+++ b/depends.py
@@ -6,9 +6,13 @@ from fastapi import Header, HTTPException, Request
 import settings
 
 
-async def admin_required(pwd: Union[str, None] = Header(default=None)):
-    if not pwd or pwd != settings.ADMIN_PASSWORD:
-        raise HTTPException(status_code=401, detail="密码错误")
+async def admin_required(pwd: Union[str, None] = Header(default=None), request: Request = None):
+    if 'share' in request.url.path:
+        if pwd != settings.ADMIN_PASSWORD and not settings.ENABLE_UPLOAD:
+            raise HTTPException(status_code=403, detail='本站上传功能已关闭，仅管理员可用')
+    else:
+        if not pwd or pwd != settings.ADMIN_PASSWORD:
+            raise HTTPException(status_code=401, detail="密码错误，请重新登录")
 
 
 class IPRateLimit:
diff --git a/main.py b/main.py
index 3d643ca..e7ab99a 100644
--- a/main.py
+++ b/main.py
@@ -121,16 +121,14 @@ async def index(code: str, ip: str = Depends(error_ip_limit), s: AsyncSession =
     }
 
 
-@app.post('/share')
-async def share(background_tasks: BackgroundTasks, pwd: str = Header(default=None), text: str = Form(default=None),
+@app.post('/share', dependencies=[Depends(admin_required)], description='分享文件')
+async def share(background_tasks: BackgroundTasks, text: str = Form(default=None),
                 style: str = Form(default='2'), value: int = Form(default=1), file: UploadFile = File(default=None),
                 ip: str = Depends(upload_ip_limit), s: AsyncSession = Depends(get_session)):
-    if not settings.ENABLE_UPLOAD and pwd != settings.ADMIN_PASSWORD:
-        raise HTTPException(status_code=403, detail="上传功能已关闭")
     code = await get_code(s)
     if style == '2':
-        if value > 7:
-            raise HTTPException(status_code=400, detail="最大有效天数为7天")
+        if value > settings.MAX_DAYS:
+            raise HTTPException(status_code=400, detail=f"最大有效天数为{settings.MAX_DAYS}天")
         exp_time = datetime.datetime.now() + datetime.timedelta(days=value)
         exp_count = -1
     elif style == '1':
diff --git a/settings.py b/settings.py
index ecb74f9..21e9b47 100644
--- a/settings.py
+++ b/settings.py
@@ -17,6 +17,8 @@ DATA_ROOT = './data/' + config('DATA_ROOT', cast=str, default=f"static")
 STATIC_URL = config('STATIC_URL', cast=str, default="/static")
 # 开启上传
 ENABLE_UPLOAD = config('ENABLE_UPLOAD', cast=bool, default=False)
+# 最长天数
+MAX_DAYS = config('MAX_DAYS', cast=int, default=7)
 # 错误次数
 ERROR_COUNT = config('ERROR_COUNT', cast=int, default=5)
 # 错误限制分钟数