Underground-Digital/Workflow-Engine

1from base64 import b64encode2from functools import wraps3from hashlib import sha14from hmac import new as hmac_new5 6from flask import abort, request7 8from configs import dify_config9from extensions.ext_database import db10from models.model import EndUser11 12 13def inner_api_only(view):14    @wraps(view)15    def decorated(*args, **kwargs):16        if not dify_config.INNER_API:17            abort(404)18 19        # get header 'X-Inner-Api-Key'20        inner_api_key = request.headers.get("X-Inner-Api-Key")21        if not inner_api_key or inner_api_key != dify_config.INNER_API_KEY:22            abort(401)23 24        return view(*args, **kwargs)25 26    return decorated27 28 29def inner_api_user_auth(view):30    @wraps(view)31    def decorated(*args, **kwargs):32        if not dify_config.INNER_API:33            return view(*args, **kwargs)34 35        # get header 'X-Inner-Api-Key'36        authorization = request.headers.get("Authorization")37        if not authorization:38            return view(*args, **kwargs)39 40        parts = authorization.split(":")41        if len(parts) != 2:42            return view(*args, **kwargs)43 44        user_id, token = parts45        if " " in user_id:46            user_id = user_id.split(" ")[1]47 48        inner_api_key = request.headers.get("X-Inner-Api-Key")49 50        data_to_sign = f"DIFY {user_id}"51 52        signature = hmac_new(inner_api_key.encode("utf-8"), data_to_sign.encode("utf-8"), sha1)53        signature = b64encode(signature.digest()).decode("utf-8")54 55        if signature != token:56            return view(*args, **kwargs)57 58        kwargs["user"] = db.session.query(EndUser).filter(EndUser.id == user_id).first()59 60        return view(*args, **kwargs)61 62    return decorated63