1 files changed, 38 insertions, 0 deletions
diff --git a/ATRI/plugins/console/driver/depends.py b/ATRI/plugins/console/driver/depends.py
new file mode 100644
index 0000000..8205f44
--- /dev/null
+++ b/ATRI/plugins/console/driver/depends.py
@@ -0,0 +1,38 @@
+from typing import Union
+from datetime import datetime
+
+from fastapi import Query, HTTPException, status
+from starlette.websockets import WebSocket
+
+from ..data_source import AuthDealer
+
+
+def http_author(token: Union[str, None] = Query(default=None)):
+    data = AuthDealer.get()
+    if data is None:
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="验证信息不存在")
+
+    now_time = datetime.now().timestamp()
+    if token != data.token:
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="密钥不匹配, 请检查")
+    elif now_time > data.dead_time:
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="密钥已过期")
+    else:
+        return token
+
+
+async def websocket_author(
+    websocket: WebSocket, token: Union[str, None] = Query(default=None)
+):
+    data = AuthDealer.get()
+    if not data:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+        return
+
+    now_time = datetime.now().timestamp()
+    if token != data.token:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+    elif now_time > data.dead_time:
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
+    else:
+        return token