Remove useless code (#4416)

api/controllers/console/__init__.py

@@ -37,9 +37,6 @@
 # Import datasets controllers
 from .datasets import data_source, datasets, datasets_document, datasets_segments, file, hit_testing
 
-# Import enterprise controllers
-from .enterprise import enterprise_sso
-
 # Import explore controllers
 from .explore import (
     audio,

api/controllers/console/feature.py

@@ -1,7 +1,6 @@
 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
@@ -15,10 +14,10 @@ def get(self):
         return FeatureService.get_features(current_user.current_tenant_id).dict()
 
 
-class EnterpriseFeatureApi(Resource):
+class SystemFeatureApi(Resource):
     def get(self):
-        return EnterpriseFeatureService.get_enterprise_features().dict()
+        return FeatureService.get_system_features().dict()
 
 
 api.add_resource(FeatureApi, '/features')
-api.add_resource(EnterpriseFeatureApi, '/enterprise-features')
+api.add_resource(SystemFeatureApi, '/system-features')

api/controllers/web/__init__.py

@@ -6,4 +6,4 @@
 api = ExternalApi(bp)
 
 
-from . import app, audio, completion, conversation, file, message, passport, saved_message, site, workflow
+from . import app, audio, completion, conversation, feature, file, message, passport, saved_message, site, workflow

api/controllers/web/app.py

@@ -1,14 +1,10 @@
-import json
-
 from flask import current_app
 from flask_restful import fields, marshal_with
 
 from controllers.web import api
 from controllers.web.error import AppUnavailableError
 from controllers.web.wraps import WebApiResource
-from extensions.ext_database import db
-from models.model import App, AppMode, AppModelConfig
-from models.tools import ApiToolProvider
+from models.model import App, AppMode
 from services.app_service import AppService
 
 

api/controllers/web/error.py

@@ -115,3 +115,9 @@ class UnsupportedFileTypeError(BaseHTTPException):
     error_code = 'unsupported_file_type'
     description = "File type not allowed."
     code = 415
+
+
+class WebSSOAuthRequiredError(BaseHTTPException):
+    error_code = 'web_sso_auth_required'
+    description = "Web SSO authentication required."
+    code = 401

api/controllers/web/feature.py

@@ -0,0 +1,12 @@
+from flask_restful import Resource
+
+from controllers.web import api
+from services.feature_service import FeatureService
+
+
+class SystemFeatureApi(Resource):
+    def get(self):
+        return FeatureService.get_system_features().dict()
+
+
+api.add_resource(SystemFeatureApi, '/system-features')

api/controllers/web/passport.py

@@ -5,14 +5,21 @@
 from werkzeug.exceptions import NotFound, Unauthorized
 
 from controllers.web import api
+from controllers.web.error import WebSSOAuthRequiredError
 from extensions.ext_database import db
 from libs.passport import PassportService
 from models.model import App, EndUser, Site
+from services.feature_service import FeatureService
 
 
 class PassportResource(Resource):
     """Base resource for passport."""
     def get(self):
+
+        system_features = FeatureService.get_system_features()
+        if system_features.sso_enforced_for_web:
+            raise WebSSOAuthRequiredError()
+
         app_code = request.headers.get('X-App-Code')
         if app_code is None:
             raise Unauthorized('X-App-Code header is missing.')
@@ -28,14 +35,15 @@ def get(self):
         app_model = db.session.query(App).filter(App.id == site.app_id).first()
         if not app_model or app_model.status != 'normal' or not app_model.enable_site:
             raise NotFound()
-        
+
         end_user = EndUser(
             tenant_id=app_model.tenant_id,
             app_id=app_model.id,
             type='browser',
             is_anonymous=True,
             session_id=generate_session_id(),
         )
+
         db.session.add(end_user)
         db.session.commit()
 
@@ -53,8 +61,10 @@ def get(self):
             'access_token': tk,
         }
 
+
 api.add_resource(PassportResource, '/passport')
 
+
 def generate_session_id():
     """
     Generate a unique session ID.

api/controllers/web/wraps.py

@@ -2,11 +2,13 @@
 
 from flask import request
 from flask_restful import Resource
-from werkzeug.exceptions import NotFound, Unauthorized
+from werkzeug.exceptions import BadRequest, NotFound, Unauthorized
 
+from controllers.web.error import WebSSOAuthRequiredError
 from extensions.ext_database import db
 from libs.passport import PassportService
 from models.model import App, EndUser, Site
+from services.feature_service import FeatureService
 
 
 def validate_jwt_token(view=None):
@@ -21,34 +23,60 @@ def decorated(*args, **kwargs):
         return decorator(view)
     return decorator
 
+
 def decode_jwt_token():
-    auth_header = request.headers.get('Authorization')
-    if auth_header is None:
-        raise Unauthorized('Authorization header is missing.')
-
-    if ' ' not in auth_header:
-        raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
-
-    auth_scheme, tk = auth_header.split(None, 1)
-    auth_scheme = auth_scheme.lower()
-
-    if auth_scheme != 'bearer':
-        raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
-    decoded = PassportService().verify(tk)
-    app_code = decoded.get('app_code')
-    app_model = db.session.query(App).filter(App.id == decoded['app_id']).first()
-    site = db.session.query(Site).filter(Site.code == app_code).first()
-    if not app_model:
-        raise NotFound()
-    if not app_code or not site:
-        raise Unauthorized('Site URL is no longer valid.')
-    if app_model.enable_site is False:
-        raise Unauthorized('Site is disabled.')
-    end_user = db.session.query(EndUser).filter(EndUser.id == decoded['end_user_id']).first()
-    if not end_user:
-        raise NotFound()
-
-    return app_model, end_user
+    system_features = FeatureService.get_system_features()
+
+    try:
+        auth_header = request.headers.get('Authorization')
+        if auth_header is None:
+            raise Unauthorized('Authorization header is missing.')
+
+        if ' ' not in auth_header:
+            raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+
+        auth_scheme, tk = auth_header.split(None, 1)
+        auth_scheme = auth_scheme.lower()
+
+        if auth_scheme != 'bearer':
+            raise Unauthorized('Invalid Authorization header format. Expected \'Bearer <api-key>\' format.')
+        decoded = PassportService().verify(tk)
+        app_code = decoded.get('app_code')
+        app_model = db.session.query(App).filter(App.id == decoded['app_id']).first()
+        site = db.session.query(Site).filter(Site.code == app_code).first()
+        if not app_model:
+            raise NotFound()
+        if not app_code or not site:
+            raise BadRequest('Site URL is no longer valid.')
+        if app_model.enable_site is False:
+            raise BadRequest('Site is disabled.')
+        end_user = db.session.query(EndUser).filter(EndUser.id == decoded['end_user_id']).first()
+        if not end_user:
+            raise NotFound()
+
+        _validate_web_sso_token(decoded, system_features)
+
+        return app_model, end_user
+    except Unauthorized as e:
+        if system_features.sso_enforced_for_web:
+            raise WebSSOAuthRequiredError()
+
+        raise Unauthorized(e.description)
+
+
+def _validate_web_sso_token(decoded, system_features):
+    # Check if SSO is enforced for web, and if the token source is not SSO, raise an error and redirect to SSO login
+    if system_features.sso_enforced_for_web:
+        source = decoded.get('token_source')
+        if not source or source != 'sso':
+            raise WebSSOAuthRequiredError()
+
+    # Check if SSO is not enforced for web, and if the token source is SSO, raise an error and redirect to normal passport login
+    if not system_features.sso_enforced_for_web:
+        source = decoded.get('token_source')
+        if source and source == 'sso':
+            raise Unauthorized('sso token expired.')
+
 
 class WebApiResource(Resource):
     method_decorators = [validate_jwt_token]