bitbox01: implement update_firmware

Author: achow101

hwilib/devices/digitalbitbox.py

@@ -1,6 +1,7 @@
 # Digital Bitbox interaction script
 
 import hid
+import io
 import struct
 import json
 import base64
@@ -15,7 +16,7 @@
 import time
 
 from ..hwwclient import HardwareWalletClient
-from ..errors import ActionCanceledError, BadArgumentError, DeviceFailureError, DeviceAlreadyInitError, DEVICE_NOT_INITIALIZED, DeviceNotReadyError, NoPasswordError, UnavailableActionError, common_err_msgs, handle_errors
+from ..errors import ActionCanceledError, BAD_ARGUMENT, BadArgumentError, DeviceFailureError, DeviceAlreadyInitError, DEVICE_CONN_ERROR, DEVICE_NOT_INITIALIZED, DeviceNotReadyError, NoPasswordError, UnavailableActionError, common_err_msgs, handle_errors
 from ..serializations import CTransaction, ExtendedKey, hash256, ser_sig_der, ser_sig_compact, ser_compact_size
 from ..base58 import get_xpub_fingerprint, xpub_main_2_test, get_xpub_fingerprint_hex
 
@@ -188,6 +189,9 @@ def close(self):
     def get_serial_number_string(self):
         return 'dbb_fw:v5.0.0'
 
+    def get_product_string(self):
+        return 'Digital Bitbox firmware'
+
 def send_frame(data, device):
     data = bytearray(data)
     data_len = len(data)
@@ -293,6 +297,45 @@ def stretch_backup_key(password):
 def format_backup_filename(name):
     return '{}-{}.pdf'.format(name, time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime()))
 
+# ----------------------------------------------------------------------------------
+# Bootloader io
+#
+
+def sendBoot(msg, dev):
+    msg = bytearray(msg) + b'\0' * (boot_buf_size_send - len(msg))
+    serial_number = dev.get_serial_number_string()
+    if 'v1.' in serial_number or 'v2.' in serial_number:
+        dev.write(b'\0' + msg)
+    else:
+        # Split `msg` into 64-byte packets
+        n = 0
+        while n < len(msg):
+            dev.write(b'\0' + msg[n:n + usb_report_size])
+            n = n + usb_report_size
+
+def sendPlainBoot(msg, dev):
+    if type(msg) == str:
+        msg = msg.encode()
+    sendBoot(msg, dev)
+    reply = []
+    while len(reply) < boot_buf_size_reply:
+        reply = reply + dev.read(boot_buf_size_reply)
+
+    reply = bytearray(reply).rstrip(b' \t\r\n\0')
+    reply = ''.join(chr(e) for e in reply)
+    return reply
+
+def sendChunk(chunknum, data, dev):
+    b = bytearray(b"\x77\x00")
+    b[1] = chunknum % 0xFF
+    b.extend(data)
+    sendBoot(b, dev)
+    reply = []
+    while len(reply) < boot_buf_size_reply:
+        reply = reply + dev.read(boot_buf_size_reply)
+    reply = bytearray(reply).rstrip(b' \t\r\n\0')
+    reply = ''.join(chr(e) for e in reply)
+
 # This class extends the HardwareWalletClient for Digital Bitbox specific things
 class DigitalbitboxClient(HardwareWalletClient):
 
@@ -310,6 +353,21 @@ def __init__(self, path, password, expert=False):
             self.device.open_path(path.encode())
         self.password = password
 
+        # Always lock the bootloader
+        if self.device.get_product_string() != 'bootloader':
+            reply = send_encrypt('{"device":"info"}', self.password, self.device)
+            if 'error' not in reply:
+                if not reply['device']['bootlock']:
+                    reply = send_encrypt('{"bootloader":"lock"}', self.password, self.device)
+                    if 'error' in reply:
+                        raise DBBError(reply)
+            else:
+                # Check it isn't initialized
+                if reply['error']['code'] == 101 or reply['error']['code'] == '101':
+                    pass
+                else:
+                    raise DBBError(reply)
+
     # Must return a dict with the xpub
     # Retrieves the public key at the specified BIP 32 derivation path
     @digitalbitbox_exception
@@ -584,8 +642,53 @@ def send_pin(self, pin):
         raise UnavailableActionError('The Digital Bitbox does not need a PIN sent from the host')
 
     # Verify firmware file then load it onto device
+    @digitalbitbox_exception
     def update_firmware(self, file):
-        raise NotImplementedError('The Digital Bitbox does not implement this method yet')
+        if self.device.get_product_string() != 'bootloader':
+            print('Device is not in bootloader mode. Unlocking bootloader, replugging will be required', file=sys.stderr)
+            print("Touch the device for 3 seconds to unlock bootloaderr. Touch briefly to cancel", file=sys.stderr)
+            reply = send_encrypt('{"bootloader":"unlock"}', self.password, self.device)
+            if 'error' in reply:
+                raise DBBError(reply)
+            return {'error': 'Digital Bitbox needs to be in bootloader mode. Unplug and replug the device and briefly touch the button within 3 seconds. Then try this command again', 'code': DEVICE_CONN_ERROR}
+
+        with open(file, "rb") as f:
+            data = bytearray()
+            while True:
+                d = f.read(chunksize)
+                if len(d) == 0:
+                    break
+                data = data + bytearray(d)
+        data = data + b'\xFF' * (applen - len(data))
+        firmware = data[448:]
+        sig = data[:448]
+        print('Hashed firmware (without signatures)', binascii.hexlify(hash256((firmware))), file=sys.stderr)
+
+        sendPlainBoot("b", self.device) # blink led
+        sendPlainBoot("v", self.device) # bootloader version
+        sendPlainBoot("e", self.device) # erase existing firmware (required)
+
+        # Send firmware
+        f = io.BytesIO(firmware)
+        cnt = 0
+        while True:
+            chunk = f.read(chunksize)
+            if len(chunk) == 0:
+                break
+            sendChunk(cnt, chunk, self.device)
+            cnt += 1
+
+        # upload sigs and verify new firmware
+        load_result = sendPlainBoot("s" + "0" + binascii.hexlify(sig).decode(), self.device)
+        if load_result[1] == 'V':
+            latest_version, = struct.unpack('>I', binascii.unhexlify(load_result[2 + 64:][:8]))
+            app_version, = struct.unpack('>I', binascii.unhexlify(load_result[2 + 64 + 8:][:8]))
+            return {'error': 'firmware downgrade not allowed. Got version %d, but must be equal or higher to %d' % (app_version, latest_version), 'code': BAD_ARGUMENT}
+        elif load_result[1] != '0':
+            return {'error': 'invalid firmware signature', 'code': BAD_ARGUMENT}
+
+        print('Please unplug and replug your device. The bootloader will be locked next time you use HWI with it.', file=sys.stderr)
+        return {'success': True}
 
 def enumerate(password=''):
     results = []