feat: service api auth

api/controllers/inner_api/app.py

@@ -7,12 +7,12 @@ from flask.helpers import stream_with_context
 
 from controllers.console.setup import setup_required
 from controllers.inner_api import api
-from controllers.inner_api.wraps import inner_api_only
+from controllers.inner_api.wraps import inner_api_only, inner_api_user_auth
 from services.completion_service import CompletionService
 from core.entities.application_entities import InvokeFrom
 
 from extensions.ext_database import db
-from models.model import App
+from models.model import App, EndUser
 
 from typing import Union, Generator
 from werkzeug.exceptions import InternalServerError, NotFound
@@ -35,7 +35,8 @@ class EnterpriseAppInvokeApi(Resource):
 
     @setup_required
     @inner_api_only
-    def post(self):
+    @inner_api_user_auth
+    def post(self, **kwargs: dict):
         request_parser = reqparse.RequestParser()
         request_parser.add_argument('app_id', type=str, required=True, nullable=False, location='json')
         request_parser.add_argument('query', type=str, required=True, nullable=False, location='json')
@@ -45,7 +46,6 @@ class EnterpriseAppInvokeApi(Resource):
 
         args = request_parser.parse_args()
 
-        
         try:
             app_id = args['app_id']
             app_model: App = db.session.query(App).filter(App.id == app_id).first()
@@ -57,7 +57,7 @@ class EnterpriseAppInvokeApi(Resource):
             
             response = CompletionService.completion(
                 app_model=app_model,
-                user=current_user,
+                user=kwargs['user'] if 'user' in kwargs else current_user,
                 args=args,
                 invoke_from=InvokeFrom.INNER_API,
                 streaming=args['stream'] if 'stream' in args else False,

api/controllers/inner_api/wraps.py

 from functools import wraps
 
 from flask import abort, current_app, request
+from hmac import new as hmac_new
+from hashlib import sha1
+from base64 import b64encode
+from json import dumps
 
+from models.model import EndUser
+
+from extensions.ext_database import db
 
 def inner_api_only(view):
     @wraps(view)
@@ -17,3 +24,38 @@ def inner_api_only(view):
         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
\ No newline at end of file