122 lines
4.2 KiB
Python
122 lines
4.2 KiB
Python
|
from typing import Dict, Optional, Tuple, Type, Union
|
||
|
|
|
||
|
|
import dns.name
|
||
|
|
|
||
|
|
try:
|
||
|
|
from dns.dnssecalgs.base import GenericPrivateKey
|
||
|
|
from dns.dnssecalgs.dsa import PrivateDSA, PrivateDSANSEC3SHA1
|
||
|
|
from dns.dnssecalgs.ecdsa import PrivateECDSAP256SHA256, PrivateECDSAP384SHA384
|
||
|
|
from dns.dnssecalgs.eddsa import PrivateED448, PrivateED25519
|
||
|
|
from dns.dnssecalgs.rsa import (
|
||
|
|
PrivateRSAMD5,
|
||
|
|
PrivateRSASHA1,
|
||
|
|
PrivateRSASHA1NSEC3SHA1,
|
||
|
|
PrivateRSASHA256,
|
||
|
|
PrivateRSASHA512,
|
||
|
|
)
|
||
|
|
|
||
|
|
_have_cryptography = True
|
||
|
|
except ImportError:
|
||
|
|
_have_cryptography = False
|
||
|
|
|
||
|
|
from dns.dnssectypes import Algorithm
|
||
|
|
from dns.exception import UnsupportedAlgorithm
|
||
|
|
from dns.rdtypes.ANY.DNSKEY import DNSKEY
|
||
|
|
|
||
|
|
AlgorithmPrefix = Optional[Union[bytes, dns.name.Name]]
|
||
|
|
|
||
|
|
algorithms: Dict[Tuple[Algorithm, AlgorithmPrefix], Type[GenericPrivateKey]] = {}
|
||
|
|
if _have_cryptography:
|
||
|
|
algorithms.update(
|
||
|
|
{
|
||
|
|
(Algorithm.RSAMD5, None): PrivateRSAMD5,
|
||
|
|
(Algorithm.DSA, None): PrivateDSA,
|
||
|
|
(Algorithm.RSASHA1, None): PrivateRSASHA1,
|
||
|
|
(Algorithm.DSANSEC3SHA1, None): PrivateDSANSEC3SHA1,
|
||
|
|
(Algorithm.RSASHA1NSEC3SHA1, None): PrivateRSASHA1NSEC3SHA1,
|
||
|
|
(Algorithm.RSASHA256, None): PrivateRSASHA256,
|
||
|
|
(Algorithm.RSASHA512, None): PrivateRSASHA512,
|
||
|
|
(Algorithm.ECDSAP256SHA256, None): PrivateECDSAP256SHA256,
|
||
|
|
(Algorithm.ECDSAP384SHA384, None): PrivateECDSAP384SHA384,
|
||
|
|
(Algorithm.ED25519, None): PrivateED25519,
|
||
|
|
(Algorithm.ED448, None): PrivateED448,
|
||
|
|
}
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_algorithm_cls(
|
||
|
|
algorithm: Union[int, str], prefix: AlgorithmPrefix = None
|
||
|
|
) -> Type[GenericPrivateKey]:
|
||
|
|
"""Get Private Key class from Algorithm.
|
||
|
|
|
||
|
|
*algorithm*, a ``str`` or ``int`` specifying the DNSKEY algorithm.
|
||
|
|
|
||
|
|
Raises ``UnsupportedAlgorithm`` if the algorithm is unknown.
|
||
|
|
|
||
|
|
Returns a ``dns.dnssecalgs.GenericPrivateKey``
|
||
|
|
"""
|
||
|
|
algorithm = Algorithm.make(algorithm)
|
||
|
|
cls = algorithms.get((algorithm, prefix))
|
||
|
|
if cls:
|
||
|
|
return cls
|
||
|
|
raise UnsupportedAlgorithm(
|
||
|
|
'algorithm "%s" not supported by dnspython' % Algorithm.to_text(algorithm)
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def get_algorithm_cls_from_dnskey(dnskey: DNSKEY) -> Type[GenericPrivateKey]:
|
||
|
|
"""Get Private Key class from DNSKEY.
|
||
|
|
|
||
|
|
*dnskey*, a ``DNSKEY`` to get Algorithm class for.
|
||
|
|
|
||
|
|
Raises ``UnsupportedAlgorithm`` if the algorithm is unknown.
|
||
|
|
|
||
|
|
Returns a ``dns.dnssecalgs.GenericPrivateKey``
|
||
|
|
"""
|
||
|
|
prefix: AlgorithmPrefix = None
|
||
|
|
if dnskey.algorithm == Algorithm.PRIVATEDNS:
|
||
|
|
prefix, _ = dns.name.from_wire(dnskey.key, 0)
|
||
|
|
elif dnskey.algorithm == Algorithm.PRIVATEOID:
|
||
|
|
length = int(dnskey.key[0])
|
||
|
|
prefix = dnskey.key[0 : length + 1]
|
||
|
|
return get_algorithm_cls(dnskey.algorithm, prefix)
|
||
|
|
|
||
|
|
|
||
|
|
def register_algorithm_cls(
|
||
|
|
algorithm: Union[int, str],
|
||
|
|
algorithm_cls: Type[GenericPrivateKey],
|
||
|
|
name: Optional[Union[dns.name.Name, str]] = None,
|
||
|
|
oid: Optional[bytes] = None,
|
||
|
|
) -> None:
|
||
|
|
"""Register Algorithm Private Key class.
|
||
|
|
|
||
|
|
*algorithm*, a ``str`` or ``int`` specifying the DNSKEY algorithm.
|
||
|
|
|
||
|
|
*algorithm_cls*: A `GenericPrivateKey` class.
|
||
|
|
|
||
|
|
*name*, an optional ``dns.name.Name`` or ``str``, for for PRIVATEDNS algorithms.
|
||
|
|
|
||
|
|
*oid*: an optional BER-encoded `bytes` for PRIVATEOID algorithms.
|
||
|
|
|
||
|
|
Raises ``ValueError`` if a name or oid is specified incorrectly.
|
||
|
|
"""
|
||
|
|
if not issubclass(algorithm_cls, GenericPrivateKey):
|
||
|
|
raise TypeError("Invalid algorithm class")
|
||
|
|
algorithm = Algorithm.make(algorithm)
|
||
|
|
prefix: AlgorithmPrefix = None
|
||
|
|
if algorithm == Algorithm.PRIVATEDNS:
|
||
|
|
if name is None:
|
||
|
|
raise ValueError("Name required for PRIVATEDNS algorithms")
|
||
|
|
if isinstance(name, str):
|
||
|
|
name = dns.name.from_text(name)
|
||
|
|
prefix = name
|
||
|
|
elif algorithm == Algorithm.PRIVATEOID:
|
||
|
|
if oid is None:
|
||
|
|
raise ValueError("OID required for PRIVATEOID algorithms")
|
||
|
|
prefix = bytes([len(oid)]) + oid
|
||
|
|
elif name:
|
||
|
|
raise ValueError("Name only supported for PRIVATEDNS algorithm")
|
||
|
|
elif oid:
|
||
|
|
raise ValueError("OID only supported for PRIVATEOID algorithm")
|
||
|
|
algorithms[(algorithm, prefix)] = algorithm_cls
|