tokens.py

import os
import time
import hashlib
import base64
import subprocess
from jose import jwt

SIGNING_KEY_PATH = os.environ.get('SIGNING_KEY_PATH')
SIGNING_KEY_TYPE = os.environ.get('SIGNING_KEY_TYPE')
SIGNING_KEY_ALG = os.environ.get('SIGNING_KEY_ALG')
SIGNING_KEY = open(SIGNING_KEY_PATH).read()

ISSUER = os.environ.get('ISSUER')
TOKEN_EXPIRATION = os.environ.get('TOKEN_EXPIRATION')
TOKEN_TYPE = os.environ.get('TOKEN_TYPE')


def run_command(command):
    process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return process.communicate()


def key_id_encode(the_bytes):
    source = base64.b32encode(the_bytes).rstrip("=")
    result = []
    for i in xrange(0, len(source), 4):
        start = i
        end = start+4
        result.append(source[start:end])
    return ":".join(result)


def kid_from_crypto_key(private_key_path, key_type='EC'):
    """
    python implementation of
    https://github.com/jlhawn/libtrust/blob/master/util.go#L192
    returns a distinct identifier which is unique to
    the public key derived from this private key.
    The format generated by this library is a base32 encoding of a 240 bit
    hash of the public key data divided into 12 groups like so:
    ABCD:EFGH:IJKL:MNOP:QRST:UVWX:YZ23:4567:ABCD:EFGH:IJKL:MNOP
    """
    algorithm = hashlib.sha256()
    if key_type == 'EC':
        der, msg = run_command(['openssl', 'ec', '-in', private_key_path,
                                '-pubout', '-outform', 'DER'])

    elif key_type == 'RSA':
        der, msg = run_command(['openssl', 'rsa', '-in', private_key_path,
                                '-pubout', '-outform', 'DER'])

    else:
        raise Exception("Key type not supported")

    if not der:
        raise Exception(msg)

    algorithm.update(der)
    return key_id_encode(algorithm.digest()[:30])


class Token(object):
    def __init__(self, service, access_type="", access_name="",
                 access_actions=None, subject=''):
        if access_actions is None:
            access_actions = []

        self.issuer = ISSUER
        self.signing_key = SIGNING_KEY
        self.signing_key_path = SIGNING_KEY_PATH
        self.signing_key_type = SIGNING_KEY_TYPE
        self.signing_key_alg = SIGNING_KEY_ALG
        self.token_expiration = TOKEN_EXPIRATION
        self.token_type = TOKEN_TYPE
        self.header = {
            'typ': self.token_type,
            'alg': self.signing_key_alg,
            'kid': kid_from_crypto_key(self.signing_key_path,
                                       self.signing_key_type)
        }
        self.claim = {
            'iss': self.issuer,
            'sub': subject,
            'aud': service,
            'exp': int(time.time()) + int(self.token_expiration),
            'nbf': int(time.time()) - 30,
            'iat': int(time.time()),
            'access': [
                {
                    'type': access_type,
                    'name': access_name,
                    'actions': access_actions
                }
            ]
        }

    def set_header(self, header):
        self.header = header

    def get_header(self):
        return self.header

    def set_claim(self, claim):
        self.claim = claim

    def get_claim(self):
        return self.claim

    def encode_token(self):
        token = jwt.encode(self.claim, self.signing_key,
                           algorithm=self.signing_key_alg,
                           headers=self.header)

        return token