Feat/enterprise sso (#3602)

api/controllers/console/__init__.py

@@ -19,4 +19,6 @@ from .datasets import data_source, datasets, datasets_document, datasets_segment
 from .explore import (audio, completion, conversation, installed_app, message, parameter, recommended_app,
                       saved_message, workflow)
 # Import workspace controllers
-from .workspace import account, members, model_providers, models, tool_providers, workspace
+from .workspace import account, members, model_providers, models, tool_providers, workspace
+# Import enterprise controllers
+from .enterprise import enterprise_sso

api/controllers/console/auth/login.py

@@ -26,10 +26,13 @@ class LoginApi(Resource):
 
         try:
             account = AccountService.authenticate(args['email'], args['password'])
-        except services.errors.account.AccountLoginError:
-            return {'code': 'unauthorized', 'message': 'Invalid email or password'}, 401
+        except services.errors.account.AccountLoginError as e:
+            return {'code': 'unauthorized', 'message': str(e)}, 401
 
-        TenantService.create_owner_tenant_if_not_exist(account)
+        # SELF_HOSTED only have one workspace
+        tenants = TenantService.get_join_tenants(account)
+        if len(tenants) == 0:
+            return {'result': 'fail', 'data': 'workspace not found, please contact system admin to invite you to join in a workspace'}
 
         AccountService.update_last_login(account, request)
 

api/controllers/console/enterprise/enterprise_sso.py

@@ -0,0 +1,59 @@
+from flask import current_app, redirect
+from flask_restful import Resource, reqparse
+
+from controllers.console import api
+from controllers.console.setup import setup_required
+from services.enterprise.enterprise_sso_service import EnterpriseSSOService
+
+
+class EnterpriseSSOSamlLogin(Resource):
+
+    @setup_required
+    def get(self):
+        return EnterpriseSSOService.get_sso_saml_login()
+
+
+class EnterpriseSSOSamlAcs(Resource):
+
+    @setup_required
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('SAMLResponse', type=str, required=True, location='form')
+        args = parser.parse_args()
+        saml_response = args['SAMLResponse']
+
+        try:
+            token = EnterpriseSSOService.post_sso_saml_acs(saml_response)
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}/signin?console_token={token}')
+        except Exception as e:
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}/signin?message={str(e)}')
+
+
+class EnterpriseSSOOidcLogin(Resource):
+
+    @setup_required
+    def get(self):
+        return EnterpriseSSOService.get_sso_oidc_login()
+
+
+class EnterpriseSSOOidcCallback(Resource):
+
+    @setup_required
+    def get(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('state', type=str, required=True, location='args')
+        parser.add_argument('code', type=str, required=True, location='args')
+        parser.add_argument('oidc-state', type=str, required=True, location='cookies')
+        args = parser.parse_args()
+
+        try:
+            token = EnterpriseSSOService.get_sso_oidc_callback(args)
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}/signin?console_token={token}')
+        except Exception as e:
+            return redirect(f'{current_app.config.get("CONSOLE_WEB_URL")}/signin?message={str(e)}')
+
+
+api.add_resource(EnterpriseSSOSamlLogin, '/enterprise/sso/saml/login')
+api.add_resource(EnterpriseSSOSamlAcs, '/enterprise/sso/saml/acs')
+api.add_resource(EnterpriseSSOOidcLogin, '/enterprise/sso/oidc/login')
+api.add_resource(EnterpriseSSOOidcCallback, '/enterprise/sso/oidc/callback')

api/controllers/console/feature.py

@@ -1,6 +1,7 @@
 from flask_login import current_user
 from flask_restful import Resource
 
+from services.enterprise.enterprise_feature_service import EnterpriseFeatureService
 from services.feature_service import FeatureService
 
 from . import api
@@ -14,4 +15,10 @@ class FeatureApi(Resource):
         return FeatureService.get_features(current_user.current_tenant_id).dict()
 
 
+class EnterpriseFeatureApi(Resource):
+    def get(self):
+        return EnterpriseFeatureService.get_enterprise_features().dict()
+
+
 api.add_resource(FeatureApi, '/features')
+api.add_resource(EnterpriseFeatureApi, '/enterprise-features')

api/controllers/console/setup.py

@@ -58,6 +58,8 @@ class SetupApi(Resource):
             password=args['password']
         )
 
+        TenantService.create_owner_tenant_if_not_exist(account)
+
         setup()
         AccountService.update_last_login(account, request)
 

api/controllers/console/workspace/workspace.py

@@ -3,6 +3,7 @@ import logging
 from flask import request
 from flask_login import current_user
 from flask_restful import Resource, fields, inputs, marshal, marshal_with, reqparse
+from werkzeug.exceptions import Unauthorized
 
 import services
 from controllers.console import api
@@ -19,7 +20,7 @@ from controllers.console.wraps import account_initialization_required, cloud_edi
 from extensions.ext_database import db
 from libs.helper import TimestampField
 from libs.login import login_required
-from models.account import Tenant
+from models.account import Tenant, TenantStatus
 from services.account_service import TenantService
 from services.file_service import FileService
 from services.workspace_service import WorkspaceService
@@ -116,6 +117,16 @@ class TenantApi(Resource):
 
         tenant = current_user.current_tenant
 
+        if tenant.status == TenantStatus.ARCHIVE:
+            tenants = TenantService.get_join_tenants(current_user)
+            # if there is any tenant, switch to the first one
+            if len(tenants) > 0:
+                TenantService.switch_tenant(current_user, tenants[0].id)
+                tenant = tenants[0]
+            # else, raise Unauthorized
+            else:
+                raise Unauthorized('workspace is archived')
+
         return WorkspaceService.get_tenant_info(tenant), 200