#!/usr/bin/env python3

# Credits: https://github.com/jpf/okta-jwks-to-pem

import argparse
import base64
import json
import struct
import sys

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization

input_key = ""

parser = argparse.ArgumentParser()
parser.add_argument("--kid", default="example", help="JWK Key ID to include in output.")
parser.add_argument(
    "key", metavar="FILE", nargs="?", help='PEM encoded public key. Use "-" for STDIN.'
)
args = parser.parse_args()

input_file = args.key

if input_file is None:
    parser.print_help()
    sys.exit(1)

with open(input_file, "rb") as key_file:
    input_key = key_file.read()


def long2intarr(long_int):
    _bytes = []
    while long_int:
        long_int, r = divmod(long_int, 256)
        _bytes.insert(0, r)
    return _bytes


def long_to_base64(n):
    bys = long2intarr(n)
    data = struct.pack("%sB" % len(bys), *bys)
    if not len(data):
        data = "\x00"
    s = base64.urlsafe_b64encode(data).rstrip(b"=")
    return s.decode("ascii")


public_key = serialization.load_pem_public_key(input_key, default_backend())

public_numbers = public_key.public_numbers()

jwk = {
    "alg": "RS256",
    "e": None,
    "n": None,
    "kid": args.kid,
    "kty": "RSA",
    "use": "sig",
}

jwk["n"] = long_to_base64(public_numbers.n)
jwk["e"] = long_to_base64(public_numbers.e)

print(json.dumps(jwk))
