codekingpro/portable-devtools

1from authlib.common.encoding import (2    to_bytes, to_unicode,3    urlsafe_b64encode, urlsafe_b64decode,4)5from authlib.common.security import generate_token6from ..rfc7517 import Key7 8 9class OctKey(Key):10    """Key class of the ``oct`` key type."""11 12    kty = 'oct'13    REQUIRED_JSON_FIELDS = ['k']14 15    def __init__(self, raw_key=None, options=None):16        super().__init__(options)17        self.raw_key = raw_key18 19    @property20    def public_only(self):21        return False22 23    def get_op_key(self, operation):24        """Get the raw key for the given key_op. This method will also25        check if the given key_op is supported by this key.26 27        :param operation: key operation value, such as "sign", "encrypt".28        :return: raw key29        """30        self.check_key_op(operation)31        if not self.raw_key:32            self.load_raw_key()33        return self.raw_key34 35    def load_raw_key(self):36        self.raw_key = urlsafe_b64decode(to_bytes(self.tokens['k']))37 38    def load_dict_key(self):39        k = to_unicode(urlsafe_b64encode(self.raw_key))40        self._dict_data = {'kty': self.kty, 'k': k}41 42    def as_dict(self, is_private=False, **params):43        tokens = self.tokens44        if 'kid' not in tokens:45            tokens['kid'] = self.thumbprint()46 47        tokens.update(params)48        return tokens49 50    @classmethod51    def validate_raw_key(cls, key):52        return isinstance(key, bytes)53 54    @classmethod55    def import_key(cls, raw, options=None):56        """Import a key from bytes, string, or dict data."""57        if isinstance(raw, cls):58            if options is not None:59                raw.options.update(options)60            return raw61 62        if isinstance(raw, dict):63            cls.check_required_fields(raw)64            key = cls(options=options)65            key._dict_data = raw66        else:67            raw_key = to_bytes(raw)68            key = cls(raw_key=raw_key, options=options)69        return key70 71    @classmethod72    def generate_key(cls, key_size=256, options=None, is_private=True):73        """Generate a ``OctKey`` with the given bit size."""74        if not is_private:75            raise ValueError('oct key can not be generated as public')76 77        if key_size % 8 != 0:78            raise ValueError('Invalid bit size for oct key')79 80        return cls.import_key(generate_token(key_size // 8), options)81