Coverage for tests/ledger/test_protocol_v1.py: 100%

1# The MIT License (MIT) 

2# 

3# Copyright (c) 2021 RSK Labs Ltd 

4# 

5# Permission is hereby granted, free of charge, to any person obtaining a copy of 

6# this software and associated documentation files (the "Software"), to deal in 

7# the Software without restriction, including without limitation the rights to 

8# use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies 

9# of the Software, and to permit persons to whom the Software is furnished to do 

10# so, subject to the following conditions: 

11# 

12# The above copyright notice and this permission notice shall be included in all 

13# copies or substantial portions of the Software. 

14# 

15# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

16# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

17# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

18# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

19# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

20# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

21# SOFTWARE. 

22 

23from unittest import TestCase 

24from unittest.mock import Mock, call, patch 

25from parameterized import parameterized 

26from comm.protocol import HSM2ProtocolError 

27from ledger.protocol_v1 import HSM1ProtocolLedger 

28from ledger.hsm2dongle import ( 

29 HSM2Dongle, 

30 HSM2FirmwareVersion, 

31 HSM2DongleError, 

32 HSM2DongleErrorResult, 

33 HSM2DongleTimeoutError, 

34 HSM2DongleCommError, 

35) 

36 

37import logging 

38 

39logging.disable(logging.CRITICAL) 

40 

41 

42class TestHSM1ProtocolLedger(TestCase): 

43 def setUp(self): 

44 self.pin = Mock() 

45 self.dongle = Mock() 

46 self.dongle.connect = Mock() 

47 self.dongle.disconnect = Mock() 

48 self.dongle.is_onboarded = Mock(return_value=True) 

49 self.dongle.get_current_mode = Mock(return_value=HSM2Dongle.MODE.SIGNER) 

50 self.dongle.get_version = Mock(return_value=HSM2FirmwareVersion(5, 5, 1)) 

51 self.dongle.get_signer_parameters = Mock(return_value=Mock( 

52 min_required_difficulty=123)) 

53 self.protocol = HSM1ProtocolLedger(self.pin, self.dongle) 

54 self.protocol.initialize_device() 

55 

56 @patch("comm.protocol.BIP32Path") 

57 def test_get_pubkey_ok(self, BIP32PathMock): 

58 BIP32PathMock.return_value = "the-key-id" 

59 self.dongle.get_public_key.return_value = "this-is-the-public-key" 

60 

61 self.assertEqual( 

62 { 

63 "errorcode": 0, 

64 "pubKey": "this-is-the-public-key" 

65 }, 

66 self.protocol.handle_request({ 

67 "version": 1, 

68 "command": "getPubKey", 

69 "keyId": "m/44'/1'/2'/3/4" 

70 }), 

71 ) 

72 self.assertEqual([call("the-key-id")], self.dongle.get_public_key.call_args_list) 

73 self.assertFalse(self.dongle.disconnect.called) 

74 

75 @patch("comm.protocol.BIP32Path") 

76 def test_get_pubkey_error(self, BIP32PathMock): 

77 BIP32PathMock.return_value = "the-key-id" 

78 self.dongle.get_public_key.side_effect = HSM2DongleErrorResult() 

79 

80 self.assertEqual( 

81 {"errorcode": -2}, 

82 self.protocol.handle_request({ 

83 "version": 1, 

84 "command": "getPubKey", 

85 "keyId": "m/44'/1'/2'/3/4" 

86 }), 

87 ) 

88 self.assertEqual([call("the-key-id")], self.dongle.get_public_key.call_args_list) 

89 self.assertFalse(self.dongle.disconnect.called) 

90 

91 @patch("comm.protocol.BIP32Path") 

92 def test_get_pubkey_timeout(self, BIP32PathMock): 

93 BIP32PathMock.return_value = "the-key-id" 

94 self.dongle.get_public_key.side_effect = HSM2DongleTimeoutError() 

95 

96 self.assertEqual( 

97 {"errorcode": -2}, 

98 self.protocol.handle_request({ 

99 "version": 1, 

100 "command": "getPubKey", 

101 "keyId": "m/44'/1'/2'/3/4" 

102 }), 

103 ) 

104 self.assertEqual([call("the-key-id")], self.dongle.get_public_key.call_args_list) 

105 self.assertFalse(self.dongle.disconnect.called) 

106 

107 @patch("comm.protocol.BIP32Path") 

108 def test_get_pubkey_commerror_reconnection(self, BIP32PathMock): 

109 BIP32PathMock.return_value = "the-key-id" 

110 self.dongle.get_public_key.side_effect = HSM2DongleCommError() 

111 

112 self.assertEqual( 

113 {"errorcode": -2}, 

114 self.protocol.handle_request({ 

115 "version": 1, 

116 "command": "getPubKey", 

117 "keyId": "m/44'/1'/2'/3/4" 

118 }), 

119 ) 

120 self.assertEqual([call("the-key-id")], self.dongle.get_public_key.call_args_list) 

121 self.assertFalse(self.dongle.disconnect.called) 

122 

123 # Reconnection logic 

124 self.dongle.get_public_key.side_effect = None 

125 self.dongle.get_public_key.return_value = "this-is-the-public-key" 

126 

127 self.assertEqual( 

128 { 

129 "errorcode": 0, 

130 "pubKey": "this-is-the-public-key" 

131 }, 

132 self.protocol.handle_request({ 

133 "version": 1, 

134 "command": "getPubKey", 

135 "keyId": "m/44'/1'/2'/3/4" 

136 }), 

137 ) 

138 

139 self._assert_reconnected() 

140 

141 @patch("comm.protocol.BIP32Path") 

142 def test_get_pubkey_unexpected_error(self, BIP32PathMock): 

143 BIP32PathMock.return_value = "the-key-id" 

144 self.dongle.get_public_key.side_effect = HSM2DongleError() 

145 

146 with self.assertRaises(HSM2ProtocolError): 

147 self.protocol.handle_request({ 

148 "version": 1, 

149 "command": "getPubKey", 

150 "keyId": "m/44'/1'/2'/3/4" 

151 }) 

152 

153 self.assertEqual([call("the-key-id")], self.dongle.get_public_key.call_args_list) 

154 self.assertFalse(self.dongle.disconnect.called) 

155 

156 @patch("comm.protocol.BIP32Path") 

157 def test_sign_ok(self, BIP32PathMock): 

158 BIP32PathMock.return_value = "the-key-id" 

159 signature = Mock(r="this-is-r", s="this-is-s") 

160 self.dongle.sign_unauthorized.return_value = (True, signature) 

161 

162 self.assertEqual( 

163 { 

164 "errorcode": 0, 

165 "signature": { 

166 "r": "this-is-r", 

167 "s": "this-is-s" 

168 } 

169 }, 

170 self.protocol.handle_request({ 

171 "version": 1, 

172 "command": "sign", 

173 "keyId": "m/44'/1'/2'/3/4", 

174 "message": "aa"*32, 

175 }), 

176 ) 

177 

178 self.assertEqual( 

179 [call(key_id="the-key-id", hash="aa"*32)], 

180 self.dongle.sign_unauthorized.call_args_list, 

181 ) 

182 self.assertFalse(self.dongle.disconnect.called) 

183 

184 @parameterized.expand([ 

185 ("path", -1, -2), 

186 ("hash", -5, -2), 

187 ("unexpected", -10, -2), 

188 ("unknown", -100, -2), 

189 ]) 

190 @patch("comm.protocol.BIP32Path") 

191 def test_sign_error(self, _, dongle_error_code, protocol_error_code, BIP32PathMock): 

192 BIP32PathMock.return_value = "the-key-id" 

193 self.dongle.sign_unauthorized.return_value = (False, dongle_error_code) 

194 

195 self.assertEqual( 

196 {"errorcode": protocol_error_code}, 

197 self.protocol.handle_request({ 

198 "version": 1, 

199 "command": "sign", 

200 "keyId": "m/44'/1'/2'/3/4", 

201 "message": "aa"*32, 

202 }), 

203 ) 

204 

205 self.assertEqual( 

206 [call(key_id="the-key-id", hash="aa"*32)], 

207 self.dongle.sign_unauthorized.call_args_list, 

208 ) 

209 self.assertFalse(self.dongle.disconnect.called) 

210 

211 @patch("comm.protocol.BIP32Path") 

212 def test_sign_timeout(self, BIP32PathMock): 

213 BIP32PathMock.return_value = "the-key-id" 

214 self.dongle.sign_unauthorized.side_effect = HSM2DongleTimeoutError() 

215 

216 self.assertEqual( 

217 {"errorcode": -2}, 

218 self.protocol.handle_request({ 

219 "version": 1, 

220 "command": "sign", 

221 "keyId": "m/44'/1'/2'/3/4", 

222 "message": "aa"*32, 

223 }), 

224 ) 

225 

226 self.assertEqual( 

227 [call(key_id="the-key-id", hash="aa"*32)], 

228 self.dongle.sign_unauthorized.call_args_list, 

229 ) 

230 self.assertFalse(self.dongle.disconnect.called) 

231 

232 @patch("comm.protocol.BIP32Path") 

233 def test_sign_commerror_reconnection(self, BIP32PathMock): 

234 BIP32PathMock.return_value = "the-key-id" 

235 self.dongle.sign_unauthorized.side_effect = HSM2DongleCommError() 

236 

237 self.assertEqual( 

238 {"errorcode": -2}, 

239 self.protocol.handle_request({ 

240 "version": 1, 

241 "command": "sign", 

242 "keyId": "m/44'/1'/2'/3/4", 

243 "message": "aa"*32, 

244 }), 

245 ) 

246 

247 self.assertEqual( 

248 [call(key_id="the-key-id", hash="aa"*32)], 

249 self.dongle.sign_unauthorized.call_args_list, 

250 ) 

251 self.assertFalse(self.dongle.disconnect.called) 

252 

253 # Reconnection logic 

254 self.dongle.sign_unauthorized.side_effect = None 

255 signature = Mock(r="this-is-r", s="this-is-s") 

256 self.dongle.sign_unauthorized.return_value = (True, signature) 

257 

258 self.assertEqual( 

259 { 

260 "errorcode": 0, 

261 "signature": { 

262 "r": "this-is-r", 

263 "s": "this-is-s" 

264 } 

265 }, 

266 self.protocol.handle_request({ 

267 "version": 1, 

268 "command": "sign", 

269 "keyId": "m/44'/1'/2'/3/4", 

270 "message": "aa"*32, 

271 }), 

272 ) 

273 

274 self._assert_reconnected() 

275 

276 @patch("comm.protocol.BIP32Path") 

277 def test_sign_exception(self, BIP32PathMock): 

278 BIP32PathMock.return_value = "the-key-id" 

279 self.dongle.sign_unauthorized.side_effect = HSM2DongleError() 

280 

281 with self.assertRaises(HSM2ProtocolError): 

282 self.protocol.handle_request({ 

283 "version": 1, 

284 "command": "sign", 

285 "keyId": "m/44'/1'/2'/3/4", 

286 "message": "aa"*32, 

287 }) 

288 

289 self.assertEqual( 

290 [call(key_id="the-key-id", hash="aa"*32)], 

291 self.dongle.sign_unauthorized.call_args_list, 

292 ) 

293 self.assertFalse(self.dongle.disconnect.called) 

294 

295 @patch("comm.protocol.BIP32Path") 

296 def test_sign_message_invalid(self, BIP32PathMock): 

297 BIP32PathMock.return_value = "the-key-id" 

298 

299 self.assertEqual( 

300 {"errorcode": -2}, 

301 self.protocol.handle_request({ 

302 "version": 1, 

303 "command": "sign", 

304 "keyId": "m/44'/1'/2'/3/4", 

305 "message": "not-a-hexadecimal-string", 

306 }), 

307 ) 

308 

309 self.assertFalse(self.dongle.sign_unauthorized.called) 

310 self.assertFalse(self.dongle.disconnect.called) 

311 

312 def _assert_reconnected(self): 

313 self.assertTrue(self.dongle.disconnect.called) 

314 self.assertEqual(2, self.dongle.connect.call_count)