from authlib.common.encoding import to_bytes
from cryptography.hazmat.primitives.serialization import (
Encoding, PrivateFormat, PublicFormat,
BestAvailableEncryption, NoEncryption,
)
from ._cryptography_key import load_pem_key
from .base_key import Key
class AsymmetricKey(Key):
"""This is the base class for a JSON Web Key."""
PUBLIC_KEY_FIELDS = []
PRIVATE_KEY_FIELDS = []
PRIVATE_KEY_CLS = bytes
PUBLIC_KEY_CLS = bytes
SSH_PUBLIC_PREFIX = b''
def __init__(self, private_key=None, public_key=None, options=None):
super().__init__(options)
self.private_key = private_key
self.public_key = public_key
@property
def public_only(self):
if self.private_key:
return False
if 'd' in self.tokens:
return False
return True
def get_op_key(self, operation):
"""Get the raw key for the given key_op. This method will also
check if the given key_op is supported by this key.
:param operation: key operation value, such as "sign", "encrypt".
:return: raw key
"""
self.check_key_op(operation)
if operation in self.PUBLIC_KEY_OPS:
return self.get_public_key()
return self.get_private_key()
def get_public_key(self):
if self.public_key:
return self.public_key
private_key = self.get_private_key()
if private_key:
return private_key.public_key()
return self.public_key
def get_private_key(self):
if self.private_key:
return self.private_key
if self.tokens:
self.load_raw_key()
return self.private_key
def load_raw_key(self):
if 'd' in self.tokens:
self.private_key = self.load_private_key()
else:
self.public_key = self.load_public_key()
def load_dict_key(self):
if self.private_key:
self._dict_data.update(self.dumps_private_key())
else:
self._dict_data.update(self.dumps_public_key())
def dumps_private_key(self):
raise NotImplementedError()
def dumps_public_key(self):
raise NotImplementedError()
def load_private_key(self):
raise NotImplementedError()
def load_public_key(self):
raise NotImplementedError()
def as_dict(self, is_private=False, **params):
"""Represent this key as a dict of the JSON Web Key."""
tokens = self.tokens
if is_private and 'd' not in tokens:
raise ValueError('This is a public key')
kid = tokens.get('kid')
if 'd' in tokens and not is_private:
# filter out private fields
tokens = {k: tokens[k] for k in tokens if k in self.PUBLIC_KEY_FIELDS}
tokens['kty'] = self.kty
if kid:
tokens['kid'] = kid
if not kid:
tokens['kid'] = self.thumbprint()
tokens.update(params)
return tokens
def as_key(self, is_private=False):
"""Represent this key as raw key."""
if is_private:
return self.get_private_key()
return self.get_public_key()
def as_bytes(self, encoding=None, is_private=False, password=None):
"""Export key into PEM/DER format bytes.
:param encoding: "PEM" or "DER"
:param is_private: export private key or public key
:param password: encrypt private key with password
:return: bytes
"""
if encoding is None or encoding == 'PEM':
encoding = Encoding.PEM
elif encoding == 'DER':
encoding = Encoding.DER
else:
raise ValueError(f'Invalid encoding: {encoding!r}')
raw_key = self.as_key(is_private)
if is_private:
if not raw_key:
raise ValueError('This is a public key')
if password is None:
encryption_algorithm = NoEncryption()
else:
encryption_algorithm = BestAvailableEncryption(to_bytes(password))
return raw_key.private_bytes(
encoding=encoding,
format=PrivateFormat.PKCS8,
encryption_algorithm=encryption_algorithm,
)
return raw_key.public_bytes(
encoding=encoding,
format=PublicFormat.SubjectPublicKeyInfo,
)
def as_pem(self, is_private=False, password=None):
return self.as_bytes(is_private=is_private, password=password)
def as_der(self, is_private=False, password=None):
return self.as_bytes(encoding='DER', is_private=is_private, password=password)
@classmethod
def import_dict_key(cls, raw, options=None):
cls.check_required_fields(raw)
key = cls(options=options)
key._dict_data = raw
return key
@classmethod
def import_key(cls, raw, options=None):
if isinstance(raw, cls):
if options is not None:
raw.options.update(options)
return raw
if isinstance(raw, cls.PUBLIC_KEY_CLS):
key = cls(public_key=raw, options=options)
elif isinstance(raw, cls.PRIVATE_KEY_CLS):
key = cls(private_key=raw, options=options)
elif isinstance(raw, dict):
key = cls.import_dict_key(raw, options)
else:
if options is not None:
password = options.pop('password', None)
else:
password = None
raw_key = load_pem_key(raw, cls.SSH_PUBLIC_PREFIX, password=password)
if isinstance(raw_key, cls.PUBLIC_KEY_CLS):
key = cls(public_key=raw_key, options=options)
elif isinstance(raw_key, cls.PRIVATE_KEY_CLS):
key = cls(private_key=raw_key, options=options)
else:
raise ValueError('Invalid data for importing key')
return key
@classmethod
def validate_raw_key(cls, key):
return isinstance(key, cls.PUBLIC_KEY_CLS) or isinstance(key, cls.PRIVATE_KEY_CLS)
@classmethod
def generate_key(cls, crv_or_size, options=None, is_private=False):
raise NotImplementedError()