coldcard: Implement update_firmware

Author: achow101

hwilib/devices/coldcard.py

@@ -6,6 +6,8 @@
 from .ckcc.client import ColdcardDevice, COINKITE_VID, CKCC_PID
 from .ckcc.protocol import CCProtocolPacker, CCBusyError, CCProtoError, CCUserRefused
 from .ckcc.constants import MAX_BLK_LEN, AF_P2WPKH, AF_CLASSIC, AF_P2WPKH_P2SH
+from .ckcc.sigheader import FW_HEADER_SIZE, FW_HEADER_OFFSET, FW_HEADER_MAGIC
+from .ckcc.utils import dfu_parse
 from ..base58 import get_xpub_fingerprint, xpub_main_2_test
 from ..serializations import ExtendedKey, PSBT
 from hashlib import sha256
@@ -244,7 +246,67 @@ def send_pin(self, pin):
 
     # Verify firmware file then load it onto device
     def update_firmware(self, file):
-        raise NotImplementedError('The Coldcard does not implement this method yet')
+        with open(file, 'rb') as fd:
+            # learn size (portable way)
+            offset = 0
+            sz = fd.seek(0, 2)
+            fd.seek(0)
+
+            # Unwrap DFU contents, if needed. Also handles raw binary file.
+            try:
+                if fd.read(5) == b'DfuSe':
+                    # expecting a DFU-wrapped file.
+                    fd.seek(0)
+                    offset, sz, *_ = dfu_parse(fd)
+                else:
+                    # assume raw binary
+                    pass
+
+                assert sz % 256 == 0, "un-aligned size: %s" % sz
+                fd.seek(offset + FW_HEADER_OFFSET)
+                hdr = fd.read(FW_HEADER_SIZE)
+
+                magic = struct.unpack_from("<I", hdr)[0]
+            except Exception:
+                magic = None
+
+            if magic != FW_HEADER_MAGIC:
+                raise BadArgumentError('{} has an invalid magic header for a firmware file.'.format(file))
+
+            fd.seek(offset)
+
+            left = sz
+            chk = sha256()
+            for pos in range(0, sz, MAX_BLK_LEN):
+                here = fd.read(min(MAX_BLK_LEN, left))
+                if not here:
+                    break
+                left -= len(here)
+                result = self.device.send_recv(CCProtocolPacker.upload(pos, sz, here))
+                assert result == pos, "Got back: %r" % result
+                chk.update(here)
+
+        # do a verify
+        expect = chk.digest()
+        result = self.device.send_recv(CCProtocolPacker.sha256())
+        assert len(result) == 32
+        if result != expect:
+            raise DeviceFailureError("Wrong checksum:\nexpect: %s\n   got: %s" % (b2a_hex(expect).decode('ascii'), b2a_hex(result).decode('ascii')))
+
+        # AFTER fully uploaded and verified, write a copy of the signature header
+        # onto the end of flash. Bootrom uses this to check entire file uploaded.
+        result = self.device.send_recv(CCProtocolPacker.upload(sz, sz + FW_HEADER_SIZE, hdr))
+        assert result == sz, "failed to write trailer"
+
+        # check also SHA after that!
+        chk.update(hdr)
+        expect = chk.digest()
+        final_chk = self.device.send_recv(CCProtocolPacker.sha256())
+        assert expect == final_chk, "Checksum mismatch after all that?"
+
+        self.device.send_recv(CCProtocolPacker.reboot())
+
+        return {'success': True}
 
 def enumerate(password=''):
     results = []