Coverage for /pythoncovmergedfiles/medio/medio/src/paramiko/paramiko/kex_mlkem.py: 32%

1""" 

2ML-KEM hybrid key exchange for SSH. 

3 

4Implements the post-quantum / classical hybrid key agreement methods 

5described in draft-ietf-sshm-mlkem-hybrid-kex (e.g. 

6``mlkem768x25519-sha256``). 

7 

8Each method combines an ML-KEM key encapsulation with a traditional 

9ECDH/X25519 key agreement; the final shared secret is the hash of the 

10two component secrets concatenated together. 

11""" 

12 

13from __future__ import annotations 

14 

15import hashlib 

16from typing import TYPE_CHECKING 

17 

18from cryptography.exceptions import UnsupportedAlgorithm 

19from cryptography.hazmat.primitives import constant_time 

20from cryptography.hazmat.primitives.asymmetric.x25519 import ( 

21 X25519PrivateKey, 

22 X25519PublicKey, 

23) 

24 

25try: 

26 from cryptography.hazmat.primitives.asymmetric import mlkem 

27except ImportError: # pragma: no cover - cryptography <48 

28 mlkem = None 

29 

30from paramiko.common import byte_chr 

31from paramiko.message import Message 

32from paramiko.ssh_exception import SSHException 

33 

34if TYPE_CHECKING: 

35 from paramiko.transport import Transport 

36 

37# Per draft-ietf-sshm-mlkem-hybrid-kex, the hybrid kex reuses message 

38# numbers 30 and 31 (named SSH_MSG_KEX_HYBRID_INIT / _REPLY). 

39_MSG_KEX_HYBRID_INIT, _MSG_KEX_HYBRID_REPLY = range(30, 32) 

40c_MSG_KEX_HYBRID_INIT, c_MSG_KEX_HYBRID_REPLY = [ 

41 byte_chr(c) for c in range(30, 32) 

42] 

43 

44 

45class KexMLKEM768X25519: 

46 """ 

47 ``mlkem768x25519-sha256`` hybrid key exchange. 

48 

49 Combines ML-KEM-768 (FIPS 203) with X25519. The shared secret is 

50 ``SHA-256(K_PQ || K_CL)`` where ``K_PQ`` is the ML-KEM secret and 

51 ``K_CL`` is the X25519 shared secret. 

52 """ 

53 

54 name = "mlkem768x25519-sha256" 

55 hash_algo = hashlib.sha256 

56 # Fixed sizes from FIPS 203 (ML-KEM-768) and RFC 7748 (X25519). 

57 _MLKEM_PUBKEY_BYTES = 1184 

58 _MLKEM_CIPHERTEXT_BYTES = 1088 

59 _X25519_PUBKEY_BYTES = 32 

60 _C_INIT_BYTES = _MLKEM_PUBKEY_BYTES + _X25519_PUBKEY_BYTES 

61 _S_REPLY_BYTES = _MLKEM_CIPHERTEXT_BYTES + _X25519_PUBKEY_BYTES 

62 

63 def __init__(self, transport: Transport): 

64 self.transport = transport 

65 # Client-side: our ML-KEM decapsulation key. Server-side: unused. 

66 self.mlkem_key = None 

67 # Our ephemeral X25519 private key (both roles). 

68 self.x25519_key = None 

69 

70 @classmethod 

71 def is_available(cls): 

72 if mlkem is None: 

73 return False 

74 try: 

75 mlkem.MLKEM768PrivateKey.generate() 

76 X25519PrivateKey.generate() 

77 except UnsupportedAlgorithm: 

78 return False 

79 return True 

80 

81 # ---- protocol entry points --------------------------------------- 

82 

83 def start_kex(self): 

84 self.x25519_key = X25519PrivateKey.generate() 

85 if self.transport.server_mode: 

86 self.transport._expect_packet(_MSG_KEX_HYBRID_INIT) 

87 return 

88 self.mlkem_key = mlkem.MLKEM768PrivateKey.generate() 

89 c_init = ( 

90 self.mlkem_key.public_key().public_bytes_raw() 

91 + self.x25519_key.public_key().public_bytes_raw() 

92 ) 

93 m = Message() 

94 m.add_byte(c_MSG_KEX_HYBRID_INIT) 

95 m.add_string(c_init) 

96 self.transport._send_message(m) 

97 self.transport._expect_packet(_MSG_KEX_HYBRID_REPLY) 

98 

99 def parse_next(self, ptype, m): 

100 if self.transport.server_mode and ptype == _MSG_KEX_HYBRID_INIT: 

101 return self._parse_hybrid_init(m) 

102 if not self.transport.server_mode and ptype == _MSG_KEX_HYBRID_REPLY: 

103 return self._parse_hybrid_reply(m) 

104 raise SSHException( 

105 "{} asked to handle packet type {:d}".format( 

106 self.__class__.__name__, ptype 

107 ) 

108 ) 

109 

110 def _x25519_exchange(self, peer_pub_bytes): 

111 peer = X25519PublicKey.from_public_bytes(peer_pub_bytes) 

112 secret = self.x25519_key.exchange(peer) 

113 # Per RFC 8731 (and reaffirmed by the hybrid draft), reject the 

114 # all-zero output that signals a small-order public value. 

115 if constant_time.bytes_eq(secret, b"\x00" * 32): 

116 raise SSHException( 

117 "peer's curve25519 public value has wrong order" 

118 ) 

119 return secret 

120 

121 # ---- server side ------------------------------------------------- 

122 

123 def _parse_hybrid_init(self, m): 

124 c_init = m.get_string() 

125 if len(c_init) != self._C_INIT_BYTES: 

126 raise SSHException( 

127 "Invalid C_INIT length for {}: got {}, expected {}".format( 

128 self.name, len(c_init), self._C_INIT_BYTES 

129 ) 

130 ) 

131 c_pk2 = c_init[: self._MLKEM_PUBKEY_BYTES] 

132 c_pk1 = c_init[self._MLKEM_PUBKEY_BYTES :] 

133 

134 # Encapsulate against the client's ML-KEM public key. 

135 client_mlkem_pub = mlkem.MLKEM768PublicKey.from_public_bytes(c_pk2) 

136 k_pq, s_ct2 = client_mlkem_pub.encapsulate() 

137 

138 # X25519 with the client's ephemeral public value. 

139 k_cl = self._x25519_exchange(c_pk1) 

140 

141 K_bytes = self.hash_algo(k_pq + k_cl).digest() 

142 

143 s_pk1 = self.x25519_key.public_key().public_bytes_raw() 

144 s_reply = s_ct2 + s_pk1 

145 

146 K_S = self.transport.get_server_key().asbytes() 

147 

148 hm = Message() 

149 hm.add( 

150 self.transport.remote_version, 

151 self.transport.local_version, 

152 self.transport.remote_kex_init, 

153 self.transport.local_kex_init, 

154 ) 

155 hm.add_string(K_S) 

156 hm.add_string(c_init) 

157 hm.add_string(s_reply) 

158 # Per the hybrid draft: K is the hash output, encoded as a string. 

159 hm.add_string(K_bytes) 

160 H = self.hash_algo(hm.asbytes()).digest() 

161 

162 self.transport._set_K_H(K_bytes, H) 

163 

164 sig = self.transport.get_server_key().sign_ssh_data( 

165 H, self.transport.host_key_type 

166 ) 

167 

168 reply = Message() 

169 reply.add_byte(c_MSG_KEX_HYBRID_REPLY) 

170 reply.add_string(K_S) 

171 reply.add_string(s_reply) 

172 reply.add_string(sig) 

173 self.transport._send_message(reply) 

174 self.transport._activate_outbound() 

175 

176 # ---- client side ------------------------------------------------- 

177 

178 def _parse_hybrid_reply(self, m): 

179 K_S = m.get_string() 

180 s_reply = m.get_string() 

181 sig = m.get_binary() 

182 

183 if len(s_reply) != self._S_REPLY_BYTES: 

184 raise SSHException( 

185 "Invalid S_REPLY length for {}: got {}, expected {}".format( 

186 self.name, len(s_reply), self._S_REPLY_BYTES 

187 ) 

188 ) 

189 s_ct2 = s_reply[: self._MLKEM_CIPHERTEXT_BYTES] 

190 s_pk1 = s_reply[self._MLKEM_CIPHERTEXT_BYTES :] 

191 

192 k_pq = self.mlkem_key.decapsulate(s_ct2) 

193 k_cl = self._x25519_exchange(s_pk1) 

194 K_bytes = self.hash_algo(k_pq + k_cl).digest() 

195 

196 c_init = ( 

197 self.mlkem_key.public_key().public_bytes_raw() 

198 + self.x25519_key.public_key().public_bytes_raw() 

199 ) 

200 

201 hm = Message() 

202 hm.add( 

203 self.transport.local_version, 

204 self.transport.remote_version, 

205 self.transport.local_kex_init, 

206 self.transport.remote_kex_init, 

207 ) 

208 hm.add_string(K_S) 

209 hm.add_string(c_init) 

210 hm.add_string(s_reply) 

211 hm.add_string(K_bytes) 

212 H = self.hash_algo(hm.asbytes()).digest() 

213 

214 self.transport._set_K_H(K_bytes, H) 

215 self.transport._verify_key(K_S, sig) 

216 self.transport._activate_outbound()