codekingpro/portable-devtools

1"""2Ephemeral Elliptic Curve Diffie-Hellman (ECDH) key exchange3RFC 5656, Section 44"""5 6from hashlib import sha256, sha384, sha5127from paramiko.common import byte_chr8from paramiko.message import Message9from paramiko.ssh_exception import SSHException10from cryptography.hazmat.backends import default_backend11from cryptography.hazmat.primitives.asymmetric import ec12from cryptography.hazmat.primitives import serialization13from binascii import hexlify14 15_MSG_KEXECDH_INIT, _MSG_KEXECDH_REPLY = range(30, 32)16c_MSG_KEXECDH_INIT, c_MSG_KEXECDH_REPLY = [byte_chr(c) for c in range(30, 32)]17 18 19class KexNistp256:20 21    name = "ecdh-sha2-nistp256"22    hash_algo = sha25623    curve = ec.SECP256R1()24 25    def __init__(self, transport):26        self.transport = transport27        # private key, client public and server public keys28        self.P = 029        self.Q_C = None30        self.Q_S = None31 32    def start_kex(self):33        self._generate_key_pair()34        if self.transport.server_mode:35            self.transport._expect_packet(_MSG_KEXECDH_INIT)36            return37        m = Message()38        m.add_byte(c_MSG_KEXECDH_INIT)39        # SEC1: V2.0  2.3.3 Elliptic-Curve-Point-to-Octet-String Conversion40        m.add_string(41            self.Q_C.public_bytes(42                serialization.Encoding.X962,43                serialization.PublicFormat.UncompressedPoint,44            )45        )46        self.transport._send_message(m)47        self.transport._expect_packet(_MSG_KEXECDH_REPLY)48 49    def parse_next(self, ptype, m):50        if self.transport.server_mode and (ptype == _MSG_KEXECDH_INIT):51            return self._parse_kexecdh_init(m)52        elif not self.transport.server_mode and (ptype == _MSG_KEXECDH_REPLY):53            return self._parse_kexecdh_reply(m)54        raise SSHException(55            "KexECDH asked to handle packet type {:d}".format(ptype)56        )57 58    def _generate_key_pair(self):59        self.P = ec.generate_private_key(self.curve, default_backend())60        if self.transport.server_mode:61            self.Q_S = self.P.public_key()62            return63        self.Q_C = self.P.public_key()64 65    def _parse_kexecdh_init(self, m):66        Q_C_bytes = m.get_string()67        self.Q_C = ec.EllipticCurvePublicKey.from_encoded_point(68            self.curve, Q_C_bytes69        )70        K_S = self.transport.get_server_key().asbytes()71        K = self.P.exchange(ec.ECDH(), self.Q_C)72        K = int(hexlify(K), 16)73        # compute exchange hash74        hm = Message()75        hm.add(76            self.transport.remote_version,77            self.transport.local_version,78            self.transport.remote_kex_init,79            self.transport.local_kex_init,80        )81        hm.add_string(K_S)82        hm.add_string(Q_C_bytes)83        # SEC1: V2.0  2.3.3 Elliptic-Curve-Point-to-Octet-String Conversion84        hm.add_string(85            self.Q_S.public_bytes(86                serialization.Encoding.X962,87                serialization.PublicFormat.UncompressedPoint,88            )89        )90        hm.add_mpint(int(K))91        H = self.hash_algo(hm.asbytes()).digest()92        self.transport._set_K_H(K, H)93        sig = self.transport.get_server_key().sign_ssh_data(94            H, self.transport.host_key_type95        )96        # construct reply97        m = Message()98        m.add_byte(c_MSG_KEXECDH_REPLY)99        m.add_string(K_S)100        m.add_string(101            self.Q_S.public_bytes(102                serialization.Encoding.X962,103                serialization.PublicFormat.UncompressedPoint,104            )105        )106        m.add_string(sig)107        self.transport._send_message(m)108        self.transport._activate_outbound()109 110    def _parse_kexecdh_reply(self, m):111        K_S = m.get_string()112        Q_S_bytes = m.get_string()113        self.Q_S = ec.EllipticCurvePublicKey.from_encoded_point(114            self.curve, Q_S_bytes115        )116        sig = m.get_binary()117        K = self.P.exchange(ec.ECDH(), self.Q_S)118        K = int(hexlify(K), 16)119        # compute exchange hash and verify signature120        hm = Message()121        hm.add(122            self.transport.local_version,123            self.transport.remote_version,124            self.transport.local_kex_init,125            self.transport.remote_kex_init,126        )127        hm.add_string(K_S)128        # SEC1: V2.0  2.3.3 Elliptic-Curve-Point-to-Octet-String Conversion129        hm.add_string(130            self.Q_C.public_bytes(131                serialization.Encoding.X962,132                serialization.PublicFormat.UncompressedPoint,133            )134        )135        hm.add_string(Q_S_bytes)136        hm.add_mpint(K)137        self.transport._set_K_H(K, self.hash_algo(hm.asbytes()).digest())138        self.transport._verify_key(K_S, sig)139        self.transport._activate_outbound()140 141 142class KexNistp384(KexNistp256):143    name = "ecdh-sha2-nistp384"144    hash_algo = sha384145    curve = ec.SECP384R1()146 147 148class KexNistp521(KexNistp256):149    name = "ecdh-sha2-nistp521"150    hash_algo = sha512151    curve = ec.SECP521R1()152