Handle KeepKey bad arguments

hwilib/devices/keepkey.py

@@ -1,7 +1,7 @@
 # KeepKey interaction script
 
 from ..hwwclient import HardwareWalletClient
-from ..errors import BadArgumentError, DeviceAlreadyUnlockedError, UnavailableActionError, DeviceNotReadyError
+from ..errors import BadArgumentError, DeviceAlreadyUnlockedError, DeviceConnectionError, UnavailableActionError, DeviceNotReadyError
 from keepkeylib.transport_hid import HidTransport
 from keepkeylib.transport_udp import UDPTransport
 from keepkeylib.client import BaseClient, DebugWireMixin, DebugLinkMixin, ProtocolMixin, TextUIMixin
@@ -120,6 +120,16 @@ class KeepKey(NoInitMixin, TextUIMixin, BaseClient):
 class KeepKeyDebug(NoInitMixin, DebugLinkMixin, DebugWireMixin, BaseClient):
     pass
 
+def keepkey_exception(f):
+    def func(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except ValueError as e:
+            raise BadArgumentError(str(e))
+        except OSError as e:
+            raise DeviceConnectionError(str(e))
+    return func
+
 # This class extends the HardwareWalletClient for Digital Bitbox specific things
 class KeepkeyClient(HardwareWalletClient):
 
@@ -156,6 +166,7 @@ class KeepkeyClient(HardwareWalletClient):
 
     # Must return a dict with the xpub
     # Retrieves the public key at the specified BIP 32 derivation path
+    @keepkey_exception
     def get_pubkey_at_path(self, path):
         self._check_unlocked()
         path = path.replace('h', '\'')
@@ -169,6 +180,7 @@ class KeepkeyClient(HardwareWalletClient):
 
     # Must return a hex string with the signed transaction
     # The tx must be in the combined unsigned transaction format
+    @keepkey_exception
     def sign_tx(self, tx):
         self._check_unlocked()
 
@@ -321,6 +333,7 @@ class KeepkeyClient(HardwareWalletClient):
 
     # Must return a base64 encoded string with the signed message
     # The message can be any string
+    @keepkey_exception
     def sign_message(self, message, keypath):
         self._check_unlocked()
         keypath = keypath.replace('h', '\'')
@@ -330,6 +343,7 @@ class KeepkeyClient(HardwareWalletClient):
         return {'signature': base64.b64encode(result.signature).decode('utf-8')}
 
     # Display address of specified type on the device. Only supports single-key based addresses.
+    @keepkey_exception
     def display_address(self, keypath, p2sh_p2wpkh, bech32):
         self._check_unlocked()
         keypath = keypath.replace('h', '\'')
@@ -344,6 +358,7 @@ class KeepkeyClient(HardwareWalletClient):
         return {'address': address}
 
     # Setup a new device
+    @keepkey_exception
     def setup_device(self, label='', passphrase=''):
         if self.client.features.initialized:
             raise DeviceAlreadyInitError('Device is already initialized. Use wipe first and try again')
@@ -351,17 +366,20 @@ class KeepkeyClient(HardwareWalletClient):
         return {'success': True}
 
     # Wipe this device
+    @keepkey_exception
     def wipe_device(self):
         self._check_unlocked()
         self.client.wipe_device()
         return {'success': True}
 
     # Restore device from mnemonic or xprv
+    @keepkey_exception
     def restore_device(self, label=''):
         self.client.recovery_device(False, 24, bool(self.password), True, label, 'english')
         return {'success': True}
 
     # Begin backup process
+    @keepkey_exception
     def backup_device(self, label='', passphrase=''):
         raise UnavailableActionError('The Keepkey does not support creating a backup via software')
 
@@ -370,6 +388,7 @@ class KeepkeyClient(HardwareWalletClient):
         self.client.close()
 
     # Prompt for a pin on device
+    @keepkey_exception
     def prompt_pin(self):
         self.client.init_device()
         if not self.client.features.pin_protection:
@@ -382,6 +401,7 @@ class KeepkeyClient(HardwareWalletClient):
         return {'success': True}
 
     # Send the pin
+    @keepkey_exception
     def send_pin(self, pin):
         if not pin.isdigit():
             raise BadArgumentError("Non-numeric PIN provided")