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
 
 

api/controllers/inner_api/__init__.py

@@ -0,0 +1,8 @@
+from flask import Blueprint
+from libs.external_api import ExternalApi
+
+bp = Blueprint('inner_api', __name__, url_prefix='/inner/api')
+api = ExternalApi(bp)
+
+from .workspace import workspace
+

api/controllers/inner_api/workspace/workspace.py

@@ -0,0 +1,37 @@
+from flask_restful import Resource, reqparse
+
+from controllers.console.setup import setup_required
+from controllers.inner_api import api
+from controllers.inner_api.wraps import inner_api_only
+from events.tenant_event import tenant_was_created
+from models.account import Account
+from services.account_service import TenantService
+
+
+class EnterpriseWorkspace(Resource):
+
+    @setup_required
+    @inner_api_only
+    def post(self):
+        parser = reqparse.RequestParser()
+        parser.add_argument('name', type=str, required=True, location='json')
+        parser.add_argument('owner_email', type=str, required=True, location='json')
+        args = parser.parse_args()
+
+        account = Account.query.filter_by(email=args['owner_email']).first()
+        if account is None:
+            return {
+                'message': 'owner account not found.'
+            }, 404
+
+        tenant = TenantService.create_tenant(args['name'])
+        TenantService.create_tenant_member(tenant, account, role='owner')
+
+        tenant_was_created.send(tenant)
+
+        return {
+            'message': 'enterprise workspace created.'
+        }
+
+
+api.add_resource(EnterpriseWorkspace, '/enterprise/workspace')

api/controllers/inner_api/wraps.py

@@ -0,0 +1,61 @@
+from base64 import b64encode
+from functools import wraps
+from hashlib import sha1
+from hmac import new as hmac_new
+
+from flask import abort, current_app, request
+
+from extensions.ext_database import db
+from models.model import EndUser
+
+
+def inner_api_only(view):
+    @wraps(view)
+    def decorated(*args, **kwargs):
+        if not current_app.config['INNER_API']:
+            abort(404)
+
+        # get header 'X-Inner-Api-Key'
+        inner_api_key = request.headers.get('X-Inner-Api-Key')
+        if not inner_api_key or inner_api_key != current_app.config['INNER_API_KEY']:
+            abort(404)
+
+        return view(*args, **kwargs)
+
+    return decorated
+
+
+def inner_api_user_auth(view):
+    @wraps(view)
+    def decorated(*args, **kwargs):
+        if not current_app.config['INNER_API']:
+            return view(*args, **kwargs)
+
+        # get header 'X-Inner-Api-Key'
+        authorization = request.headers.get('Authorization')
+        if not authorization:
+            return view(*args, **kwargs)
+
+        parts = authorization.split(':')
+        if len(parts) != 2:
+            return view(*args, **kwargs)
+
+        user_id, token = parts
+        if ' ' in user_id:
+            user_id = user_id.split(' ')[1]
+
+        inner_api_key = request.headers.get('X-Inner-Api-Key')
+
+        data_to_sign = f'DIFY {user_id}'
+
+        signature = hmac_new(inner_api_key.encode('utf-8'), data_to_sign.encode('utf-8'), sha1)
+        signature = b64encode(signature.digest()).decode('utf-8')
+
+        if signature != token:
+            return view(*args, **kwargs)
+
+        kwargs['user'] = db.session.query(EndUser).filter(EndUser.id == user_id).first()
+
+        return view(*args, **kwargs)
+
+    return decorated

api/controllers/service_api/wraps.py

@@ -12,7 +12,7 @@ from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
 
 from extensions.ext_database import db
 from libs.login import _get_user
-from models.account import Account, Tenant, TenantAccountJoin
+from models.account import Account, Tenant, TenantAccountJoin, TenantStatus
 from models.model import ApiToken, App, EndUser
 from services.feature_service import FeatureService
 
@@ -47,6 +47,10 @@ def validate_app_token(view: Optional[Callable] = None, *, fetch_user_arg: Optio
             if not app_model.enable_api:
                 raise NotFound()
 
+            tenant = db.session.query(Tenant).filter(Tenant.id == app_model.tenant_id).first()
+            if tenant.status == TenantStatus.ARCHIVE:
+                raise NotFound()
+
             kwargs['app_model'] = app_model
 
             if fetch_user_arg:
@@ -137,6 +141,7 @@ def validate_dataset_token(view=None):
                 .filter(Tenant.id == api_token.tenant_id) \
                 .filter(TenantAccountJoin.tenant_id == Tenant.id) \
                 .filter(TenantAccountJoin.role.in_(['owner'])) \
+                .filter(Tenant.status == TenantStatus.NORMAL) \
                 .one_or_none() # TODO: only owner information is required, so only one is returned.
             if tenant_account_join:
                 tenant, ta = tenant_account_join

api/controllers/web/site.py

@@ -6,6 +6,7 @@ from werkzeug.exceptions import Forbidden
 from controllers.web import api
 from controllers.web.wraps import WebApiResource
 from extensions.ext_database import db
+from models.account import TenantStatus
 from models.model import Site
 from services.feature_service import FeatureService
 
@@ -54,6 +55,9 @@ class AppSiteApi(WebApiResource):
         if not site:
             raise Forbidden()
 
+        if app_model.tenant.status == TenantStatus.ARCHIVE:
+            raise Forbidden()
+
         can_replace_logo = FeatureService.get_features(app_model.tenant_id).can_replace_logo
 
         return AppSiteInfo(app_model.tenant, app_model, site, end_user.id, can_replace_logo)