monkey/monkey_island/cc/auth.py

54 lines
1.4 KiB
Python

from functools import wraps
from flask import current_app, abort
from flask_jwt import JWT, _jwt_required, JWTError
from werkzeug.security import safe_str_cmp
from cc.environment.environment import env
__author__ = 'itay.mizeretz'
class User(object):
def __init__(self, id, username, secret):
self.id = id
self.username = username
self.secret = secret
def __str__(self):
return "User(id='%s')" % self.id
def init_jwt(app):
users = env.get_auth_users()
username_table = {u.username: u for u in users}
userid_table = {u.id: u for u in users}
def authenticate(username, secret):
user = username_table.get(username, None)
if user and safe_str_cmp(user.secret.encode('utf-8'), secret.encode('utf-8')):
return user
def identity(payload):
user_id = payload['identity']
return userid_table.get(user_id, None)
if env.is_auth_enabled():
JWT(app, authenticate, identity)
def jwt_required(realm=None):
def wrapper(fn):
@wraps(fn)
def decorator(*args, **kwargs):
if env.is_auth_enabled():
try:
_jwt_required(realm or current_app.config['JWT_DEFAULT_REALM'])
except JWTError:
abort(401)
return fn(*args, **kwargs)
return decorator
return wrapper