spi_flash: Support firmware dictionary validation

Updates firmware validation to use a dictionary if provided (or found
when updating from the default out/ directory).

Validation without a dictionary still checks the following (in order):
1. Active firmware's raw dictionary changed after update
2. Checksum of firmware.cur matches expected

Signed-off-by: Justin Schuh <code@justinschuh.com>

scripts/flash-sdcard.sh

@@ -6,6 +6,8 @@
 KLIPPY_ENV="${HOME}/klippy-env/bin/python"
 SRCDIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )"/.. && pwd )"
 KLIPPER_BIN="${SRCDIR}/out/klipper.bin"
+KLIPPER_BIN_DEFAULT=$KLIPPER_BIN
+KLIPPER_DICT_DEFAULT="${SRCDIR}/out/klipper.dict"
 SPI_FLASH="${SRCDIR}/scripts/spi_flash/spi_flash.py"
 BAUD_ARG=""
 # Force script to exit if an error occurs
@@ -15,7 +17,7 @@ print_help_message()
 {
     echo "SD Card upload utility for Klipper"
     echo
-    echo "usage: flash_sdcard.sh [-h] [-l] [-b <baud>] [-f <firmware>]"
+    echo "usage: flash_sdcard.sh [-h] [-l] [-b <baud>] [-f <firmware>] [-d <dictionary>]"
     echo "                       <device> <board>"
     echo
     echo "positional arguments:"
@@ -27,10 +29,11 @@ print_help_message()
     echo "  -l                list available boards"
     echo "  -b <baud>         serial baud rate (default is 250000)"
     echo "  -f <firmware>     path to klipper.bin"
+    echo "  -d <dictionary>   path to klipper.dict for firmware validation"
 }
 
 # Parse command line "optional args"
-while getopts "hlb:f:" arg; do
+while getopts "hlb:f:d:" arg; do
     case $arg in
         h)
             print_help_message
@@ -42,6 +45,7 @@ while getopts "hlb:f:" arg; do
             ;;
         b) BAUD_ARG="-b ${OPTARG}";;
         f) KLIPPER_BIN=$OPTARG;;
+        d) KLIPPER_DICT=$OPTARG;;
     esac
 done
 
@@ -64,6 +68,18 @@ if [ ! -e $DEVICE ]; then
     exit -1
 fi
 
+if [ ! $KLIPPER_DICT ] && [ $KLIPPER_BIN == $KLIPPER_BIN_DEFAULT ] ; then
+    KLIPPER_DICT=$KLIPPER_DICT_DEFAULT
+fi
+
+if [ $KLIPPER_DICT ]; then
+    if [ ! -f $KLIPPER_DICT ]; then
+        echo "No file found at '${KLIPPER_BIN}'"
+        exit -1
+    fi
+    KLIPPER_DICT="-d ${KLIPPER_DICT}"
+fi
+
 # Run Script
 echo "Flashing ${KLIPPER_BIN} to ${DEVICE}"
-${KLIPPY_ENV} ${SPI_FLASH} ${BAUD_ARG} ${DEVICE} ${BOARD} ${KLIPPER_BIN}
+${KLIPPY_ENV} ${SPI_FLASH} ${BAUD_ARG} ${KLIPPER_DICT} ${DEVICE} ${BOARD} ${KLIPPER_BIN}

scripts/spi_flash/spi_flash.py

@@ -13,6 +13,7 @@ import logging
 import collections
 import time
 import traceback
+import json
 import board_defs
 import fatfs_lib
 import reactor
@@ -951,13 +952,22 @@ class MCUConnection:
             % (fw_path, sd_size, sd_chksm))
         return sd_chksm
 
-    def verify_flash(self, req_chksm, old_dictionary):
+    def verify_flash(self, req_chksm, old_dictionary, req_dictionary):
         output("Verifying Flash...")
         validation_passed = False
         msgparser = self._serial.get_msgparser()
         cur_dictionary = msgparser.get_raw_data_dictionary()
-        # Check that the version changed
-        if cur_dictionary != old_dictionary:
+        # If we have a dictionary, check that it matches.
+        if req_dictionary:
+            if cur_dictionary != req_dictionary:
+                raise SPIFlashError("Version Mismatch: Got '%s...', "
+                                    "expected '%s...'"
+                                    % (msgparser.get_version_info()[0],
+                                       json.loads(req_dictionary)['version']))
+            output("Version matched...")
+            validation_passed = True
+        # Otherwise check that the MCU dictionary changed
+        elif cur_dictionary != old_dictionary:
             output("Version updated...")
             validation_passed = True
         else:
@@ -1018,6 +1028,14 @@ class SPIFlash:
         self.task_complete = False
         self.need_upload = True
         self.old_dictionary = None
+        self.new_dictionary = None
+        if args['klipper_dict_path'] is not None:
+            try:
+                with open(args['klipper_dict_path'], 'rb') as dict_f:
+                    self.new_dictionary = dict_f.read(32*1024)
+            except Exception:
+                raise SPIFlashError("Missing or invalid dictionary at '%s'"
+                                    % (args['klipper_dict_path'],))
 
     def _wait_for_reconnect(self):
         output("Waiting for device to reconnect...")
@@ -1062,7 +1080,8 @@ class SPIFlash:
         # Reconnect and verify
         self.mcu_conn.connect()
         self.mcu_conn.configure_mcu()
-        self.mcu_conn.verify_flash(self.firmware_checksum, self.old_dictionary)
+        self.mcu_conn.verify_flash(self.firmware_checksum, self.old_dictionary,
+                                   self.new_dictionary)
         self.mcu_conn.reset()
         self.task_complete = True
 
@@ -1112,6 +1131,9 @@ def main():
     parser.add_argument(
         "-v", "--verbose", action="store_true",
         help="Enable verbose output")
+    parser.add_argument(
+        "-d", "--dict_path", metavar="<klipper.dict>", type=str,
+        default=None, help="Klipper firmware dictionary")
     parser.add_argument(
         "device", metavar="<device>", help="Device Serial Port")
     parser.add_argument(
@@ -1129,6 +1151,7 @@ def main():
     flash_args['device'] = args.device
     flash_args['baud'] = args.baud
     flash_args['klipper_bin_path'] = args.klipper_bin_path
+    flash_args['klipper_dict_path'] = args.dict_path
     check_need_convert(args.board, flash_args)
     fatfs_lib.check_fatfs_build(output)
     try: