codekingpro/portable-devtools

1from authlib.common.encoding import to_bytes2from cryptography.hazmat.primitives.serialization import (3    Encoding, PrivateFormat, PublicFormat,4    BestAvailableEncryption, NoEncryption,5)6from ._cryptography_key import load_pem_key7from .base_key import Key8 9 10class AsymmetricKey(Key):11    """This is the base class for a JSON Web Key."""12    PUBLIC_KEY_FIELDS = []13    PRIVATE_KEY_FIELDS = []14    PRIVATE_KEY_CLS = bytes15    PUBLIC_KEY_CLS = bytes16    SSH_PUBLIC_PREFIX = b''17 18    def __init__(self, private_key=None, public_key=None, options=None):19        super().__init__(options)20        self.private_key = private_key21        self.public_key = public_key22 23    @property24    def public_only(self):25        if self.private_key:26            return False27        if 'd' in self.tokens:28            return False29        return True30 31    def get_op_key(self, operation):32        """Get the raw key for the given key_op. This method will also33        check if the given key_op is supported by this key.34 35        :param operation: key operation value, such as "sign", "encrypt".36        :return: raw key37        """38        self.check_key_op(operation)39        if operation in self.PUBLIC_KEY_OPS:40            return self.get_public_key()41        return self.get_private_key()42 43    def get_public_key(self):44        if self.public_key:45            return self.public_key46 47        private_key = self.get_private_key()48        if private_key:49            return private_key.public_key()50 51        return self.public_key52 53    def get_private_key(self):54        if self.private_key:55            return self.private_key56 57        if self.tokens:58            self.load_raw_key()59        return self.private_key60 61    def load_raw_key(self):62        if 'd' in self.tokens:63            self.private_key = self.load_private_key()64        else:65            self.public_key = self.load_public_key()66 67    def load_dict_key(self):68        if self.private_key:69            self._dict_data.update(self.dumps_private_key())70        else:71            self._dict_data.update(self.dumps_public_key())72 73    def dumps_private_key(self):74        raise NotImplementedError()75 76    def dumps_public_key(self):77        raise NotImplementedError()78 79    def load_private_key(self):80        raise NotImplementedError()81 82    def load_public_key(self):83        raise NotImplementedError()84 85    def as_dict(self, is_private=False, **params):86        """Represent this key as a dict of the JSON Web Key."""87        tokens = self.tokens88        if is_private and 'd' not in tokens:89            raise ValueError('This is a public key')90 91        kid = tokens.get('kid')92        if 'd' in tokens and not is_private:93            # filter out private fields94            tokens = {k: tokens[k] for k in tokens if k in self.PUBLIC_KEY_FIELDS}95            tokens['kty'] = self.kty96            if kid:97                tokens['kid'] = kid98 99        if not kid:100            tokens['kid'] = self.thumbprint()101 102        tokens.update(params)103        return tokens104 105    def as_key(self, is_private=False):106        """Represent this key as raw key."""107        if is_private:108            return self.get_private_key()109        return self.get_public_key()110 111    def as_bytes(self, encoding=None, is_private=False, password=None):112        """Export key into PEM/DER format bytes.113 114        :param encoding: "PEM" or "DER"115        :param is_private: export private key or public key116        :param password: encrypt private key with password117        :return: bytes118        """119 120        if encoding is None or encoding == 'PEM':121            encoding = Encoding.PEM122        elif encoding == 'DER':123            encoding = Encoding.DER124        else:125            raise ValueError(f'Invalid encoding: {encoding!r}')126 127        raw_key = self.as_key(is_private)128        if is_private:129            if not raw_key:130                raise ValueError('This is a public key')131            if password is None:132                encryption_algorithm = NoEncryption()133            else:134                encryption_algorithm = BestAvailableEncryption(to_bytes(password))135            return raw_key.private_bytes(136                encoding=encoding,137                format=PrivateFormat.PKCS8,138                encryption_algorithm=encryption_algorithm,139            )140        return raw_key.public_bytes(141            encoding=encoding,142            format=PublicFormat.SubjectPublicKeyInfo,143        )144 145    def as_pem(self, is_private=False, password=None):146        return self.as_bytes(is_private=is_private, password=password)147 148    def as_der(self, is_private=False, password=None):149        return self.as_bytes(encoding='DER', is_private=is_private, password=password)150 151    @classmethod152    def import_dict_key(cls, raw, options=None):153        cls.check_required_fields(raw)154        key = cls(options=options)155        key._dict_data = raw156        return key157 158    @classmethod159    def import_key(cls, raw, options=None):160        if isinstance(raw, cls):161            if options is not None:162                raw.options.update(options)163            return raw164 165        if isinstance(raw, cls.PUBLIC_KEY_CLS):166            key = cls(public_key=raw, options=options)167        elif isinstance(raw, cls.PRIVATE_KEY_CLS):168            key = cls(private_key=raw, options=options)169        elif isinstance(raw, dict):170            key = cls.import_dict_key(raw, options)171        else:172            if options is not None:173                password = options.pop('password', None)174            else:175                password = None176            raw_key = load_pem_key(raw, cls.SSH_PUBLIC_PREFIX, password=password)177            if isinstance(raw_key, cls.PUBLIC_KEY_CLS):178                key = cls(public_key=raw_key, options=options)179            elif isinstance(raw_key, cls.PRIVATE_KEY_CLS):180                key = cls(private_key=raw_key, options=options)181            else:182                raise ValueError('Invalid data for importing key')183        return key184 185    @classmethod186    def validate_raw_key(cls, key):187        return isinstance(key, cls.PUBLIC_KEY_CLS) or isinstance(key, cls.PRIVATE_KEY_CLS)188 189    @classmethod190    def generate_key(cls, crv_or_size, options=None, is_private=False):191        raise NotImplementedError()192