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}, }, ), ),