Sfoglia il codice sorgente

add option to prompt for a validation password when initializing admin user (#2302)

Chenhe Gu 1 anno fa
parent
commit
09acf215f0

+ 10 - 0
api/controllers/console/error.py

@@ -13,6 +13,16 @@ class NotSetupError(BaseHTTPException):
                   "Please proceed with the initialization and installation process first."
     code = 401
 
+class NotInitValidateError(BaseHTTPException):
+    error_code = 'not_init_validated'
+    description = "Init validation has not been completed yet. " \
+                  "Please proceed with the init validation process first."
+    code = 401
+
+class InitValidateFailedError(BaseHTTPException):
+    error_code = 'init_validate_failed'
+    description = "Init validation failed. Please check the password and try again."
+    code = 401
 
 class AccountNotLinkTenantError(BaseHTTPException):
     error_code = 'account_not_link_tenant'

+ 47 - 0
api/controllers/console/init_validate.py

@@ -0,0 +1,47 @@
+import os
+from flask import current_app, session
+from flask_restful import Resource, reqparse
+from libs.helper import str_len
+from models.model import DifySetup
+from services.account_service import TenantService
+
+from . import api
+from .error import AlreadySetupError, InitValidateFailedError
+from .wraps import only_edition_self_hosted
+
+
+class InitValidateAPI(Resource):
+
+    def get(self):
+        init_status = get_init_validate_status()
+        if init_status:
+            return { 'status': 'finished' }
+        return {'status': 'not_started' }
+
+    @only_edition_self_hosted
+    def post(self):
+        # is tenant created
+        tenant_count = TenantService.get_tenant_count()
+        if tenant_count > 0:
+            raise AlreadySetupError()
+
+        parser = reqparse.RequestParser()
+        parser.add_argument('password', type=str_len(30),
+                            required=True, location='json')
+        input_password = parser.parse_args()['password']
+
+        if input_password != os.environ.get('INIT_PASSWORD'):
+            session['is_init_validated'] = False
+            raise InitValidateFailedError()
+            
+        session['is_init_validated'] = True
+        return {'result': 'success'}, 201
+
+def get_init_validate_status():
+    if current_app.config['EDITION'] == 'SELF_HOSTED':
+        if os.environ.get('INIT_PASSWORD'):
+            return session.get('is_init_validated') or DifySetup.query.first()
+    
+    return True
+
+api.add_resource(InitValidateAPI, '/init')

+ 10 - 3
api/controllers/console/setup.py

@@ -10,7 +10,8 @@ from models.model import DifySetup
 from services.account_service import AccountService, RegisterService, TenantService
 
 from . import api
-from .error import AlreadySetupError, NotSetupError
+from .error import AlreadySetupError, NotSetupError, NotInitValidateError
+from .init_validate import get_init_validate_status
 from .wraps import only_edition_self_hosted
 
 
@@ -24,7 +25,7 @@ class SetupApi(Resource):
                     'step': 'finished',
                     'setup_at': setup_status.setup_at.isoformat()
                 }
-            return {'step': 'not_start'}
+            return {'step': 'not_started'}
         return {'step': 'finished'}
 
     @only_edition_self_hosted
@@ -37,6 +38,9 @@ class SetupApi(Resource):
         tenant_count = TenantService.get_tenant_count()
         if tenant_count > 0:
             raise AlreadySetupError()
+    
+        if not get_init_validate_status():
+            raise NotInitValidateError()
 
         parser = reqparse.RequestParser()
         parser.add_argument('email', type=email,
@@ -71,7 +75,10 @@ def setup_required(view):
     @wraps(view)
     def decorated(*args, **kwargs):
         # check setup
-        if not get_setup_status():
+        if not get_init_validate_status():
+            raise NotInitValidateError()
+        
+        elif not get_setup_status():
             raise NotSetupError()
 
         return view(*args, **kwargs)

+ 3 - 0
docker/docker-compose.yaml

@@ -15,6 +15,9 @@ services:
       # different from api or web app domain.
       # example: http://cloud.dify.ai
       CONSOLE_WEB_URL: ''
+      # Password for admin user initialization.
+      # If left unset, admin user will not be prompted for a password when creating the initial admin account.
+      INIT_PASSWORD: ''
       # The base URL of console application api server, refers to the Console base URL of WEB service if console domain is
       # different from api or web app domain.
       # example: http://cloud.dify.ai

+ 82 - 0
web/app/init/InitPasswordPopup.tsx

@@ -0,0 +1,82 @@
+'use client'
+import { useEffect, useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import { useRouter } from 'next/navigation'
+import Toast from '../components/base/toast'
+import Loading from '../components/base/loading'
+import Button from '@/app/components/base/button'
+import { fetchInitValidateStatus, initValidate } from '@/service/common'
+import type { InitValidateStatusResponse } from '@/models/common'
+
+const InitPasswordPopup = () => {
+  const [password, setPassword] = useState('')
+  const [loading, setLoading] = useState(true)
+  const [validated, setValidated] = useState(false)
+  const router = useRouter()
+
+  const { t } = useTranslation()
+
+  const handleValidation = async () => {
+    setLoading(true)
+    try {
+      const response = await initValidate({ body: { password } })
+      if (response.result === 'success') {
+        setValidated(true)
+        router.push('/install') // or render setup form
+      }
+      else {
+        throw new Error('Validation failed')
+      }
+    }
+    catch (e: any) {
+      Toast.notify({
+        type: 'error',
+        message: e.message,
+        duration: 5000,
+      })
+      setLoading(false)
+    }
+  }
+
+  useEffect(() => {
+    fetchInitValidateStatus().then((res: InitValidateStatusResponse) => {
+      if (res.status === 'finished')
+        window.location.href = '/install'
+      else
+        setLoading(false)
+    })
+  }, [])
+
+  return (
+    loading
+      ? <Loading />
+      : <div>
+        {!validated && (
+          <div className="block mx-12 min-w-28">
+            <div className="mb-4">
+              <label htmlFor="password" className="block text-sm font-medium text-gray-700">
+                {t('login.adminInitPassword')}
+
+              </label>
+              <div className="mt-1 relative rounded-md shadow-sm">
+                <input
+                  id="password"
+                  type="password"
+                  value={password}
+                  onChange={e => setPassword(e.target.value)}
+                  className="appearance-none block w-full px-3 py-2 border border-gray-300 rounded-md shadow-sm placeholder-gray-400 focus:outline-none focus:ring-indigo-500 focus:border-indigo-500 sm:text-sm"
+                />
+              </div>
+            </div>
+            <div className="flex flex-row flex-wrap justify-stretch p-0">
+              <Button type="primary" onClick={handleValidation} className="basis-full min-w-28">
+                {t('login.validate')}
+              </Button>
+            </div>
+          </div>
+        )}
+      </div>
+  )
+}
+
+export default InitPasswordPopup

+ 22 - 0
web/app/init/page.tsx

@@ -0,0 +1,22 @@
+import React from 'react'
+import classNames from 'classnames'
+import style from '../signin/page.module.css'
+import InitPasswordPopup from './InitPasswordPopup'
+
+const Install = () => {
+  return (
+    <div className={classNames(
+      style.background,
+      'flex w-full min-h-screen',
+      'p-4 lg:p-8',
+      'gap-x-20',
+      'justify-center lg:justify-start',
+    )}>
+      <div className="block m-auto w-96">
+        <InitPasswordPopup />
+      </div>
+    </div>
+  )
+}
+
+export default Install

+ 11 - 5
web/app/install/installForm.tsx

@@ -9,8 +9,8 @@ import Loading from '../components/base/loading'
 import Button from '@/app/components/base/button'
 // import I18n from '@/context/i18n'
 
-import { fetchSetupStatus, setup } from '@/service/common'
-import type { SetupStatusResponse } from '@/models/common'
+import { fetchInitValidateStatus, fetchSetupStatus, setup } from '@/service/common'
+import type { InitValidateStatusResponse, SetupStatusResponse } from '@/models/common'
 
 const validEmailReg = /^[\w\.-]+@([\w-]+\.)+[\w-]{2,}$/
 const validPassword = /^(?=.*[a-zA-Z])(?=.*\d).{8,}$/
@@ -70,10 +70,16 @@ const InstallForm = () => {
 
   useEffect(() => {
     fetchSetupStatus().then((res: SetupStatusResponse) => {
-      if (res.step === 'finished')
+      if (res.step === 'finished') {
         window.location.href = '/signin'
-      else
-        setLoading(false)
+      }
+      else {
+        fetchInitValidateStatus().then((res: InitValidateStatusResponse) => {
+          if (res.status === 'not_started')
+            window.location.href = '/init'
+        })
+      }
+      setLoading(false)
     })
   }, [])
 

+ 5 - 3
web/i18n/lang/login.en.ts

@@ -9,7 +9,7 @@ const translation = {
   namePlaceholder: 'Your username',
   forget: 'Forgot your password?',
   signBtn: 'Sign in',
-  installBtn: 'Setting',
+  installBtn: 'Set up',
   setAdminAccount: 'Setting up an admin account',
   setAdminAccountDesc: 'Maximum privileges for admin account, which can be used to create applications and manage LLM providers, etc.',
   createAndSignIn: 'Create and sign in',
@@ -32,7 +32,7 @@ const translation = {
   tosDesc: 'By signing up, you agree to our',
   donthave: 'Don\'t have?',
   invalidInvitationCode: 'Invalid invitation code',
-  accountAlreadyInited: 'Account already inited',
+  accountAlreadyInited: 'Account already initialized',
   error: {
     emailEmpty: 'Email address is required',
     emailInValid: 'Please enter a valid email address',
@@ -51,7 +51,9 @@ const translation = {
   explore: 'Explore Dify',
   activatedTipStart: 'You have joined the',
   activatedTipEnd: 'team',
-  activated: 'Sign In Now',
+  activated: 'Sign in now',
+  adminInitPassword: 'Admin initialization password',
+  validate: 'Validate',
 }
 
 export default translation

+ 2 - 0
web/i18n/lang/login.zh.ts

@@ -52,6 +52,8 @@ const translation = {
   activatedTipStart: '您已加入',
   activatedTipEnd: '团队',
   activated: '现在登录',
+  adminInitPassword: '管理员初始化密码',
+  validate: '验证',
 }
 
 export default translation

+ 4 - 0
web/models/common.ts

@@ -13,6 +13,10 @@ export type SetupStatusResponse = {
   setup_at?: Date
 }
 
+export type InitValidateStatusResponse = {
+  status: 'finished' | 'not_started'
+}
+
 export type UserProfileResponse = {
   id: string
   name: string

+ 5 - 1
web/service/base.ts

@@ -256,7 +256,11 @@ const baseFetch = <T>(
                 }
                 const loginUrl = `${globalThis.location.origin}/signin`
                 bodyJson.then((data: ResponseError) => {
-                  if (data.code === 'not_setup' && IS_CE_EDITION)
+                  if (data.code === 'init_validate_failed' && IS_CE_EDITION)
+                    Toast.notify({ type: 'error', message: data.message, duration: 4000 })
+                  else if (data.code === 'not_init_validated' && IS_CE_EDITION)
+                    globalThis.location.href = `${globalThis.location.origin}/init`
+                  else if (data.code === 'not_setup' && IS_CE_EDITION)
                     globalThis.location.href = `${globalThis.location.origin}/install`
                   else if (location.pathname !== '/signin' || !IS_CE_EDITION)
                     globalThis.location.href = loginUrl

+ 9 - 0
web/service/common.ts

@@ -9,6 +9,7 @@ import type {
   FileUploadConfigResponse,
   ICurrentWorkspace,
   IWorkspace,
+  InitValidateStatusResponse,
   InvitationResponse,
   LangGeniusVersionResponse,
   Member,
@@ -42,6 +43,14 @@ export const setup: Fetcher<CommonResponse, { body: Record<string, any> }> = ({
   return post<CommonResponse>('/setup', { body })
 }
 
+export const initValidate: Fetcher<CommonResponse, { body: Record<string, any> }> = ({ body }) => {
+  return post<CommonResponse>('/init', { body })
+}
+
+export const fetchInitValidateStatus = () => {
+  return get<InitValidateStatusResponse>('/init')
+}
+
 export const fetchSetupStatus = () => {
   return get<SetupStatusResponse>('/setup')
 }