Generating JWT Tokens Upon Login

Welcome to the next step in securing your Flask application! In the last lesson, you learned how to create a basic login endpoint that validates user credentials. Today, we'll take it a step further by introducing JSON Web Tokens (JWT) to enhance the security of your application.

Let's briefly recap our Flask app setup and the initial code we wrote.

from flask import Flask, request, jsonify
from marshmallow import Schema, fields, ValidationError
from marshmallow.validate import Length

# Initialize a Flask app instance
app = Flask(__name__)

# Mock database of users
database = [
    {"id": 1, "username": "cosmo", "password": "space-corgi"}
]

# Define a schema for validating login data
class LoginSchema(Schema):
    username = fields.Str(required=True, validate=Length(min=1))
    password = fields.Str(required=True, validate=Length(min=1))

# Create an instance of LoginSchema
login_schema = LoginSchema()

This code initializes a Flask app, sets up a mock database, and uses Marshmallow to validate login data. If you need a refresher on any of these concepts, feel free to revisit the previous lesson.

Before we dive into JWT, let's understand Bearer Tokens. A bearer token is like a "movie ticket" that you present to gain entry. Whoever "bears" this token can access certain resources.

Instead of sending your username and password with every request—which is risky—you log in once and receive a bearer token. This token is then included in the headers of future requests to prove your identity.

JWT is a type of bearer token. By using tokens, you improve security and efficiency since your credentials are not repeatedly transmitted, and servers can quickly verify token validity.

A JSON Web Token (JWT) is a compact, URL-safe means of representing claims between two parties. It's primarily used for securely transmitting information between a client and a server. A JWT token consists of three parts: a Header, a Payload, and a Signature.

• The Header contains metadata about the token, such as the type and hashing algorithm used.
• The Payload carries the actual data or claims, such as user information and token expiration time.
• The Signature ensures the token's integrity by verifying that the content hasn't been tampered with, using a secret key.

By encoding this information, JWT protects user data and offers a reliable way to verify user identities.

Now, let's introduce the Flask-JWT-Extended library, which we'll use to generate and handle JWT tokens. In case you want to install Flask-JWT-Extended in your machine, you can use the following command:

pip install flask-jwt-extended

In the CodeSignal environment, this library is pre-installed, but it's good practice to know how to do this for your local development.

To manage JWTs in your Flask application, you'll be using the JWTManager from the Flask-JWT-Extended library. This will handle the creation, validation, and revocation of JWTs.

from flask_jwt_extended import JWTManager

# Set the secret key for signing JWTs
app.config['JWT_SECRET_KEY'] = 'super-secret' 

# Initialize the JWTManager with the Flask app
jwt = JWTManager(app)

Here’s a step-by-step explanation:

1. Set the Secret Key: app.config['JWT_SECRET_KEY'] = 'super-secret'

   • The JWT_SECRET_KEY is a critical part of your security. It's used to sign the JWTs so they can be verified by your Flask app. Make sure to use a strong, unpredictable secret key in a production environment.

2. Initialize JWTManager: jwt = JWTManager(app)

   • By initializing JWTManager with your Flask app, you integrate JWT handling capabilities, like creating and validating tokens.

By configuring JWTManager, you prepare your Flask app to securely handle JWTs. Now, let's enhance our /login endpoint to generate and return a JWT upon successful authentication.

Before we start generating JWT tokens, let's revisit the login endpoint we made in the last lesson. Here, we validate user credentials and return a success message if they are correct.

# Define a route to receive login credentials
@app.route('/login', methods=['POST'])
def login():
    try:
        # Validate and deserialize the input data according to the schema
        data = login_schema.load(request.get_json())
    except ValidationError as err:
        # If validation fails, return an error message and a 400 status code
        return jsonify(error=err.messages), 400

    # Extract username and password from the validated data
    username = data['username']
    password = data['password']

    # Find the user in the mock database
    user = next((user for user in database if user["username"] == username), None)
    
    # Check if the user exists and if the password matches
    if user and user["password"] == password:
        # Login successful
        return jsonify(message="Login successful"), 200
    else:
        # Return an error if the user does not exist or the password is incorrect
        return jsonify(error="Bad username or password"), 401

Now, let's enhance our /login endpoint to generate and return a JWT token upon successful authentication.

Here’s our adapted endpoint:

from flask_jwt_extended import create_access_token

# Define a route to receive login credentials
@app.route('/login', methods=['POST'])
def login():
    # -- previous input validation and user verification code goes here --

    # Check if the user exists and if the password matches
    if user and user["password"] == password:
        # Create an access token for the user
        access_token = create_access_token(identity=username)
        # Return the access token with a success message
        return jsonify(access_token=access_token), 200
    else:
        # Return an error if the user does not exist or the password is incorrect
        return jsonify(error="Bad username or password"), 401

Key changes:

1. Import create_access_token: To generate a JWT token.
2. Generate Access Token: Create an access token using create_access_token(identity=username), by setting the identity to the username, the token can later be used to identify the user making requests.
3. Return Access Token: Return the generated JWT token in the JSON response.

With these changes, our login endpoint now provides a JWT token, enhancing the security and efficiency of our application.

When a client successfully logs in by providing valid credentials to the /login endpoint, the response will include the generated JWT token. Here's what the response might look like:

{
    "access_token": "eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9..."
}

The client should keep this token securely and send it in the Authorization header as Bearer <token> for future requests to access protected resources.

In this lesson, you learned how to enhance your Flask login endpoint by generating and returning JWT tokens. We covered:

• What JWT is and why it's useful for authentication.
• How to install and configure Flask-JWT-Extended.
• Setting up JWTManager in your Flask app.
• Modifying the login endpoint to generate and return a JWT on successful authentication.

Now that you're equipped with the knowledge of JWT-based authentication, it's time to put it into practice. Keep testing and experimenting with the code, and always double-check the security aspects of your implementation.

Happy coding!