Feat/enterprise sso (#3602)

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