Encode and decode JSON Web Token/JWT (with RS256) from private/public key files

jwt_tools.py

"""
Encode and decode JSON Web Token/JWT (with RS256) from private/public key files
Requires https://cryptography.io:
pip install cryptography
"""

__author__ = 'cuzi'
__email__ = 'cuzi@openmail.cc'
__all__ = ['generate_jwt', 'validate_jwt']

from base64 import urlsafe_b64encode, urlsafe_b64decode
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding
import json

def b64encode(s):
    return urlsafe_b64encode(s).replace(b'=', b'').decode('ascii')

def b64decode(s):
    return urlsafe_b64decode(s.encode('ascii') + b'======')

def jdumps(obj):
    return json.dumps(obj, separators=(',', ':'))

def generate_jwt(payload, key_file, algorithm='RS256'):
    """Generate a signed JSON Web Token/JWT

    Returns:
    The signed token as a str

    Keyword arguments:
    payload -- A dict
    key_file -- Path to the .pem private key file
    algorithm -- Currently only RS256

    Usage:
    token = generate_jwt(payload, 'private_key.pem', 'RS256')
    """
    headers = {
        'typ': 'JWT',
        'alg': algorithm
    }

    headers_base64 = b64encode(jdumps(headers).encode('ascii'))
    payload_base64 = b64encode(jdumps(payload).encode('utf-8'))
    message = headers_base64 + '.' + payload_base64

    if algorithm == 'RS256':
        with open(key_file, 'rb') as f:
            private_key = serialization.load_pem_private_key(
                f.read(),
                password=None,
                backend=default_backend()
            )
        signature = private_key.sign(
            message.encode('ascii'),
            padding.PKCS1v15(),
            hashes.SHA256()
        )
    else:
        raise ValueError('Unsupported algorithm')

    signature_base64 = b64encode(signature)
    return message + '.' + signature_base64


def validate_jwt(token, public_key, algorithm='RS256'):
    """Validate and decode a JSON Web Token/JWT

    Returns:
    The decoded payload i.e. a dict

    Keyword arguments:
    token -- The signed token
    public_key -- A RSAPublicKey
    algorithm -- Currently only RS256

    Raises:
    ValueError - For wrong algorithm
    InvalidSignature - For signature mismatch

    Usage:
    with open('public_key.pub', 'rb') as f:
        public_key = serialization.load_pem_public_key(
            f.read(),
            backend=default_backend()
        )
    result = validate_jwt(token, public_key, 'RS256')
    """
    headers_base64, payload_base64, signature_base64 = token.split('.')

    headers = json.loads(b64decode(headers_base64).decode('ascii'))
    signature = b64decode(signature_base64)
    payload = b64decode(payload_base64).decode('utf-8')

    if headers['alg'] != algorithm:
        raise ValueError(
            'Provided algorithm does not match algorithm from token header')

    if algorithm == 'RS256':
        public_key.verify(
            signature,
            (headers_base64 + '.' + payload_base64).encode('ascii'),
            padding.PKCS1v15(),
            hashes.SHA256()
        )
    else:
        raise ValueError('Unsupported algorithm')

    return json.loads(payload)


if __name__ == '__main__':
    # Example, requires private_key.pem and public_key.pub

    payload = {
        'some': 'dict',
        'data': 1.5
    }

    token = generate_jwt(payload, 'private_key.pem', 'RS256')
    print("Encoded:")
    print(token)

    with open('public_key.pub', 'rb') as f:
        public_key = serialization.load_pem_public_key(
            f.read(),
            backend=default_backend()
        )
    result = validate_jwt(token, public_key, 'RS256')
    print("Decoded:")
    print(result)