codekingpro/portable-devtools

1from authlib.common.encoding import json_loads2from .key_set import KeySet3from ._cryptography_key import load_pem_key4 5 6class JsonWebKey:7    JWK_KEY_CLS = {}8 9    @classmethod10    def generate_key(cls, kty, crv_or_size, options=None, is_private=False):11        """Generate a Key with the given key type, curve name or bit size.12 13        :param kty: string of ``oct``, ``RSA``, ``EC``, ``OKP``14        :param crv_or_size: curve name or bit size15        :param options: a dict of other options for Key16        :param is_private: create a private key or public key17        :return: Key instance18        """19        key_cls = cls.JWK_KEY_CLS[kty]20        return key_cls.generate_key(crv_or_size, options, is_private)21 22    @classmethod23    def import_key(cls, raw, options=None):24        """Import a Key from bytes, string, PEM or dict.25 26        :return: Key instance27        """28        kty = None29        if options is not None:30            kty = options.get('kty')31 32        if kty is None and isinstance(raw, dict):33            kty = raw.get('kty')34 35        if kty is None:36            raw_key = load_pem_key(raw)37            for _kty in cls.JWK_KEY_CLS:38                key_cls = cls.JWK_KEY_CLS[_kty]39                if key_cls.validate_raw_key(raw_key):40                    return key_cls.import_key(raw_key, options)41 42        key_cls = cls.JWK_KEY_CLS[kty]43        return key_cls.import_key(raw, options)44 45    @classmethod46    def import_key_set(cls, raw):47        """Import KeySet from string, dict or a list of keys.48 49        :return: KeySet instance50        """51        raw = _transform_raw_key(raw)52        if isinstance(raw, dict) and 'keys' in raw:53            keys = raw.get('keys')54            return KeySet([cls.import_key(k) for k in keys])55        raise ValueError('Invalid key set format')56 57 58def _transform_raw_key(raw):59    if isinstance(raw, str) and \60            raw.startswith('{') and raw.endswith('}'):61        return json_loads(raw)62    elif isinstance(raw, (tuple, list)):63        return {'keys': raw}64    return raw65