File size: 1,773 Bytes
4304c6d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
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
|