first commit
This commit is contained in:
linkitapps
@@ -0,0 +1,216 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user