How to implement login required decorator in Flask

Question

I have 2 Flask apps (different projects) that work together . One implements some API which uses tokens for auth. The second one consumes the API and makes a web interface for it. Now I have a login function that sends the username and password to the API, and if correct, gets the auth token in return. Once I have the token, I save it to the session of the user and the user should now be considered as logged in/ autheticated. How can I implement the login_required decorator for such a case.

Here is my login function -

 def login(self):
        response = make_request(BASE_URL + 'login/', clean_data(self.data))
        if response.status_code == 200:
            session['auth_token'] = response.json().get('auth_token')
            return True
        return False

How can I make the login_required decorator?

Also I am using Redis to store sessions if that matters.

Solution 1

Have a look at the official flask docs regarding decorators: https://flask.palletsprojects.com/en/1.1.x/patterns/viewdecorators/ or the python docs https://www.python.org/dev/peps/pep-0318/ as well.

Your decorator should look something like:

from functools import wraps
from flask import abort
import jwt

def authorize(f):
    @wraps(f)
    def decorated_function(*args, **kws):
            if not 'Authorization' in request.headers:
               abort(401)

            user = None
            data = request.headers['Authorization'].encode('ascii','ignore')
            token = str.replace(str(data), 'Bearer ','')
            try:
                user = jwt.decode(token, JWT_SECRET, algorithms=['HS256'])['sub']
            except:
                abort(401)

            return f(user, *args, **kws)            
    return decorated_function

... and then in your app.py you may have:

@app.route('/api/game', methods=['POST'])
@authorize
def create(user):
    data = json.loads(request.data)
    ....

In this particular case I have used JWT as token and your token can be different respectively the decoding of the token can be your custom implementation, but the basic mechanisms are pretty much as on the example above.

Solution 2

Given that each subsequent request will contain the API token, the decorator should do the following

• Accept a generic request. You can use *args and **kargs for that
• Extract the token from the header and compare it with the token stored in db (not Redis, but wherever the token generated is stored in the backend)
• If authenticated, the *args and **kargs should be passed on to the decorated function
• The output of the decorated function should then be returned as is
• If the authentication failed, an error message should be returned.

For explanation on decorators, check out this link: http://thecodeship.com/patterns/guide-to-python-function-decorators/

Solution 3

I would place the following decorator function in somewhere common

def validate_api_token(validation_func):
    def decorator(f):
        @wraps(f)
        def decorated_function(*args, **kws):
            api_token = request.headers.get('Authorization')
            is_valid_api_token = validation_func(api_token)
            if is_valid_api_token:
                return f(*args, **kws)

            return 'Invalid API Token', 401

        return decorated_function

    return decorator

For small POC flask apps, if you're ok with storing the tokens in a non-versioned file, the following can work:

# tokens are read from a non-versioned `.tokens` file and loaded into a set
api_tokens = load_api_tokens() 


def simple_api_token_validation(api_token):
    return api_token in api_tokens


@app.route("/v1/my/secret/function", methods=['POST'])
@validate_api_token(simple_api_token_validation)
def my_secret_function():
   body = request.get_json()
   # ...

Another simple option is to query against a database (e.g. redis):

redis_session = Redis(host=REDIS_HOST, password=REDIS_PASSWORD)


def redis_api_token_validation(api_token):
    if not api_token:
        return False

    api_token_hash = hashlib.sha256(api_token.encode()).hexdigest()
    return redis_session.exists(f'api:tokens:{api_token_hash}')


@app.route("/v1/my/secret/function", methods=['POST'])
@validate_api_token(redis_api_token_validation)
def my_secret_function():
   body = request.get_json()
   # ...

Best IMO as @Velin answered is to use jwt to validate the token

Comments

[Celeo]

Are you looking to make your own decorator or use something like flask-security?

[utkbansal]

Why should I not compare the token to the one stored in the redis? I'll be using the decorator on almost all views, and if I compare it with the value in the db, I'll have an extra api call in each view.

[TheGeorgeous]

The token in redis is the one you got from the auth server, correct? So you should be comparing it with the one that is there in the auth server.

[TheGeorgeous]

Also there won't be any extra API call, the token should be part of your API request header

[utkbansal]

I'm unable to get what you are trying to explain. The auth server returns the token when I send the correct email and password. And for all the other api requests, I have to send the auth token to the api server as a header.

[utkbansal]

So when I first recive the auth token, I store it in the session which is in redis.

[TheGeorgeous]

Correct. Ideally the token should have an expiry time set on the server for security.