feat(oj): 实现oj的用户管理

example.env

@@ -0,0 +1 @@
+SECRET="THISISASECRET"

main.py

@@ -1,8 +1,12 @@
+from manager import cli_app, start_server_background
+from ccf_parser import CCF
 import json
 import pathlib
 
-from ccf_parser import CCF
-from manager import cli_app, start_server_background
+import dotenv
+
+dotenv.load_dotenv()
+
 
 if '__main__' == __name__:
     cli_app()

manager/cli/base.py

@@ -1,4 +1,5 @@
 import shutil
+import time
 import zipfile
 from pathlib import Path
 from typing import *
@@ -13,6 +14,7 @@
 from rich.progress import track
 from rich.text import Text
 
+from online_judge import start_oj_background
 from utils import manager_logger
 
 from ..base import _start_server, start_server_background
@@ -90,8 +92,16 @@ def intro():
 
 
 @app.command(name='server')
-def start_server_command():  # pragma: no cover
-    download_ited()
+def start_server_command(manager: bool = True, oj: bool = True):  # pragma: no cover
+    if manager:
+        download_ited()
 
-    manager_logger.info('访问 http://localhost:2568/editor 以访问ItsWA Manager。')
-    _start_server()
+        manager_logger.info(
+            '访问 http://localhost:2568/editor 以访问ItsWA Manager。')
+        start_server_background()
+
+    if oj:
+        start_oj_background()
+
+    while True:
+        time.sleep(10**9)

online_judge/__init__.py

@@ -0,0 +1,2 @@
+from .base import app as oj_app
+from .base import start_oj_background

online_judge/auth.py

@@ -0,0 +1,95 @@
+import datetime
+import os
+from typing import *
+
+import fastapi
+import jwt
+import pydantic
+from fastapi import APIRouter, Body, Depends, HTTPException
+from fastapi.security import APIKeyCookie
+from tinydb import Query
+
+from .models.user import User
+from .utils import usercol
+
+router = APIRouter(prefix='/auth', tags=['用户验证'])
+apikey_schema = APIKeyCookie(name='itswa-oj-apikey')
+
+
+def get_apikey_decoded(apikey: Optional[str] = Depends(apikey_schema)) -> Dict[Any, Any]:
+    if not apikey:
+        raise HTTPException(status_code=401, detail="请提供API Key")
+
+    try:
+        decoded = jwt.decode(apikey, algorithms=['HS256'])
+    except jwt.DecodeError:
+        raise HTTPException(status_code=401, detail="API Key无效")
+    except jwt.ExpiredSignatureError:
+        raise HTTPException(status_code=401, detail="API Key已过期")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"未知错误: {e}")
+
+    return decoded
+
+
+def get_user(decoded: Dict[Any, Any] = Depends(get_apikey_decoded)) -> User:
+    try:
+        decoded: User = User(**decoded)
+    except pydantic.ValidationError:
+        raise HTTPException(status_code=401, detail="API Key无效")
+
+    User_Query = Query()
+    result = usercol.search(User_Query.username == decoded.username)[0]
+
+    return User(**result)
+
+
+def get_token(user: User) -> str:
+    return jwt.encode(
+        {
+            **user.model_dump(mode='json'),
+            'exp': datetime.datetime.now() + datetime.timedelta(days=7)
+        },
+        os.environ['SECRET']
+    )
+
+
+@router.post('/login', name='登录', responses={
+    200: {
+        "description": "登录成功",
+        "content": {
+            "application/json": {
+                "example": {
+                    'token': 'user_token'
+                }
+            }
+        }
+    }
+})
+async def user_login(username: Annotated[str, Body()], password: Annotated[str, Body()]):
+    User_Query = Query()
+    results = usercol.search(User_Query.username ==
+                             username and User_Query.password == password)
+
+    if results.__len__() >= 1:
+        return {
+            'token': get_token(User.model_validate(results[0]))
+        }
+    else:
+        raise HTTPException(status_code=401, detail="用户名或密码错误")
+
+
+@router.post('/register', name='注册', response_model=User)
+async def user_register(username: Annotated[str, Body()], password: Annotated[str, Body()]):
+    User_Query = Query()
+
+    if usercol.search(User_Query.username == username).__len__() >= 1:
+        raise HTTPException(status_code=409, detail="用户名已存在")
+
+    usercol.insert(User(
+        username=username,
+        password=password,
+        role='default'
+    ).model_dump(mode='json'))
+
+    return usercol.search(User_Query.username == username)[0]

online_judge/base.py

@@ -0,0 +1,26 @@
+import multiprocessing
+
+import fastapi
+import uvicorn
+
+from utils import online_judge_logger as logger
+
+from .auth import router as auth_router
+from .contests import router as contests_router
+
+app = fastapi.FastAPI(title='ItsWA Online Judge API')
+app.include_router(contests_router)
+app.include_router(auth_router)
+
+
+def _start_oj():  # proagma: no cover
+    logger.info('Online Judge API 启动, 地址 http://0.0.0.0:6572/')
+    uvicorn.run('online_judge:oj_app', host="0.0.0.0", port=6572,
+                workers=6, log_level='warning')
+
+
+def start_oj_background():  # pragma: no cover
+    process = multiprocessing.Process(target=_start_oj)
+    process.start()
+
+    return process

online_judge/contests.py

@@ -0,0 +1,31 @@
+import json
+from pathlib import Path
+from typing import *
+
+import fastapi
+from fastapi import APIRouter, HTTPException
+
+import ccf_parser
+
+router = APIRouter(prefix='/contests', tags=['比赛'])
+
+
+@router.get('/', response_model=List[ccf_parser.CCF])
+async def get_contests():
+    contest_indexes_path = Path('./config/contests.json')
+    contest_indexes = [
+        ccf_parser.ContestIndex.model_validate(x)
+        for x in json.loads(contest_indexes_path.read_text('utf-8'))
+    ]
+
+    ccfs: List[ccf_parser.CCF] = [
+        ccf_parser.CCF.model_validate_json(
+            x.ccf_file.joinpath('ccf.json').read_text('utf-8'))
+        for x in contest_indexes
+    ]
+
+    # 抹除题目数据
+    for ccf in ccfs:
+        ccf.contest.problems = []
+
+    return ccfs

online_judge/models/user.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+
+class User(BaseModel):
+    username: str
+    password: str
+    role: str

online_judge/utils/__init__.py

@@ -0,0 +1 @@
+from .database import db, usercol

online_judge/utils/database.py

@@ -0,0 +1,6 @@
+import tinydb
+from tinydb import TinyDB
+
+db = TinyDB('./assets/oj_db.json', indent=4,
+            ensure_ascii=False, sort_keys=True)
+usercol = db.table('users')

pyproject.toml

@@ -18,6 +18,8 @@ requests = "^2.31.0"
 pytest-html = "^4.1.1"
 dominate = "^2.9.1"
 tinydb = "^4.8.0"
+pyjwt = "^2.8.0"
+python-dotenv = "^1.0.1"
 
 
 [[tool.poetry.source]]