Playing with RSA JWTs

jwt_demo.py

import jwt


def create_token(payload, private_key):
    return jwt.encode(payload, private_key, algorithm="RS256")


def decode_token(token, public_key, default=None):
    try:
        return jwt.decode(token, public_key, algorithms="RS256")
    except jwt.exceptions.InvalidSignatureError:
        return default


if __name__ == "__main__":
    import rsa

    def test():
        private_key = rsa.generate_key()
        public_key = rsa.generate_public_key(private_key)

        payload = {"Hello": "World"}
        token = create_token(payload, private_key)

        parts = token.decode().split(".")
        corrupt = f"{parts[0]}.eyJIZWxsbyBXb3JsZCI6IldvcmxkIn0.{parts[2]}".encode()

        assert decode_token(token, public_key) == payload
        assert decode_token(corrupt, public_key) is None

        print("Success")

    test()

requirements.txt

cffi==1.14.0
cryptography==2.9
pycparser==2.20
PyJWT==1.7.1
six==1.14.0

rsa.py

import cryptography.exceptions
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import rsa
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat import backends


def generate_key():
    return rsa.generate_private_key(
        public_exponent=65537,
        key_size=2048,
        backend=backends.default_backend()
    )


def generate_public_key(private_key):
    return private_key.public_key()


def generate_private_pem(private_key, password):
    return private_key.private_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.BestAvailableEncryption(password)
    )


def generate_public_pem(public_key):
    return public_key.public_bytes(
        encoding=serialization.Encoding.PEM,
        format=serialization.PublicFormat.SubjectPublicKeyInfo
    )


def sign(message, private_key):
    return private_key.sign(
        message,
        padding.PSS(
            mgf=padding.MGF1(hashes.SHA256()),
            salt_length=padding.PSS.MAX_LENGTH
        ),
        hashes.SHA256()
    )


def load_private_key(private_pem, password):
    return serialization.load_pem_private_key(
        private_pem,
        password=password,
        backend=backends.default_backend()
    )


def load_public_key(public_pem):
    return serialization.load_pem_public_key(
        public_pem,
        backend=backends.default_backend()
    )


def valid_signature(signature, message, public_key):
    try:
        public_key.verify(
            signature,
            message,
            padding.PSS(
                mgf=padding.MGF1(hashes.SHA256()),
                salt_length=padding.PSS.MAX_LENGTH
            ),
            hashes.SHA256()
        )
    except cryptography.exceptions.InvalidSignature:
        return False
    else:
        return True


if __name__ == "__main__":
    def test():
        private = generate_key()
        public = generate_public_key(private)
        private_pem = generate_private_pem(private, b"password123")
        public_pem = generate_public_pem(public)

        loaded_private = load_private_key(private_pem, b"password123")
        loaded_public = load_public_key(public_pem)

        message = b"Hello"
        signature = sign(message, private)

        # VERIFY MESSAGE USING LOADED PRIVATE KEY
        assert valid_signature(signature, message, generate_public_key(loaded_private))

        # VERIFY MESSAGE USING LOADED PUBLIC KEY
        assert valid_signature(signature, message, loaded_public)

        # VERIFY INVALID MESSAGE USING LOADED PRIVATE KEY
        assert not valid_signature(signature, b"Hello World", generate_public_key(loaded_private))

        # VERIFY INVALID MESSAGE USING LOADED PUBLIC KEY
        assert not valid_signature(signature, b"Hello World", loaded_public)

        print("Success")

    test()