UPDATE WEB UI

plugins/web_ui/auth/__init__.py

@@ -1,74 +1,32 @@
 import json
-from datetime import datetime, timedelta
-from typing import Optional
+from datetime import timedelta
 
 import nonebot
-from fastapi import Depends, HTTPException
-from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm
-from jose import JWTError, jwt
-from pydantic import BaseModel
-from starlette import status
+from fastapi import APIRouter, Depends
+from fastapi.security import OAuth2PasswordRequestForm
 
-from configs.config import Config
-from configs.path_config import DATA_PATH
-
-from ..config import Result, router
+from ..models.model import Result
+from ..utils import (
+    ACCESS_TOKEN_EXPIRE_MINUTES,
+    create_token,
+    get_user,
+    token_data,
+    token_file,
+)
 
 app = nonebot.get_app()
 
 
-SECRET_KEY = "09d25e094faa6ca2556c818166b7a9563b93f7099f6f0f4caa6cf63b88e8d3e7"
-ALGORITHM = "HS256"
-ACCESS_TOKEN_EXPIRE_MINUTES = 30
-
-oauth2_scheme = OAuth2PasswordBearer(tokenUrl="api/login")
-
-
-token_file = DATA_PATH / "web_ui" / "token.json"
-token_file.parent.mkdir(parents=True, exist_ok=True)
-token_data = {"token": []}
-if token_file.exists():
-    token_data = json.load(open(token_file, "r", encoding="utf8"))
-
-
-class User(BaseModel):
-    username: str
-    password: str
-
-
-class Token(BaseModel):
-    access_token: str
-    token_type: str
-
-
-def get_user(uname: str) -> Optional[User]:
-    username = Config.get_config("web-ui", "username")
-    password = Config.get_config("web-ui", "password")
-    if username and password and uname == username:
-        return User(username=username, password=password)
-
-
-form_exception = HTTPException(
-    status_code=status.HTTP_401_UNAUTHORIZED,
-    detail="Could not validate credentials",
-    headers={"WWW-Authenticate": "Bearer"},
-)
-
-
-def create_token(user: User, expires_delta: Optional[timedelta] = None):
-    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
-    return jwt.encode(
-        claims={"sub": user.username, "exp": expire},
-        key=SECRET_KEY,
-        algorithm=ALGORITHM,
-    )
+router = APIRouter()
 
 
 @router.post("/login")
 async def login_get_token(form_data: OAuth2PasswordRequestForm = Depends()):
-    user = get_user(form_data.username)
-    if not user or user.password != form_data.password:
-        raise form_exception
+    if user := get_user(form_data.username):
+        if user.password != form_data.password:
+            return Result.fail("真笨, 密码都能记错!", 999)
+    else:
+        return Result.fail("你滴配置文件里用户名密码配置项为空", 998)
     access_token = create_token(
         user=user, expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
     )
@@ -77,31 +35,6 @@ async def login_get_token(form_data: OAuth2PasswordRequestForm = Depends()):
         token_data["token"] = token_data["token"][1:]
     with open(token_file, "w", encoding="utf8") as f:
         json.dump(token_data, f, ensure_ascii=False, indent=4)
-    return {"access_token": access_token, "token_type": "bearer"}
-
-
-credentials_exception = HTTPException(
-    status_code=status.HTTP_401_UNAUTHORIZED,
-    detail="Could not validate credentials",
-    headers={"WWW-Authenticate": "Bearer"},
-)
-
-
-@app.post("/auth")
-def token_to_user(token: str = Depends(oauth2_scheme)):
-    if token not in token_data["token"]:
-        try:
-            payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
-            username, expire = payload.get("sub"), payload.get("exp")
-            user = get_user(username)  # type: ignore
-            if user is None:
-                raise JWTError
-        except JWTError:
-            return Result(code=401)
-    return Result(code=200, info="登录成功")
-
-
-if __name__ == "__main__":
-    import uvicorn
-
-    uvicorn.run(app, host="127.0.0.1", port=8080)
+    return Result.ok(
+        {"access_token": access_token, "token_type": "bearer"}, "欢迎回家, 欧尼酱!"
+    )