Handle Coldcard canceled actions and bad arguments

hwilib/devices/coldcard.py

@@ -1,9 +1,9 @@
 # Coldcard interaction script
 
 from ..hwwclient import HardwareWalletClient
-from ..errors import UnavailableActionError, DeviceFailureError
+from ..errors import ActionCanceledError, BadArgumentError, DeviceBusyError, UnavailableActionError, DeviceFailureError
 from ckcc.client import ColdcardDevice, COINKITE_VID, CKCC_PID
-from ckcc.protocol import CCProtocolPacker, CCProtoError
+from ckcc.protocol import CCProtocolPacker, CCBusyError, CCProtoError, CCUserRefused
 from ckcc.constants import MAX_BLK_LEN, AF_P2WPKH, AF_CLASSIC, AF_P2WPKH_P2SH
 from ..base58 import xpub_main_2_test, get_xpub_fingerprint_hex
 from hashlib import sha256
@@ -18,6 +18,18 @@ import time
 CC_SIMULATOR_SOCK = '/tmp/ckcc-simulator.sock'
 # Using the simulator: https://github.com/Coldcard/firmware/blob/master/unix/README.md
 
+def coldcard_exception(f):
+    def func(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except CCProtoError as e:
+            raise BadArgumentError(str(e))
+        except CCUserRefused as e:
+            raise ActionCanceledError('{} canceled'.format(f.__name__))
+        except CCBusyError as e:
+            raise DeviceBusyError(str(e))
+    return func
+
 # This class extends the HardwareWalletClient for ColdCard specific things
 class ColdcardClient(HardwareWalletClient):
 
@@ -33,6 +45,7 @@ class ColdcardClient(HardwareWalletClient):
 
     # Must return a dict with the xpub
     # Retrieves the public key at the specified BIP 32 derivation path
+    @coldcard_exception
     def get_pubkey_at_path(self, path):
         self.device.check_mitm()
         path = path.replace('h', '\'')
@@ -45,6 +58,7 @@ class ColdcardClient(HardwareWalletClient):
 
     # Must return a hex string with the signed transaction
     # The tx must be in the combined unsigned transaction format
+    @coldcard_exception
     def sign_tx(self, tx):
         self.device.check_mitm()
 
@@ -98,6 +112,7 @@ class ColdcardClient(HardwareWalletClient):
 
     # Must return a base64 encoded string with the signed message
     # The message can be any string. keypath is the bip 32 derivation path for the key to sign with
+    @coldcard_exception
     def sign_message(self, message, keypath):
         self.device.check_mitm()
         keypath = keypath.replace('h', '\'')
@@ -128,6 +143,7 @@ class ColdcardClient(HardwareWalletClient):
         return {"signature": sig}
 
     # Display address of specified type on the device. Only supports single-key based addresses.
+    @coldcard_exception
     def display_address(self, keypath, p2sh_p2wpkh, bech32):
         self.device.check_mitm()
         keypath = keypath.replace('h', '\'')
@@ -157,6 +173,7 @@ class ColdcardClient(HardwareWalletClient):
         raise UnavailableActionError('The Coldcard does not support restoring via software')
 
     # Begin backup process
+    @coldcard_exception
     def backup_device(self, label='', passphrase=''):
         self.device.check_mitm()