How to implement json web token (jwt) in flask

March 23, 2024

What is json web token?

JSON Web Token (JWT) is like a digital ID card for websites and apps. It’s often issued after a successful login and securely tells a website about the user and their accessibility

Secure Cookies also do the same thing. but there is some core differences

now create our json web token

The structure

Now create our json webtoken with pyhon

We will use hashlib and hmac library for this

As json web token contains three part. We will add this three part in the token. Also we will apply encode decode. So only we will able to know whats inside token payload.

import base64
import hashlib
import hmac
import json
import os
import secrets
import time
import re

# Define constants
SECRET_KEY = '123'
EXPECTED_ISSUER = 'expected_issuer'
EXPECTED_AUDIENCE = 'expected_audience'

# Define regular expression pattern for issuer name
ISSUER_PATTERN = r'^[a-zA-Z0-9_-]+$'
#ISSUER_PATTERN = r'^[a-z0-9_-]{5,20}$'  # 5-20 characters, lowercase letters, numbers, underscore, hyphen

def encode_jwt(payload, expiration_time):
    header = {'typ': 'JWT', 'alg': 'HS256'}
    header_json = json.dumps(header, separators=(',', ':')).encode('utf-8')
    
    # Include expiration time in the payload
    payload['exp'] = int(time.time()) + expiration_time
    
    payload_json = json.dumps(payload, separators=(',', ':')).encode('utf-8')

    encoded_header = base64.urlsafe_b64encode(header_json).decode('utf-8').rstrip('=')
    encoded_payload = base64.urlsafe_b64encode(payload_json).decode('utf-8').rstrip('=')

    signature = hmac.new(SECRET_KEY.encode('utf-8'), f"{encoded_header}.{encoded_payload}".encode('utf-8'), hashlib.sha256)
    encoded_signature = base64.urlsafe_b64encode(signature.digest()).decode('utf-8').rstrip('=')

    jwt_token = f"{encoded_header}.{encoded_payload}.{encoded_signature}"
    return jwt_token

def decode_jwt(jwt_token):
    # Check if the token contains three parts
    parts = jwt_token.split('.')
    if len(parts) != 3:
        return {'error': 'Invalid token format'}

    encoded_header, encoded_payload, encoded_signature = parts

    # Validate and decode the payload
    try:
        payload = base64.urlsafe_b64decode(encoded_payload + '=' * (-len(encoded_payload) % 4)).decode('utf-8')
        decoded_payload = json.loads(payload)
    except (json.JSONDecodeError, UnicodeDecodeError):
        return {'error': 'Invalid payload format'}

    # Validate expiration time
    if 'exp' in decoded_payload:
        if not isinstance(decoded_payload['exp'], int):
            return {'error': 'Expiration time must be an integer'}
        if decoded_payload['exp'] < int(time.time()):
            return {'error': 'Token expired'}

    # Validate issuer using regular expression
    if 'iss' in decoded_payload:
        issuer = decoded_payload['iss']
        if not re.match(ISSUER_PATTERN, issuer):
            return {'error': 'Invalid issuer name format'}

        if issuer != EXPECTED_ISSUER:
            return {'error': 'Invalid issuer'}

    # Validate audience
    if 'aud' in decoded_payload:
        if not isinstance(decoded_payload['aud'], str):
            return {'error': 'Audience must be a string'}
        if EXPECTED_AUDIENCE not in decoded_payload['aud']:
            return {'error': 'Invalid audience'}

    # Validate signature
    expected_signature = base64.urlsafe_b64encode(hmac.new(SECRET_KEY.encode('utf-8'), f"{encoded_header}.{encoded_payload}".encode('utf-8'), hashlib.sha256).digest()).decode('utf-8').rstrip('=')
    if not secrets.compare_digest(encoded_signature, expected_signature):
        return {'error': 'Invalid signature'}

    return decoded_payload
    
payload = {'user_id': 1, 'username': 'Our UserName', 'iss': 'expected_issuer', 'aud': 'expected_audience'}


# Set expiration time in seconds (e.g., 1 hour)
expiration_time = 3600

# Encoding the payload into a JWT with expiration time
jwt_token = encode_jwt(payload, expiration_time)
print("Encoded JWT with expiration:", jwt_token)

# Decoding the JWT to retrieve the payload
decoded_payload = decode_jwt(jwt_token)
print("Decoded Payload:", decoded_payload)





Copy

Simple implementation in flask

from flask import Flask, request, jsonify, render_template, redirect, url_for, make_response
import base64
import hashlib
import hmac
import json
import datetime

app = Flask(__name__)

# Secret key for encoding/decoding JWT
SECRET_KEY = '234231465457'

def encode_jwt(payload):
	header = {'typ': 'JWT', 'alg': 'HS256'}
	header_json = json.dumps(header, separators=(',', ':')).encode('utf-8')
	payload_json = json.dumps(payload, separators=(',', ':')).encode('utf-8')

	# Base64 encoding header and payload
	encoded_header = base64.urlsafe_b64encode(header_json).decode('utf-8')
	encoded_payload = base64.urlsafe_b64encode(payload_json).decode('utf-8')

	# Creating signature
	signature = hmac.new(SECRET_KEY.encode('utf-8'), f"{encoded_header}.{encoded_payload}".encode('utf-8'), hashlib.sha256)
	encoded_signature = base64.urlsafe_b64encode(signature.digest()).decode('utf-8')

	# Combining all parts to form JWT
	jwt_token = f"{encoded_header}.{encoded_payload}.{encoded_signature}"

	return jwt_token

def decode_jwt(jwt_token):
	encoded_header, encoded_payload, encoded_signature = jwt_token.split('.')
	payload = base64.urlsafe_b64decode(encoded_payload.encode('utf-8')).decode('utf-8')
	return json.loads(payload)

@app.route('/')
def home():
	# Check if user is already logged in
	jwt_token = request.cookies.get('jwt_token')
	if jwt_token:
		try:
			payload = decode_jwt(jwt_token)
			print(payload)
			return redirect(url_for('dashboard'))
		except:
			pass  # Invalid token, proceed to login
	return render_template('index.html')

@app.route('/login', methods=['POST'])
def login():
	# Dummy authentication
	if request.form['username'] == 'username' and request.form['password'] == 'password':
		# Create JWT token with payload
		payload = {'username': request.form['username']}
		jwt_token = encode_jwt(payload)
		
		# Set the JWT token in a cookie
		response = make_response(redirect(url_for('dashboard')))
		response.set_cookie('jwt_token', jwt_token, httponly=True, expires=datetime.datetime.now() + datetime.timedelta(minutes=30))
		
		return response
	else:
		return jsonify({'message': 'Invalid username or password'}), 401

@app.route('/dashboard')
def dashboard():
	# Retrieve JWT token from the cookie
	jwt_token = request.cookies.get('jwt_token')
	if jwt_token:
		try:
			payload = decode_jwt(jwt_token)
			return render_template('dashboard.html', username=payload['username'])
		except:
			return jsonify({'message': 'Invalid token'}), 401
	else:
		return jsonify({'message': 'Token missing'}), 401
		
@app.route('/logout')
def logout():
	# Clear the JWT token from the cookie
	response = make_response(redirect(url_for('home')))
	response.set_cookie('jwt_token', '', expires=0)
	return response
		
		

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

index.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Login</title>
</head>
<body>
    <h1>Login</h1>
    <p>our example code username is username and password is password</p>
    <form action="/login" method="post">
        <label for="username">Username:</label><br>
        <input type="text" id="username" name="username"><br>
        <label for="password">Password:</label><br>
        <input type="password" id="password" name="password"><br><br>
        <input type="submit" value="Submit">
    </form>
</body>
</html>

dashboard.html

<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Protected Dashboard</title>
</head>
<body>
    <h1>Welcome, {{ username }}!</h1>
    <p>This is your protected dashboard.</p>
    <p>Feel free to explore.</p>
    <a href="/logout">Logout</a>
</body>
</html>

Get this complete example code from github