Specter-DIY: import xpubs without the 3s timeout, one confirmation per key set, wrong-network handling

src/cryptoadvance/specter/devices/hwi/specter_diy.py

@@ -23,6 +23,86 @@
 py_enumerate = enumerate
 logger = logging.getLogger(__name__)
 
+# Prefix of the machine-parseable error the firmware sends back (as
+# "error: <this><network>") when it refuses a plain `xpub <path>` request
+# because the path's implied coin type doesn't match the network currently
+# active on the device. Keep in sync with specter-diy's
+# apps/xpubs/xpubs.py.
+_NETWORK_MISMATCH_PREFIX = "network mismatch: device is on "
+
+# Human-readable labels for the network identifiers the firmware reports,
+# matching the names shown on the device's own screen (embit's
+# NETWORKS[...]["name"]).
+_NETWORK_DISPLAY_NAMES = {
+    "main": "Mainnet",
+    "test": "Testnet",
+    "regtest": "Regtest",
+    "signet": "Signet",
+    "liquidv1": "Liquid",
+    "liquidtestnet": "Liquid Testnet",
+    "elementsregtest": "Liquid Regtest",
+}
+
+# Standard BIP44-style purposes that carry an implied coin type (and hence
+# an implied network). Keep in sync with specter-diy's apps/xpubs/scope.py.
+_STANDARD_PURPOSES = {44, 48, 49, 84, 86, 87}
+
+# By-name switch-to hint for a requested path's coin type, mirroring
+# specter-diy's network_hint_for_path(). Coin type 1' is shared by
+# testnet/signet/regtest/liquidtestnet/elementsregtest, so it can only be
+# named as the group.
+_COIN_TYPE_NETWORK_HINTS = {
+    0: "Mainnet",
+    1: "Testnet, Signet or Regtest",
+    1776: "Liquid",
+}
+
+
+def _requested_network_label(bip32_path):
+    """
+    The network a standard-purpose derivation path implies, by name, from
+    its coin type - or None if the path isn't standard-purpose (and so
+    carries no network meaning). Matches the firmware rule: the purpose
+    and coin type must both be *hardened* for the path to carry network
+    semantics.
+    """
+    parts = [p for p in str(bip32_path).strip().split("/") if p and p.lower() != "m"]
+    if len(parts) < 2:
+        return None
+
+    def _hardened_index(component):
+        component = component.strip()
+        if component[-1:] not in ("h", "H", "'"):
+            raise ValueError("not hardened")
+        return int(component[:-1])
+
+    try:
+        purpose = _hardened_index(parts[0])
+        coin_type = _hardened_index(parts[1])
+    except ValueError:
+        return None
+    if purpose not in _STANDARD_PURPOSES:
+        return None
+    return _COIN_TYPE_NETWORK_HINTS.get(coin_type)
+
+
+class SpecterDIYNetworkMismatchError(BadArgumentError):
+    """
+    The device refused an xpub request because the requested derivation's
+    coin type doesn't match the network currently active on the device
+    (e.g. asking for a mainnet-style key while the device is set to
+    Testnet). Raised instead of a generic BadArgumentError so callers can
+    tell the user exactly which network to switch the device to, the same
+    thing the device's own screen already told the person holding it.
+
+    :ivar device_network: the raw network identifier the device reported
+        (e.g. "test", "regtest") - not a display name.
+    """
+
+    def __init__(self, device_network: str, message: str) -> None:
+        super().__init__(message)
+        self.device_network = device_network
+
 
 class SpecterClient(HardwareWalletClient):
     """Create a client for a HID device that has already been opened.
@@ -31,7 +111,10 @@ class SpecterClient(HardwareWalletClient):
     that hardware wallet subclasses should implement.
     """
 
-    # timeout large enough to handle xpub derivations
+    # bounded timeout for non-interactive operations (e.g. fingerprint,
+    # used for device discovery/identification). xpub now requires an
+    # on-device user confirmation, so get_pubkey_at_path() does not use
+    # this - it waits indefinitely, like sign_tx() already does.
     TIMEOUT = 3
 
     def __init__(self, path: str, password: str = "", expert: bool = False) -> None:
@@ -87,8 +170,12 @@ def get_pubkey_at_path(self, bip32_path: str) -> ExtendedKey:
         :param bip32_path: The BIP 32 derivation path
         :return: The extended public key
         """
-        # this should be fast
-        xpub = self.query("xpub %s" % bip32_path, timeout=self.TIMEOUT)
+        # xpub now requires on-device user confirmation and may take
+        # arbitrarily long - wait indefinitely, same as sign_tx()
+        try:
+            xpub = self.query("xpub %s" % bip32_path)
+        except BadArgumentError as e:
+            raise self._network_mismatch_error(e, bip32_path) from e
         hd = ExtendedKey.deserialize(xpub)
         # Specter returns xpub with a prefix
         # for a network currently selected on the device
@@ -97,6 +184,95 @@ def get_pubkey_at_path(self, bip32_path: str) -> ExtendedKey:
         )
         return hd
 
+    def _network_mismatch_error(
+        self, error: BadArgumentError, bip32_path: str
+    ) -> Exception:
+        """
+        Recognize the firmware's "network mismatch: device is on <network>"
+        response and turn it into a SpecterDIYNetworkMismatchError with a
+        message naming both networks involved. Any other BadArgumentError
+        is returned unchanged.
+        """
+        msg = str(error)
+        if not msg.startswith(_NETWORK_MISMATCH_PREFIX):
+            return error
+        device_network = msg[len(_NETWORK_MISMATCH_PREFIX) :].strip()
+        # The device controls this string; never put it in the message
+        # verbatim (it's rendered, not escaped, in the UI). Only ever show
+        # a known, whitelisted label.
+        device_label = (
+            _NETWORK_DISPLAY_NAMES.get(device_network) or "an unknown network"
+        )
+        # Name the network to switch to from the requested path's own coin
+        # type - unambiguous for mainnet (0') and Liquid (1776'), the
+        # shared group for 1'. If the coin type has no name we recognize,
+        # don't guess from the client's chain (that can produce a "switch
+        # to Mainnet" while the device already is on Mainnet) - say what
+        # the device's own screen says instead.
+        requested_label = _requested_network_label(bip32_path)
+        if requested_label is not None:
+            message = (
+                "This Specter-DIY is currently set to %s, so it can't share a "
+                "%s key from here. Switch the device to %s and try again."
+                % (device_label, requested_label, requested_label)
+            )
+        else:
+            message = (
+                "This Specter-DIY is currently set to %s, which doesn't match "
+                "the requested derivation. Switch the device to the network "
+                "that derivation is for and try again." % device_label
+            )
+        return SpecterDIYNetworkMismatchError(device_network, message)
+
+    def begin_xpub_authorization(self, paths) -> bool:
+        """
+        Authorize a whole set of xpub derivation paths with a *single*
+        on-device confirmation instead of one Confirm per
+        ``get_pubkey_at_path()`` call.
+
+        Uses the Specter-DIY ``xpubauth begin <scope>`` protocol: after one
+        approval, ``xpub <derivation>`` requests that fall inside the
+        approved scope are answered without a further prompt (each
+        authorized path exactly once). Call :meth:`end_xpub_authorization`
+        when done.
+
+        Returns ``True`` when the device authorized the scope, ``False``
+        when the firmware is too old to know ``xpubauth`` or rejected the
+        scope. The device only authorizes a scope for the network it is
+        currently on, so a ``False`` for one network's standard paths is
+        also the signal that the device is on a *different* network -
+        callers can use that instead of blindly asking for both.
+
+        Raises ``ActionCanceledError`` if the user declined on the device.
+        All ``paths`` must belong to the same network.
+        """
+        if not paths:
+            return False
+        try:
+            # no timeout: waits for the user to Confirm/Cancel on the device
+            self.query("xpubauth begin %s" % ";".join(paths))
+            return True
+        except UnavailableActionError:
+            logger.info(
+                "Specter-DIY firmware without 'xpubauth' support - "
+                "falling back to per-xpub confirmation"
+            )
+            return False
+        except BadArgumentError as e:
+            logger.info(
+                "Specter-DIY did not authorize the xpubauth scope (%s) - "
+                "falling back to per-xpub confirmation",
+                e,
+            )
+            return False
+
+    def end_xpub_authorization(self) -> None:
+        """Discard any active ``xpubauth`` authorization (best effort)."""
+        try:
+            self.query("xpubauth end")
+        except Exception as e:
+            logger.debug("xpubauth end failed: %s", e)
+
     def sign_b64psbt(self, psbt: str) -> str:
         # works with both PSBT and PSET
         return self.query("sign %s" % psbt)
@@ -308,6 +484,11 @@ def read_until(self, eol, timeout=None):
                 raw = self.ser.read(1)
                 res += raw
             except Exception as e:
+                raw = b""
+            if not raw:
+                # non-blocking read returned nothing (e.g. waiting for the
+                # user to confirm on the device) - yield the CPU instead
+                # of spinning on read(1)
                 time.sleep(0.01)
             if timeout is not None and time.time() > t0 + timeout:
                 self.ser.close()