Graham Morby

Posted on May 14, 2022 • Updated on May 23, 2022

JWT Auth in Flask Python

#beginners #tutorial #python #flask

Ever wondered how you use auth in Flask and secure some endpoints? Well, it's not that crazy, and it's pretty easy to implement.

So let's first start by creating a new folder - /JWTExample

Pop open a terminal and cd into that folder from your location and run the following command to set up a virtual environment.

python -m venv env

Ok before we carry on if you have python three installed you will need to use 'python3' and 'pip3' in your commands, if not its just 'python' and 'pip'.

Now we have our virtual environment squared away we need to grab some packages from pip to allow for what we need to do. Now this will be a very simple set up and you can go a lot deeper into flask auth.

Lets make sure we are in our virtual environment by running the following:

source env/bin/activate

So lets install:

pip3 install flask flask_jwt_extended flask_sqlalchemy

Ok so we have what we need to get started! PERFECT!

Lets first create a basic Flask app and make sure thats working all ok. So create a file called app.py and lets add the following code.

from flask import Flask, jsonify, request

app = Flask(__name__)

# Planet Routes
@app.route('/', methods=['GET'])
def index():
    return jsonify(message='Welcome to flask!')


if __name__ == '__main__':
    app.run()

If we then use

flask run

we should be greeted with our wonderful message of 'Welcome to Flask'.

So we have done the first step, lets push onto the next. So the next step is allowing a user to create an account. Now for this tutorial we are going to use SQLite as a database but you can of course use any one of your favorites.

Ok so this next bit is a bit long but hang in there and it won't take long.

So first we need a package called marshmallow

pip3 install flask-marshmallow

Let make sure that is being used

from flask_marshmallow import Marshmallow
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import Column, Integer, String, Float

db = SQLAlchemy(app)
ma = Marshmallow(app)

Next, we need to add some flask CLI commands - these we can use to help users create the database and seed it with some data.

# DB set up and seeders
@app.cli.command('db_create')
def db_create():
    db.create_all()
    print('Database created')


@app.cli.command('db_drop')
def db_drop():
    db.drop_all()
    print('Database dropped')


@app.cli.command('db_seed')
def db_seed():
    test_user = User(first_name='Stephen',
                        last_name='Hawking',
                        email='admin@admin.com',
                        password='admin')
    db.session.add(test_user)
    db.session.commit()
    print('Database seeded')

Then we create a User database model, this tells SQL alchemy how to layout our Users table and what type of columns should be.

# Database models
class User(db.Model):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    first_name = Column(String)
    last_name = Column(String)
    email = Column(String, unique=True)
    password = Column(String)

# DB Schemas
class UserSchema(ma.Schema):
    class Meta:
        fields = ('id', 'first_name', 'last_name', 'email', 'password')

and finally lets add the database connection strings

import os
basedir = os.path.abspath(os.path.dirname(__file__))
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + os.path.join(basedir, 'jwt.db')

Ok, so what is flask marshmallow? I guess the best way to describe it is it allows us to output our data in views in a neat way.

Wow, that was a lot of code! Well done you! If you pop down into the terminal and write

flask db_create

you should see a message 'database created'!

Then we test the database with

flask db_seed

that will place our test data in the database and we are ready to look at the register endpoint.

So we need to create a registered route and save any new user to the database.

@app.route('/register', methods=['POST'])
def register():
    email = request.form['email']
    test = User.query.filter_by(email=email).first()
    if test:
        return jsonify(message='That email already exists'), 409
    else:
        first_name = request.form['first_name']
        last_name = request.form['last_name']
        password = request.form['password']
        user = User(first_name=first_name, last_name=last_name, email=email, password=password)
        db.session.add(user)
        db.session.commit()
        return jsonify(message='User created successfully'), 201

Ok so a lot is going here, firstly we take some form data and check if that email already exists, then if it does we return a JSON object with a message stating an error. If we don't have that email we build a User object from all the form data sent in and then save it to the database. Then we return a JSON object with a 201 status and a message.

And now finally we are ready to install flask JWT! So let's grab the package and go for it

pip3 install flask-jwt-extended

So we have our JWT package and now we need to use it. This is a three ponged attack:

Import the package
Set up a super secret key
Create an instance

So import as follows:

from flask_jwt_extended import JWTManager, jwt_required, create_access_token

Set up a super secret key (Don't tell anyone!):

app.config['JWT_SECRET_KEY'] = 'super-secret'  # Change on production

Create an instance:

jwt = JWTManager(app)

So we are fully set up to use out JWT package, so now we build the login route.

@app.route('/login', methods=['POST'])
def login():
    if request.is_json:
        email = request.json['email']
        password = request.json['password']
    else:
        email = request.form['email']
        password = request.form['password']

    test = User.query.filter_by(email=email, password=password).first()
    if test:
        access_token = create_access_token(identity=email)
        return jsonify(message='Login Successful', access_token=access_token)
    else:
        return jsonify('Bad email or Password'), 401

So again loads going on here! So lets start walking through the code. In this code we are accepting JSON or Form Data, we take our email and password and we assign it to variables. We then test to make sure they match and if they do we create a JWT token! and we return a message and the token itself! Woop!! We made it.

So how do we test it I hear you say! Well if we add @jwt-required to our index method it will require us to have a JWT token to access it

@app.route('/', methods=['GET'])
@jwt_required
def index():
    return jsonify(message='Welcome to flask!')

So lets do some testing

First let's load postman and try and access the home route, we should get an error message as follows:

Ok so that works! Let's login and get a JWT token returned, now remember we have already seeded our database and can use that user to login with:

And now we need to add our Bearer token to our auth section and if we hit go on our post method we will get our welcome message! And we have fully installed Flask JWT Auth!

Here is the full code in full:

from flask import Flask, jsonify, request
from flask_sqlalchemy import SQLAlchemy
from sqlalchemy import Column, Integer, String, Float
import os
from flask_marshmallow import Marshmallow
from flask_jwt_extended import JWTManager, jwt_required, create_access_token

app = Flask(__name__)
basedir = os.path.abspath(os.path.dirname(__file__))
app.config['SQLALCHEMY_DATABASE_URI'] = 'sqlite:///' + os.path.join(basedir, 'jwt.db')
app.config['JWT_SECRET_KEY'] = 'super-secret'  # Change on production



db = SQLAlchemy(app)
ma = Marshmallow(app)
jwt = JWTManager(app)


# DB set up and seeders
@app.cli.command('db_create')
def db_create():
    db.create_all()
    print('Database created')


@app.cli.command('db_drop')
def db_drop():
    db.drop_all()
    print('Database dropped')


@app.cli.command('db_seed')
def db_seed():
    test_user = User(first_name='Stephen',
                     last_name='Hawking',
                     email='admin@admin.com',
                     password='admin')
    db.session.add(test_user)
    db.session.commit()
    print('Database seeded')


# Planet Routes
@app.route('/', methods=['GET'])
@jwt_required()
def index():
    return jsonify(message="Hello Flask!")



# User routes
@app.route('/register', methods=['POST'])
def register():
    email = request.form['email']
    test = User.query.filter_by(email=email).first()
    if test:
        return jsonify(message='That email already exists'), 409
    else:
        first_name = request.form['first_name']
        last_name = request.form['last_name']
        password = request.form['password']
        user = User(first_name=first_name, last_name=last_name, email=email, password=password)
        db.session.add(user)
        db.session.commit()
        return jsonify(message='User created successfully'), 201


@app.route('/login', methods=['POST'])
def login():
    if request.is_json:
        email = request.json['email']
        password = request.json['password']
    else:
        email = request.form['email']
        password = request.form['password']

    test = User.query.filter_by(email=email, password=password).first()
    if test:
        access_token = create_access_token(identity=email)
        return jsonify(message='Login Successful', access_token=access_token)
    else:
        return jsonify('Bad email or Password'), 401


# Database models
class User(db.Model):
    __tablename__ = 'users'
    id = Column(Integer, primary_key=True)
    first_name = Column(String)
    last_name = Column(String)
    email = Column(String, unique=True)
    password = Column(String)


# DB Schemas
class UserSchema(ma.Schema):
    class Meta:
        fields = ('id', 'first_name', 'last_name', 'email', 'password')



# Marsh mellow db adds
user_schema = UserSchema()
users_schema = UserSchema(many=True)


if __name__ == '__main__':
    app.run()