Merge branch 'feat/0.2.2.1' into release

docker/bisheng/config/config.yaml

@@ -4,9 +4,9 @@ database_url:
 
 # 缓存配置  redis://[[username]:[password]]@localhost:6379/0
 redis_url: "redis://redis:6379/0"
-redis_host: [("redis")]
-redis_master:
-redis_password:
+# redis_host: [("redis")]
+# redis_master:
+# redis_password:
 
 environment:
   env: dev

src/backend/bisheng/api/services/captcha.py

@@ -0,0 +1,11 @@
+from bisheng.cache.redis import redis_client
+
+
+async def verify_captcha(captcha: str, captcha_key: str):
+    # check captcha
+    captcha_value = redis_client.get(captcha_key)
+    if captcha_value:
+        redis_client.delete(captcha_key)
+        return captcha_value.lower() == captcha.lower()
+    else:
+        return False

src/backend/bisheng/api/v1/chat.py

@@ -118,13 +118,20 @@ async def chat(
         *,
         flow_id: str,
         websocket: WebSocket,
+        t: Optional[str] = None,
         chat_id: Optional[str] = None,
         Authorize: AuthJWT = Depends(),
 ):
     """Websocket endpoint for chat."""
     try:
-        Authorize.jwt_required(auth_from='websocket', websocket=websocket)
-        payload = json.loads(Authorize.get_jwt_subject())
+        if t:
+            Authorize.jwt_required(auth_from='websocket', token=t)
+            Authorize._token = t
+        else:
+            Authorize.jwt_required(auth_from='websocket', websocket=websocket)
+
+        payload = Authorize.get_jwt_subject()
+        payload = json.loads(payload)
         user_id = payload.get('user_id')
         if chat_id:
             with next(get_session()) as session:

src/backend/bisheng/api/v1/schemas.py

@@ -9,6 +9,11 @@
 from pydantic import BaseModel, Field, validator
 
 
+class CaptchaInput(BaseModel):
+    captcha_key: str
+    captcha: str
+
+
 class ChunkInput(BaseModel):
     knowledge_id: int
     documents: List[Document]

src/backend/bisheng/api/v1/skillcenter.py

@@ -37,7 +37,7 @@ def create_template(*, session: Session = Depends(get_session), template: Templa
     return resp_200(db_template)
 
 
-@router.get('/template/', response_model=UnifiedResponseModel[list[Template]], status_code=200)
+@router.get('/template', response_model=UnifiedResponseModel[list[Template]], status_code=200)
 def read_template(*,
                   page_size: Optional[int] = None,
                   page_name: Optional[int] = None,
@@ -47,15 +47,16 @@ def read_template(*,
     """Read all flows."""
     sql = select(Template.id, Template.name, Template.description, Template.update_time)
     if id:
-        sql = select(Template).where(Template.id == id)
+        template = session.get(Template, id)
+        return resp_200([template])
     if name:
         sql.where(Template.name == name)
     sql.order_by(Template.order_num.desc())
     if page_size and page_name:
         sql.offset(page_size * (page_name - 1)).limit(page_size)
-
     try:
-        templates = session.exec(sql).all()
+        template_session = session.exec(sql)
+        templates = template_session.mappings().all()
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e)) from e
     return resp_200(templates)

src/backend/bisheng/api/v1/user.py

@@ -1,17 +1,28 @@
 import hashlib
 import json
+import random
+import string
+import uuid
+from base64 import b64decode, b64encode
+from io import BytesIO
 from typing import Optional
 from uuid import UUID
 
+import rsa
+from bisheng.api.services.captcha import verify_captcha
 from bisheng.api.v1.schemas import UnifiedResponseModel, resp_200
+from bisheng.cache.redis import redis_client
 from bisheng.database.base import get_session
 from bisheng.database.models.flow import Flow
 from bisheng.database.models.knowledge import Knowledge
 from bisheng.database.models.role import Role, RoleCreate, RoleUpdate
 from bisheng.database.models.role_access import AccessType, RoleAccess, RoleRefresh
 from bisheng.database.models.user import User, UserCreate, UserLogin, UserRead, UserUpdate
 from bisheng.database.models.user_role import UserRole, UserRoleCreate
+from bisheng.settings import settings
+from bisheng.utils.constants import CAPTCHA_PREFIX, RSA_KEY
 from bisheng.utils.logger import logger
+from captcha.image import ImageCaptcha
 from fastapi import APIRouter, Depends, HTTPException
 from fastapi.encoders import jsonable_encoder
 from fastapi.security import OAuth2PasswordBearer
@@ -27,7 +38,12 @@
 
 @router.post('/user/regist', response_model=UnifiedResponseModel[UserRead], status_code=201)
 async def regist(*, session: Session = Depends(get_session), user: UserCreate):
-    db_user = User.from_orm(user)
+    # 验证码校验
+    if settings.get_from_db('use_captcha'):
+        if not user.captcha_key or not await verify_captcha(user.captcha, user.captcha_key):
+            raise HTTPException(status_code=500, detail='验证码错误')
+
+    db_user = User.model_validate(user)
     # check if admin user
     admin = session.exec(select(User).where(User.user_id == 1)).all()
     if not admin:
@@ -40,7 +56,12 @@ async def regist(*, session: Session = Depends(get_session), user: UserCreate):
         raise HTTPException(status_code=500, detail='账号已存在')
     else:
         try:
-            db_user.password = md5_hash(user.password)
+            if value := redis_client.get(RSA_KEY):
+                private_key = value[1]
+                db_user.password = md5_hash(
+                    rsa.decrypt(b64decode(user.password), private_key).decode('utf-8'))
+            else:
+                db_user.password = md5_hash(user.password)
             session.add(db_user)
             session.flush()
             session.refresh(db_user)
@@ -60,8 +81,18 @@ async def login(*,
                 session: Session = Depends(get_session),
                 user: UserLogin,
                 Authorize: AuthJWT = Depends()):
+    # 验证码校验
+    if settings.get_from_db('use_captcha'):
+        if not user.captcha_key or not await verify_captcha(user.captcha, user.captcha_key):
+            raise HTTPException(status_code=500, detail='验证码错误')
+
     # check if user already exist
-    password = md5_hash(user.password)
+    if value := redis_client.get(RSA_KEY):
+        private_key = value[1]
+        password = md5_hash(rsa.decrypt(b64decode(user.password), private_key).decode('utf-8'))
+    else:
+        password = md5_hash(user.password)
+
     db_user = session.exec(
         select(User).where(User.user_name == user.user_name, User.password == password)).first()
     if db_user:
@@ -457,6 +488,49 @@ async def flow_list(*,
     })
 
 
+@router.get('/user/get_captcha', status_code=200)
+async def get_captcha():
+    # generate captcha
+    chr_all = string.ascii_letters + string.digits
+    chr_4 = ''.join(random.sample(chr_all, 4))
+    image = ImageCaptcha().generate_image(chr_4)
+    # 对image 进行base 64 编码
+    buffered = BytesIO()
+    image.save(buffered, format='PNG')
+
+    capthca_b64 = b64encode(buffered.getvalue()).decode()
+    logger.info('get_captcha captcha_char={}', chr_4)
+    # generate key, 生成简单的唯一id，
+    key = CAPTCHA_PREFIX + uuid.uuid4().hex[:8]
+    redis_client.set(key, chr_4, expiration=300)
+
+    # 增加配置，是否必须使用验证码
+    return resp_200({
+        'captcha_key': key,
+        'captcha': capthca_b64,
+        'user_capthca': settings.get_from_db('use_captcha') or False
+    })
+
+
+@router.get('/user/public_key', status_code=200)
+async def get_rsa_publish_key():
+    # redis 存储
+    key = RSA_KEY
+    # redis lock
+    if redis_client.setNx(key, 1):
+        # Generate a key pair
+        (pubkey, privkey) = rsa.newkeys(512)
+
+        # Save the keys to strings
+        redis_client.set(key, (pubkey, privkey), 3600)
+    else:
+        pubkey, privkey = redis_client.get(key)
+
+    pubkey_str = pubkey.save_pkcs1().decode()
+
+    return resp_200({'public_key': pubkey_str})
+
+
 def md5_hash(string):
     md5 = hashlib.md5()
     md5.update(string.encode('utf-8'))

src/backend/bisheng/api/v2/rpc.py

@@ -5,7 +5,7 @@
 from bisheng.database.models.user import User
 from bisheng.database.models.user_role import UserRole
 from bisheng.settings import settings
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, HTTPException
 from fastapi.responses import RedirectResponse
 from fastapi_jwt_auth import AuthJWT
 from loguru import logger
@@ -25,45 +25,38 @@ def set_cookie(*,
                session: Session = Depends(get_session),
                Authorize: AuthJWT = Depends()):
     """设置默认"""
-    if deptId:
-        # this interface should update user model, and now the main ref don't mathes
-        db_user = session.exec(select(User).where(User.dept_id == deptId)).first()
-        if not db_user:
-            db_user = User(user_name=deptName, password='none', dept_id=deptId)
-            session.add(db_user)
-            session.flush()
-            db_user_role = UserRole(user_id=db_user.user_id, role_id=2)
-            session.add(db_user_role)
-            session.commit()
-            session.refresh(db_user)
-    if not deptId:
-        dept_user_id = settings.get_from_db('default_operator').get('user')
-    else:
-        dept_user_id = db_user.user_id
 
-    payload = {'user_name': deptName, 'user_id': dept_user_id, 'role': [2]}
     # admin
-    role_admin = session.query(UserRole).where(UserRole.user_id == user_id,
-                                               UserRole.role_id == 1).first()
     try:
-        if user_id and role_id == 1:
-            if not role_admin:
-                # keep
-                admin_user = session.get(User, user_id)
-                if not admin_user:
-                    db_user = User(id=user_id, user_name=deptName, password='none', dept_id=deptId)
-                    session.add(db_user)
-                db_user_role = UserRole(user_id=user_id, role_id=1)
+        if deptId:
+            # this interface should update user model, and now the main ref don't mathes
+            db_user = session.exec(select(User).where(User.dept_id == deptId)).first()
+            if not db_user:
+                db_user = User(user_name=deptName, password='none', dept_id=deptId)
+                session.add(db_user)
+                session.flush()
+                db_user_role = UserRole(user_id=db_user.user_id, role_id=2)
+                session.add(db_user_role)
+                session.commit()
+                session.refresh(db_user)
+        else:
+            raise ValueError('deptId 必须传递')
+        payload = {'user_name': deptName, 'user_id': db_user.user_id, 'role': [2]}
+        if role_id == 1:
+            admin_user = session.query(User).where(User.user_name == 'root').first()
+            if not admin_user:
+                admin_user = User(user_name='root', password='none')
+                session.add(admin_user)
+                session.flush()
+                session.refresh(admin_user)
+                db_user_role = UserRole(user_id=admin_user.user_id, role_id=1)
                 session.add(db_user_role)
-            payload = {'user_name': deptName, 'user_id': user_id, 'role': 'admin'}
-        elif user_id:
-            if role_admin:
-                # delete
-                session.delete(role_admin)
-        session.commit()
+            payload = {'user_name': 'root', 'user_id': admin_user.user_id, 'role': 'admin'}
+            session.commit()
     except Exception as e:
         logger.error(str(e))
         session.rollback()
+        return HTTPException(status_code=500, detail=str(e))
 
     # Create the tokens and passing to set_access_cookies or set_refresh_cookies
     access_token = Authorize.create_access_token(subject=json.dumps(payload), expires_time=864000)

src/backend/bisheng/database/base.py

@@ -68,7 +68,7 @@ def session_getter():
         session = Session(db_service.engine)
         yield session
     except Exception as e:
-        print('Session rollback because of exception:', e)
+        logger.info('Session rollback because of exception:{}', e)
         session.rollback()
         raise
     finally:

src/backend/bisheng/database/models/user.py

@@ -47,10 +47,14 @@ class UserQuery(UserBase):
 class UserLogin(UserBase):
     password: str
     user_name: str
+    captcha_key: Optional[str]
+    captcha: Optional[str]
 
 
 class UserCreate(UserBase):
     password: str
+    captcha_key: Optional[str]
+    captcha: Optional[str]
 
 
 class UserUpdate(SQLModelSerializable):

src/backend/bisheng/utils/constants.py

@@ -51,6 +51,10 @@ def python_function(text: str) -> str:
 # 用来记录 preset questions, dict: key 为作用对象的id
 PRESET_QUESTION = 'preset_question'
 
+# redis key
+CAPTCHA_PREFIX = 'cap_'
+RSA_KEY = 'rsa_'
+
 LOADERS_INFO: List[Dict[str, Any]] = [
     {
         'loader': 'AirbyteJSONLoader',

src/backend/bisheng/utils/threadpool.py

@@ -92,7 +92,8 @@ async def as_completed(self) -> List[Tuple[str, concurrent.futures.Future]]:
                             pending_count += 1
                 if len(lf) == 0:
                     self.async_task.pop(k)
-            logger.info('async_wait_count={}', pending_count)
+            if pending_count > 0:
+                logger.info('async_wait_count={}', pending_count)
             return completed_futures
 
     # async def async_done_callback(self, future):

src/backend/pyproject.toml

@@ -23,6 +23,8 @@ bisheng_pyautogen = "0.1.19"
 minio = "7.2.0"
 loguru = "^0.7.1"
 fastapi_jwt_auth = "^0.5.0"
+captcha = “^0.0.5"
+rsa = ""
 pydantic = "^1.10.13"
 numexpr = "^2.8.6"
 pypinyin = "0.50.0"

src/frontend/package.json

@@ -1,6 +1,6 @@
 {
   "name": "bisheng",
-  "version": "0.2.2",
+  "version": "0.2.2.1",
   "private": true,
   "dependencies": {
     "@emotion/react": "^11.11.1",
@@ -39,6 +39,7 @@
     "esbuild": "^0.17.19",
     "i18next": "^23.6.0",
     "i18next-http-backend": "^2.2.2",
+    "jsencrypt": "^3.3.2",
     "lodash-es": "^4.17.21",
     "lucide-react": "^0.233.0",
     "pdfjs-dist": "3.10.111",

src/frontend/public/locales/en/bs.json

@@ -14,9 +14,11 @@
         "pleaseEnterAccount": "Please enter your account",
         "pleaseEnterPassword": "Please enter your password",
         "accountTooShort": "Account is too short",
-        "passwordTooShort": "Password is too short, at least six characters",
+        "passwordTooShort": "Password is too short, at least 8 characters",
+        "passwordError": "The password must include upper and lower case letters, numbers, and special characters!",
         "passwordMismatch": "Passwords do not match",
-        "registrationSuccess": "Registration successful. Please enter your password to log in"
+        "registrationSuccess": "Registration successful. Please enter your password to log in",
+        "pleaseEnterCaptcha": "Please enter captcha"
     },
     "menu": {
         "app": "App",