diff --git a/project/server/auth/views.py b/project/server/auth/views.py index a3fe7a5..56370fd 100644 --- a/project/server/auth/views.py +++ b/project/server/auth/views.py @@ -1,7 +1,7 @@ # project/server/auth/views.py - -from flask import Blueprint, request, make_response, jsonify +from functools import wraps +from flask import g, Blueprint, request, make_response, jsonify from flask.views import MethodView from project.server import bcrypt, db @@ -10,6 +10,32 @@ auth_blueprint = Blueprint('auth', __name__) +def login_required(f): + @wraps(f) + def decorated_function(*args, **kwargs): + # get the auth token + auth_header = request.headers.get('Authorization') + auth_token = auth_token = auth_header.split(" ")[1] if auth_header else '' + if auth_token: + resp = User.decode_auth_token(auth_token) + if not isinstance(resp, str): + g.auth_token = auth_token + g.user = User.query.filter_by(id=resp).first() + return f(*args, **kwargs) + response_object = { + 'status': 'fail', + 'message': resp + } + return make_response(jsonify(response_object)), 401 + else: + response_object = { + 'status': 'fail', + 'message': 'Provide a valid auth token.' + } + return make_response(jsonify(response_object)), 401 + return decorated_function + + class RegisterAPI(MethodView): """ User Registration Resource @@ -94,37 +120,16 @@ class UserAPI(MethodView): User Resource """ def get(self): - # get the auth token - auth_header = request.headers.get('Authorization') - if auth_header: - auth_token = auth_header.split(" ")[1] - else: - auth_token = '' - if auth_token: - resp = User.decode_auth_token(auth_token) - if not isinstance(resp, str): - user = User.query.filter_by(id=resp).first() - responseObject = { - 'status': 'success', - 'data': { - 'user_id': user.id, - 'email': user.email, - 'admin': user.admin, - 'registered_on': user.registered_on - } - } - return make_response(jsonify(responseObject)), 200 - responseObject = { - 'status': 'fail', - 'message': resp + responseObject = { + 'status': 'success', + 'data': { + 'user_id': g.user.id, + 'email': g.user.email, + 'admin': g.user.admin, + 'registered_on': g.user.registered_on } - return make_response(jsonify(responseObject)), 401 - else: - responseObject = { - 'status': 'fail', - 'message': 'Provide a valid auth token.' - } - return make_response(jsonify(responseObject)), 401 + } + return make_response(jsonify(responseObject)), 200 class LogoutAPI(MethodView): @@ -132,44 +137,23 @@ class LogoutAPI(MethodView): Logout Resource """ def post(self): - # get auth token - auth_header = request.headers.get('Authorization') - if auth_header: - auth_token = auth_header.split(" ")[1] - else: - auth_token = '' - if auth_token: - resp = User.decode_auth_token(auth_token) - if not isinstance(resp, str): - # mark the token as blacklisted - blacklist_token = BlacklistToken(token=auth_token) - try: - # insert the token - db.session.add(blacklist_token) - db.session.commit() - responseObject = { - 'status': 'success', - 'message': 'Successfully logged out.' - } - return make_response(jsonify(responseObject)), 200 - except Exception as e: - responseObject = { - 'status': 'fail', - 'message': e - } - return make_response(jsonify(responseObject)), 200 - else: - responseObject = { - 'status': 'fail', - 'message': resp - } - return make_response(jsonify(responseObject)), 401 - else: + # mark the token as blacklisted + blacklist_token = BlacklistToken(token=g.auth_token) + try: + # insert the token + db.session.add(blacklist_token) + db.session.commit() + responseObject = { + 'status': 'success', + 'message': 'Successfully logged out.' + } + return make_response(jsonify(responseObject)), 200 + except Exception as e: responseObject = { 'status': 'fail', - 'message': 'Provide a valid auth token.' + 'message': e } - return make_response(jsonify(responseObject)), 403 + return make_response(jsonify(responseObject)), 200 # define the API resources registration_view = RegisterAPI.as_view('register_api') @@ -190,11 +174,11 @@ def post(self): ) auth_blueprint.add_url_rule( '/auth/status', - view_func=user_view, + view_func=login_required(user_view), methods=['GET'] ) auth_blueprint.add_url_rule( '/auth/logout', - view_func=logout_view, + view_func=login_required(logout_view), methods=['POST'] )