tools.checksum: new module for checksums

Checksums are used in various places, so it makes sense for them to have
their own helper module.

Author: dlech

pybricksdev/connections.py

@@ -10,22 +10,23 @@
 import struct
 
 import asyncssh
-from bleak.backends.device import BLEDevice
-from bleak import BleakClient
 import semver
+from bleak import BleakClient
+from bleak.backends.device import BLEDevice
 from tqdm.auto import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
 
 from .ble import BLEConnection
 from .ble.nus import NUS_RX_UUID, NUS_TX_UUID
 from .ble.pybricks import (
-    Event,
     PYBRICKS_CONTROL_UUID,
     PYBRICKS_PROTOCOL_VERSION,
     SW_REV_UUID,
+    Event,
     Status,
 )
 from .compile import compile_file
+from .tools.checksum import xor_bytes
 from .usbconnection import USBConnection
 
 logger = logging.getLogger(__name__)
@@ -249,9 +250,7 @@ async def send_message(self, data):
             raise ValueError("Cannot send this much data at once")
 
         # Compute expected reply
-        checksum = 0
-        for b in data:
-            checksum ^= b
+        checksum = xor_bytes(data, 0)
 
         # Clear existing checksum
         self.prepare_checksum()
@@ -746,15 +745,9 @@ async def disconnect(self):
         else:
             logger.debug("already disconnected")
 
-    def get_checksum(self, block):
-        checksum = 0
-        for b in block:
-            checksum ^= b
-        return checksum
-
     async def send_block(self, data):
         self.checksum_ready.clear()
-        self.expected_checksum = self.get_checksum(data)
+        self.expected_checksum = xor_bytes(data, 0)
         try:
             await self.client.write_gatt_char(NUS_RX_UUID, bytearray(data), False)
             await asyncio.wait_for(self.checksum_ready.wait(), timeout=0.5)

pybricksdev/flash.py

@@ -3,139 +3,29 @@
 
 import asyncio
 import io
-from collections import namedtuple
 import json
 import logging
 import os
 import platform
 import struct
 import sys
+import typing
+import zipfile
+from collections import namedtuple
 from typing import Dict, Tuple
+
 from tqdm.auto import tqdm
 from tqdm.contrib.logging import logging_redirect_tqdm
-import typing
-import zipfile
 
 from .ble import BLERequestsConnection
 from .ble.lwp3 import BootloaderCommand
 from .ble.lwp3.bytecodes import HubKind
-from .compile import save_script, compile_file
+from .compile import compile_file, save_script
+from .tools.checksum import crc32_checksum, sum_complement
 
 logger = logging.getLogger(__name__)
 
 
-def sum_complement(fw, max_size):
-    """Calculates the checksum of a firmware file using the sum complement
-    method of adding each 32-bit word and the returning the two's complement
-    as the checksum.
-
-    Arguments:
-        fw (file):
-            The firmware file (a binary buffer - e.g. a file opened in 'rb' mode)
-        max_size (int):
-            The maximum size of the firmware file.
-
-    Returns:
-        int: The correction needed to make the checksum of the file == 0.
-    """
-    checksum = 0
-    size = 0
-
-    while True:
-        word = fw.read(4)
-        if not word:
-            break
-        checksum += struct.unpack("I", word)[0]
-        size += 4
-
-    if size > max_size:
-        print('firmware + main.mpy is too large"', file=sys.stderr)
-        exit(1)
-
-    for _ in range(size, max_size, 4):
-        checksum += 0xFFFFFFFF
-
-    checksum &= 0xFFFFFFFF
-    correction = checksum and (1 << 32) - checksum or 0
-
-    return correction
-
-
-# thanks https://stackoverflow.com/a/33152544/1976323
-
-_CRC_TABLE = (
-    0x00000000,
-    0x04C11DB7,
-    0x09823B6E,
-    0x0D4326D9,
-    0x130476DC,
-    0x17C56B6B,
-    0x1A864DB2,
-    0x1E475005,
-    0x2608EDB8,
-    0x22C9F00F,
-    0x2F8AD6D6,
-    0x2B4BCB61,
-    0x350C9B64,
-    0x31CD86D3,
-    0x3C8EA00A,
-    0x384FBDBD,
-)
-
-
-def _dword(value):
-    return value & 0xFFFFFFFF
-
-
-def _crc32_fast(crc, data):
-    crc, data = _dword(crc), _dword(data)
-    crc ^= data
-    for _ in range(8):
-        crc = _dword(crc << 4) ^ _CRC_TABLE[crc >> 28]
-    return crc
-
-
-def _crc32_fast_block(crc, buffer):
-    for data in buffer:
-        crc = _crc32_fast(crc, data)
-    return crc
-
-
-def crc32_checksum(fw, max_size):
-    """Calculate the checksum of a firmware file using CRC-32 as implemented
-    in STM32 microprocessors.
-
-    Parameters
-    ----------
-    fw : file
-        The firmware file (a binary buffer - e.g. a file opened in 'rb' mode)
-    max_size : int
-        The maximum size of the firmware file.
-
-    Returns
-    -------
-    int
-        The checksum
-    """
-
-    # remove the last 4 bytes that are the placeholder for the checksum
-    try:
-        fw = fw.read()[:-4]
-    except AttributeError:
-        fw = fw[:-4]
-    if len(fw) + 4 > max_size:
-        raise ValueError("File is too large")
-
-    if len(fw) & 3:
-        raise ValueError("bytes_data length must be multiple of four")
-
-    crc = 0xFFFFFFFF
-    for index in range(0, len(fw), 4):
-        data = int.from_bytes(fw[index : index + 4], "little")
-        crc = _crc32_fast(crc, data)
-    return crc
-
-
 async def create_firmware(
     firmware_zip: typing.Union[str, os.PathLike, typing.BinaryIO]
 ) -> Tuple[bytes, dict]:

pybricksdev/tools/__init__.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: MIT
+# Copyright (c) 2021 The Pybricks Authors
+
+"""
+Helper functions for calculating checksums.
+"""
+
+from io import BytesIO
+
+
+def xor_bytes(data: bytes, init: int = 0xFF) -> int:
+    """
+    Computes a checksum by xoring each byte in ``data`` starting with ``init``.
+
+    Args:
+        data: The byte data to perform the checksum on.
+        init: The initial value of the checksum.
+
+    Returns:
+        The calculated checksum.
+    """
+    checksum = init
+
+    for b in data:
+        checksum ^= b
+
+    return checksum
+
+
+def sum_complement(data: BytesIO, max_size: int) -> int:
+    """
+    Calculates the checksum of using the sum complement method of adding each
+    32-bit word (little-endian) and the returning the two's complement as the
+    checksum.
+
+    Arguments:
+        data:
+            A binary buffer - e.g. a file opened in 'rb' mode.
+        max_size:
+            The maximum size of the firmware file.
+
+    Returns:
+        The correction needed to make the checksum == 0.
+    """
+    checksum = 0
+    size = 0
+
+    while True:
+        word = data.read(4)
+        if not word:
+            break
+        checksum += int.from_bytes(word, "little")
+        size += 4
+
+    if size > max_size:
+        raise ValueError("data is too large")
+
+    for _ in range(size, max_size, 4):
+        checksum += 0xFFFFFFFF
+
+    checksum &= 0xFFFFFFFF
+    correction = checksum and (1 << 32) - checksum or 0
+
+    return correction
+
+
+# thanks https://stackoverflow.com/a/33152544/1976323
+
+_CRC_TABLE = (
+    0x00000000,
+    0x04C11DB7,
+    0x09823B6E,
+    0x0D4326D9,
+    0x130476DC,
+    0x17C56B6B,
+    0x1A864DB2,
+    0x1E475005,
+    0x2608EDB8,
+    0x22C9F00F,
+    0x2F8AD6D6,
+    0x2B4BCB61,
+    0x350C9B64,
+    0x31CD86D3,
+    0x3C8EA00A,
+    0x384FBDBD,
+)
+
+
+def _dword(value):
+    return value & 0xFFFFFFFF
+
+
+def _crc32_fast(crc, data):
+    crc, data = _dword(crc), _dword(data)
+    crc ^= data
+    for _ in range(8):
+        crc = _dword(crc << 4) ^ _CRC_TABLE[crc >> 28]
+    return crc
+
+
+def crc32_checksum(data: BytesIO, max_size: int) -> int:
+    """
+    Calculate the checksum using CRC-32 as implemented in STM32 microprocessors.
+
+    Args:
+        data:
+            A binary buffer - e.g. a file opened in 'rb' mode.
+        max_size:
+            The maximum size of the firmware file.
+
+    Returns:
+        The checksum.
+    """
+
+    # remove the last 4 bytes that are the placeholder for the checksum
+    try:
+        data = data.read()[:-4]
+    except AttributeError:
+        data = data[:-4]
+
+    if len(data) + 4 > max_size:
+        raise ValueError("data is too large")
+
+    if len(data) & 3:
+        raise ValueError("bytes_data length must be multiple of four")
+
+    crc = 0xFFFFFFFF
+    for index in range(0, len(data), 4):
+        data = int.from_bytes(data[index : index + 4], "little")
+        crc = _crc32_fast(crc, data)
+    return crc

pybricksdev/tools/checksum.py

@@ -0,0 +1,9 @@
+from pybricksdev.tools.checksum import xor_bytes
+
+
+def test_xor_bytes():
+    assert xor_bytes(b"\x00") == 0xFF ^ 0x00
+    assert xor_bytes(b"\x00", 0x00) == 0x00 ^ 0x00
+    assert xor_bytes(b"\xFF") == 0xFF ^ 0xFF
+    assert xor_bytes(b"\xFF", 0x00) == 0x00 ^ 0xFF
+    assert xor_bytes(b"\x01\x02\x03\x04") == 0xFF ^ 0x01 ^ 0x02 ^ 0x03 ^ 0x04