From fb723878cae1009a6ec69dbd0dd29b47eacdc268 Mon Sep 17 00:00:00 2001
From: xtaodada <xtao@xtaolink.cn>
Date: Sun, 22 Sep 2024 21:53:36 +0800
Subject: [PATCH] :ambulance: hotfix: change web token to cookie

---
 pagermaid/web/api/login.py       |  5 ++++-
 pagermaid/web/api/status.py      | 25 ++++++++-----------------
 pagermaid/web/api/utils.py       |  9 +++++----
 pagermaid/web/pages/home_page.py |  3 ---
 4 files changed, 17 insertions(+), 25 deletions(-)

diff --git a/pagermaid/web/api/login.py b/pagermaid/web/api/login.py
index f41d960..7f8d006 100644
--- a/pagermaid/web/api/login.py
+++ b/pagermaid/web/api/login.py
@@ -20,9 +20,12 @@ route = APIRouter()
 async def login(user: UserModel):
     if not Config.WEB_SECRET_KEY or user.password == Config.WEB_SECRET_KEY:
         token = create_token()
-        return {
+        data = {
             "status": 0,
             "msg": "登录成功",
             "data": {"version": pgm_version_code, "token": token},
         }
+        response = JSONResponse(content=data)
+        response.set_cookie(key="token_ck", value=token, expires=1800)
+        return response
     return {"status": -100, "msg": "登录失败，请重新输入密钥"}
diff --git a/pagermaid/web/api/status.py b/pagermaid/web/api/status.py
index 48e9c74..deeadda 100644
--- a/pagermaid/web/api/status.py
+++ b/pagermaid/web/api/status.py
@@ -1,22 +1,19 @@
 import asyncio
-from typing import Union, Optional
+from typing import Union
 
-from fastapi import APIRouter, Header
+from fastapi import APIRouter
 from fastapi.responses import JSONResponse, StreamingResponse
 
 from pagermaid.common.status import get_status
 from pagermaid.common.system import run_eval
-from pagermaid.config import Config
 from pagermaid.utils import execute
 from pagermaid.web.api.utils import authentication
 
 route = APIRouter()
 
 
-@route.get("/log")
-async def get_log(token: Optional[str] = Header(...), num: Union[int, str] = 100):
-    if token != Config.WEB_SECRET_KEY:
-        return "非法请求"
+@route.get("/log", dependencies=[authentication()])
+async def get_log(num: Union[int, str] = 100):
     try:
         num = int(num)
     except ValueError:
@@ -31,11 +28,8 @@ async def get_log(token: Optional[str] = Header(...), num: Union[int, str] = 100
     return StreamingResponse(streaming_logs())
 
 
-@route.get("/run_eval")
-async def run_cmd(token: Optional[str] = Header(...), cmd: str = ""):
-    if token != Config.WEB_SECRET_KEY:
-        return "非法请求"
-
+@route.get("/run_eval", dependencies=[authentication()])
+async def run_cmd(cmd: str = ""):
     async def run_cmd_func():
         result = (await run_eval(cmd)).split("\n")
         for i in result:
@@ -45,11 +39,8 @@ async def run_cmd(token: Optional[str] = Header(...), cmd: str = ""):
     return StreamingResponse(run_cmd_func()) if cmd else "无效命令"
 
 
-@route.get("/run_sh")
-async def run_sh(token: Optional[str] = Header(...), cmd: str = ""):
-    if token != Config.WEB_SECRET_KEY:
-        return "非法请求"
-
+@route.get("/run_sh", dependencies=[authentication()])
+async def run_sh(cmd: str = ""):
     async def run_sh_func():
         result = (await execute(cmd)).split("\n")
         for i in result:
diff --git a/pagermaid/web/api/utils.py b/pagermaid/web/api/utils.py
index d30cf78..217422d 100644
--- a/pagermaid/web/api/utils.py
+++ b/pagermaid/web/api/utils.py
@@ -1,7 +1,7 @@
 import datetime
 from typing import Optional
 
-from fastapi import Header, HTTPException, Depends
+from fastapi import Header, HTTPException, Depends, Cookie
 from jose import jwt
 
 from pagermaid.config import Config
@@ -11,12 +11,13 @@ TOKEN_EXPIRE_MINUTES = 30
 
 
 def authentication():
-    def inner(token: Optional[str] = Header(None)):
+    def inner(token: Optional[str] = Header(None), token_ck: str = Cookie(None)):
+        _token = token or token_ck
         if Config.WEB_SECRET_KEY:
-            if token == Config.WEB_SECRET_KEY:
+            if _token == Config.WEB_SECRET_KEY:
                 return
             try:
-                jwt.decode(token, Config.WEB_SECRET_KEY, algorithms=ALGORITHM)
+                jwt.decode(_token, Config.WEB_SECRET_KEY, algorithms=ALGORITHM)
             except (jwt.JWTError, jwt.ExpiredSignatureError, AttributeError):
                 raise HTTPException(
                     status_code=400, detail="登录验证失败或已失效，请重新登录"
diff --git a/pagermaid/web/pages/home_page.py b/pagermaid/web/pages/home_page.py
index 84f2b17..10fa170 100644
--- a/pagermaid/web/pages/home_page.py
+++ b/pagermaid/web/pages/home_page.py
@@ -43,7 +43,6 @@ log_page = Log(
     source={
         "method": "get",
         "url": "/pagermaid/api/log?num=${log_num | raw}",
-        "headers": {"token": Config.WEB_SECRET_KEY},
     },
 )
 
@@ -69,7 +68,6 @@ cmd_input = Form(
                         source={
                             "method": "get",
                             "url": "/pagermaid/api/run_sh?cmd=${command | raw}",
-                            "headers": {"token": Config.WEB_SECRET_KEY},
                         },
                     ),
                 ),
@@ -99,7 +97,6 @@ eval_input = Form(
                         source={
                             "method": "get",
                             "url": "/pagermaid/api/run_eval?cmd=${command | raw}",
-                            "headers": {"token": Config.WEB_SECRET_KEY},
                         },
                     ),
                 ),