jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Remove useless code (#4416)

Browse source
This commit is contained in:
Garfield Dai 2024-05-15 16:14:49 +08:00 • committed by GitHub
commit dd94931116
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
26 changed files with 466 additions and 230 deletions
Download patch file Download diff file Expand all files Collapse all files

3
api/controllers/console/__init__.py
View file

@ -37,9 +37,6 @@ from .billing import billing
# 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,

0
api/controllers/console/enterprise/__init__.py
View file

59
api/controllers/console/enterprise/enterprise_sso.py
View file

@ -1,59 +0,0 @@
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')

7
api/controllers/console/feature.py
View file

@ -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 @@ class FeatureApi(Resource):
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')

2
api/controllers/web/__init__.py
View file

@ -6,4 +6,4 @@ bp = Blueprint('web', __name__, url_prefix='/api')
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

6
api/controllers/web/app.py
View file

@ -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

6
api/controllers/web/error.py
View file

@ -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

12
api/controllers/web/feature.py Normal file
View file

@ -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')

12
api/controllers/web/passport.py
View file

@ -5,14 +5,21 @@ from flask_restful import Resource
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,7 +35,7 @@ class PassportResource(Resource):
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,
@ -36,6 +43,7 @@ class PassportResource(Resource):
is_anonymous=True,
session_id=generate_session_id(),
)
db.session.add(end_user)
db.session.commit()
@ -53,8 +61,10 @@ class PassportResource(Resource):
'access_token': tk,
}
api.add_resource(PassportResource, '/passport')
def generate_session_id():
"""
Generate a unique session ID.

78
api/controllers/web/wraps.py
View file

@ -2,11 +2,13 @@ from functools import wraps
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 validate_jwt_token(view=None):
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.')
system_features = FeatureService.get_system_features()
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()
try:
auth_header = request.headers.get('Authorization')
if auth_header is None:
raise Unauthorized('Authorization header is missing.')
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()
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.')
return app_model, end_user
class WebApiResource(Resource):
method_decorators = [validate_jwt_token]

28
api/services/enterprise/enterprise_feature_service.py
View file

@ -1,28 +0,0 @@
from flask import current_app
from pydantic import BaseModel
from services.enterprise.enterprise_service import EnterpriseService
class EnterpriseFeatureModel(BaseModel):
sso_enforced_for_signin: bool = False
sso_enforced_for_signin_protocol: str = ''
class EnterpriseFeatureService:
@classmethod
def get_enterprise_features(cls) -> EnterpriseFeatureModel:
features = EnterpriseFeatureModel()
if current_app.config['ENTERPRISE_ENABLED']:
cls._fulfill_params_from_enterprise(features)
return features
@classmethod
def _fulfill_params_from_enterprise(cls, features):
enterprise_info = EnterpriseService.get_info()
features.sso_enforced_for_signin = enterprise_info['sso_enforced_for_signin']
features.sso_enforced_for_signin_protocol = enterprise_info['sso_enforced_for_signin_protocol']

60
api/services/enterprise/enterprise_sso_service.py
View file

@ -1,60 +0,0 @@
import logging
from models.account import Account, AccountStatus
from services.account_service import AccountService, TenantService
from services.enterprise.base import EnterpriseRequest
logger = logging.getLogger(__name__)
class EnterpriseSSOService:
@classmethod
def get_sso_saml_login(cls) -> str:
return EnterpriseRequest.send_request('GET', '/sso/saml/login')
@classmethod
def post_sso_saml_acs(cls, saml_response: str) -> str:
response = EnterpriseRequest.send_request('POST', '/sso/saml/acs', json={'SAMLResponse': saml_response})
if 'email' not in response or response['email'] is None:
logger.exception(response)
raise Exception('Saml response is invalid')
return cls.login_with_email(response.get('email'))
@classmethod
def get_sso_oidc_login(cls):
return EnterpriseRequest.send_request('GET', '/sso/oidc/login')
@classmethod
def get_sso_oidc_callback(cls, args: dict):
state_from_query = args['state']
code_from_query = args['code']
state_from_cookies = args['oidc-state']
if state_from_cookies != state_from_query:
raise Exception('invalid state or code')
response = EnterpriseRequest.send_request('GET', '/sso/oidc/callback', params={'code': code_from_query})
if 'email' not in response or response['email'] is None:
logger.exception(response)
raise Exception('OIDC response is invalid')
return cls.login_with_email(response.get('email'))
@classmethod
def login_with_email(cls, email: str) -> str:
account = Account.query.filter_by(email=email).first()
if account is None:
raise Exception('account not found, please contact system admin to invite you to join in a workspace')
if account.status == AccountStatus.BANNED:
raise Exception('account is banned, please contact system admin')
tenants = TenantService.get_join_tenants(account)
if len(tenants) == 0:
raise Exception("workspace not found, please contact system admin to invite you to join in a workspace")
token = AccountService.get_account_jwt_token(account)
return token

25
api/services/feature_service.py
View file

@ -2,6 +2,7 @@ from flask import current_app
from pydantic import BaseModel
from services.billing_service import BillingService
from services.enterprise.enterprise_service import EnterpriseService
class SubscriptionModel(BaseModel):
@ -30,6 +31,13 @@ class FeatureModel(BaseModel):
can_replace_logo: bool = False
class SystemFeatureModel(BaseModel):
sso_enforced_for_signin: bool = False
sso_enforced_for_signin_protocol: str = ''
sso_enforced_for_web: bool = False
sso_enforced_for_web_protocol: str = ''
class FeatureService:
@classmethod
@ -43,6 +51,15 @@ class FeatureService:
return features
@classmethod
def get_system_features(cls) -> SystemFeatureModel:
system_features = SystemFeatureModel()
if current_app.config['ENTERPRISE_ENABLED']:
cls._fulfill_params_from_enterprise(system_features)
return system_features
@classmethod
def _fulfill_params_from_env(cls, features: FeatureModel):
features.can_replace_logo = current_app.config['CAN_REPLACE_LOGO']
@ -73,3 +90,11 @@ class FeatureService:
features.docs_processing = billing_info['docs_processing']
features.can_replace_logo = billing_info['can_replace_logo']
@classmethod
def _fulfill_params_from_enterprise(cls, features):
enterprise_info = EnterpriseService.get_info()
features.sso_enforced_for_signin = enterprise_info['sso_enforced_for_signin']
features.sso_enforced_for_signin_protocol = enterprise_info['sso_enforced_for_signin_protocol']
features.sso_enforced_for_web = enterprise_info['sso_enforced_for_web']
features.sso_enforced_for_web_protocol = enterprise_info['sso_enforced_for_web_protocol']