217 lines
6.3 KiB
Python
217 lines
6.3 KiB
Python
import warnings
|
|
from typing import cast, overload, Literal, Union, Callable, Protocol
|
|
from ._keys import (
|
|
JWKRegistry,
|
|
KeySet,
|
|
Key,
|
|
KeySetSerialization,
|
|
)
|
|
from ._rfc7517.pem import import_from_pem_key, import_from_ssh_key
|
|
from ._rfc7517.types import AnyKey, DictKey, KeyParameters
|
|
from ._rfc7518.oct_key import OctKey
|
|
from ._rfc7518.rsa_key import RSAKey
|
|
from ._rfc7518.ec_key import ECKey
|
|
from ._rfc8037.okp_key import OKPKey
|
|
from ._rfc8812 import register_secp256k1
|
|
from ._rfc7638 import calculate_thumbprint as thumbprint
|
|
from ._rfc9278 import calculate_thumbprint_uri as thumbprint_uri
|
|
from .errors import SecurityWarning, InvalidKeyTypeError
|
|
from .registry import Header
|
|
from .util import to_bytes
|
|
|
|
|
|
__all__ = [
|
|
# types
|
|
"Key",
|
|
"DictKey",
|
|
"KeyParameters",
|
|
"KeyCallable",
|
|
"KeyFlexible",
|
|
"KeySetSerialization",
|
|
"KeyBase",
|
|
"GuestProtocol",
|
|
# modules
|
|
"JWKRegistry",
|
|
"OctKey",
|
|
"RSAKey",
|
|
"ECKey",
|
|
"OKPKey",
|
|
"KeySet",
|
|
# methods
|
|
"guess_key",
|
|
"import_key",
|
|
"generate_key",
|
|
"thumbprint",
|
|
"thumbprint_uri",
|
|
]
|
|
|
|
register_secp256k1()
|
|
|
|
|
|
class GuestProtocol(Protocol): # pragma: no cover
|
|
def headers(self) -> Header: ...
|
|
|
|
def set_kid(self, kid: str) -> None: ...
|
|
|
|
|
|
KeyBase = Union[Key, KeySet]
|
|
KeyCallable = Callable[[GuestProtocol], KeyBase]
|
|
KeyFlexible = Union[KeyBase, KeyCallable]
|
|
|
|
|
|
def guess_key(
|
|
key: KeyFlexible,
|
|
obj: GuestProtocol,
|
|
random: bool = False,
|
|
use: Literal["sig", "enc"] | None = None,
|
|
) -> Key:
|
|
"""Guess key from a various sources.
|
|
|
|
:param key: a very flexible key
|
|
:param obj: a protocol that has ``headers`` and ``set_kid`` methods
|
|
:param random: pick a random key from key set
|
|
:param use: optional "use" value
|
|
"""
|
|
resolved_key: KeyBase
|
|
if callable(key):
|
|
resolved_key = key(obj)
|
|
else:
|
|
resolved_key = key
|
|
|
|
if isinstance(resolved_key, (OctKey, RSAKey, ECKey, OKPKey)):
|
|
return resolved_key
|
|
elif isinstance(resolved_key, KeySet):
|
|
headers = obj.headers()
|
|
kid: str | None = headers.get("kid")
|
|
|
|
parameters: KeyParameters = {"alg": headers["alg"]}
|
|
if use:
|
|
parameters["use"] = use
|
|
|
|
if not kid and random:
|
|
# choose one key by random
|
|
return_key = resolved_key.pick_random_key(headers["alg"], parameters)
|
|
if return_key is None:
|
|
raise ValueError("Invalid key")
|
|
return_key.ensure_kid()
|
|
obj.set_kid(cast(str, return_key.kid))
|
|
else:
|
|
return_key = resolved_key.get_by_kid(kid, parameters)
|
|
return return_key
|
|
else:
|
|
raise ValueError("Invalid key")
|
|
|
|
|
|
@overload
|
|
def import_key(data: AnyKey, key_type: Literal["oct"], parameters: KeyParameters | None = None) -> OctKey: ...
|
|
|
|
|
|
@overload
|
|
def import_key(data: AnyKey, key_type: Literal["RSA"], parameters: KeyParameters | None = None) -> RSAKey: ...
|
|
|
|
|
|
@overload
|
|
def import_key(data: AnyKey, key_type: Literal["EC"], parameters: KeyParameters | None = None) -> ECKey: ...
|
|
|
|
|
|
@overload
|
|
def import_key(data: AnyKey, key_type: Literal["OKP"], parameters: KeyParameters | None = None) -> OKPKey: ...
|
|
|
|
|
|
@overload
|
|
def import_key(data: AnyKey, key_type: None = None, parameters: KeyParameters | None = None) -> Key: ...
|
|
|
|
|
|
def import_key(
|
|
data: AnyKey,
|
|
key_type: Literal["oct", "RSA", "EC", "OKP"] | None = None,
|
|
parameters: KeyParameters | None = None,
|
|
) -> Key:
|
|
"""Importing a key from bytes, string, and dict. When ``value`` is a dict,
|
|
this method can tell the key type automatically, otherwise, developers
|
|
SHOULD pass the ``key_type`` themselves.
|
|
|
|
:param data: the key data in bytes, string, or dict.
|
|
:param key_type: an optional key type in string.
|
|
:param parameters: extra key parameters
|
|
:return: OctKey, RSAKey, ECKey, or OKPKey
|
|
"""
|
|
if isinstance(data, (str, bytes)) and key_type is None:
|
|
warnings.warn("Using implicit key type is not recommended.", SecurityWarning)
|
|
|
|
value = to_bytes(data)
|
|
ssh_types = tuple(
|
|
cls.binding.ssh_type for cls in JWKRegistry.key_types.values() if hasattr(cls.binding, "ssh_type")
|
|
)
|
|
if value.startswith(ssh_types):
|
|
try:
|
|
raw_key = import_from_ssh_key(value)
|
|
except ValueError:
|
|
return OctKey.import_key(value, parameters)
|
|
else:
|
|
try:
|
|
raw_key = import_from_pem_key(value)
|
|
except ValueError:
|
|
return OctKey.import_key(value, parameters)
|
|
|
|
for cls in JWKRegistry.key_types.values():
|
|
if hasattr(cls.binding, "check_cryptography_key") and cls.binding.check_cryptography_key(raw_key):
|
|
return cls(raw_key, data, parameters)
|
|
raise InvalidKeyTypeError("Not a key of any supported type") # pragma: no cover
|
|
return JWKRegistry.import_key(data, key_type, parameters)
|
|
|
|
|
|
@overload
|
|
def generate_key(
|
|
key_type: Literal["oct"],
|
|
crv_or_size: int | None = None,
|
|
parameters: KeyParameters | None = None,
|
|
private: bool = True,
|
|
auto_kid: bool = False,
|
|
) -> OctKey: ...
|
|
|
|
|
|
@overload
|
|
def generate_key(
|
|
key_type: Literal["RSA"],
|
|
crv_or_size: int | None = None,
|
|
parameters: KeyParameters | None = None,
|
|
private: bool = True,
|
|
auto_kid: bool = False,
|
|
) -> RSAKey: ...
|
|
|
|
|
|
@overload
|
|
def generate_key(
|
|
key_type: Literal["EC"],
|
|
crv_or_size: Literal["P-256", "P-384", "P-521", "secp256k1"] | None = None,
|
|
parameters: KeyParameters | None = None,
|
|
private: bool = True,
|
|
auto_kid: bool = False,
|
|
) -> ECKey: ...
|
|
|
|
|
|
@overload
|
|
def generate_key(
|
|
key_type: Literal["OKP"],
|
|
crv_or_size: Literal["Ed25519", "Ed448", "X25519", "X448"] | None = None,
|
|
parameters: KeyParameters | None = None,
|
|
private: bool = True,
|
|
auto_kid: bool = False,
|
|
) -> OKPKey: ...
|
|
|
|
|
|
def generate_key(
|
|
key_type: Literal["oct", "RSA", "EC", "OKP"],
|
|
crv_or_size: str | int | None = None,
|
|
parameters: KeyParameters | None = None,
|
|
private: bool = True,
|
|
auto_kid: bool = False,
|
|
) -> Key:
|
|
"""Generating key according to the given key type. When ``key_type`` is
|
|
"oct" and "RSA", the second parameter SHOULD be a key size in bits.
|
|
When ``key_type`` is "EC" and "OKP", the second
|
|
parameter SHOULD be a "crv" string.
|
|
"""
|
|
return JWKRegistry.generate_key(key_type, crv_or_size, parameters, private, auto_kid)
|