TokenAuth.py

from flask import request, make_response
from functools import wraps


class TokenAuth(object):
	def __init__(self):
		self.authenticate_callback = None
		self.auth_error_callback = None

		def default_authenticate(token):
			return False

		def default_auth_error():
			return "Unauthorized Access"

		self.authenticate(default_authenticate)
		self.error_handler(default_auth_error)

	def authenticate(self, f):
		self.authenticate_callback = f
		return f

	def error_handler(self, f):
		@wraps(f)
		def decorated(*args, **kwargs):
			res = f(*args, **kwargs)
			res = make_response(res)
			if res.status_code == 200:
				# if user didn't set status code, use 401
				res.status_code = 401
			if 'Authorization' not in res.headers.keys():
				res.headers['Authorization'] = ''
			return res

		self.auth_error_callback = decorated
		return decorated

	def login_required(self, f):
		@wraps(f)
		def decorated(*args, **kwargs):
			if 'Authorization' not in request.headers:
				request.data
				return self.auth_error_callback()

			token = request.headers['Authorization']
			if self.authenticate_callback(token) is False:
				request.data
				return self.auth_error_callback()

			return f(*args, **kwargs)

		return decorated