✨ Firmware Upload via Binary Transfer (#23462)

buildroot/share/scripts/MarlinBinaryProtocol.py

@@ -0,0 +1,434 @@
+#
+# MarlinBinaryProtocol.py
+# Supporting Firmware upload via USB/Serial, saving to the attached media.
+#
+import serial
+import math
+import time
+from collections import deque
+import threading
+import sys
+import datetime
+import random
+try:
+    import heatshrink
+    heatshrink_exists = True
+except ImportError:
+    heatshrink_exists = False
+
+
+def millis():
+    return time.perf_counter() * 1000
+
+class TimeOut(object):
+    def __init__(self, milliseconds):
+        self.duration = milliseconds
+        self.reset()
+
+    def reset(self):
+        self.endtime = millis() + self.duration
+
+    def timedout(self):
+        return millis() > self.endtime
+
+class ReadTimeout(Exception):
+    pass
+class FatalError(Exception):
+    pass
+class SycronisationError(Exception):
+    pass
+class PayloadOverflow(Exception):
+    pass
+class ConnectionLost(Exception):
+    pass
+
+class Protocol(object):
+    device = None
+    baud = None
+    max_block_size = 0
+    port = None
+    block_size = 0
+
+    packet_transit = None
+    packet_status = None
+    packet_ping = None
+
+    errors = 0
+    packet_buffer = None
+    simulate_errors = 0
+    sync = 0
+    connected = False
+    syncronised = False
+    worker_thread = None
+
+    response_timeout = 1000
+
+    applications = []
+    responses = deque()
+
+    def __init__(self, device, baud, bsize, simerr, timeout):
+        print("pySerial Version:", serial.VERSION)
+        self.port = serial.Serial(device, baudrate = baud, write_timeout = 0, timeout = 1)
+        self.device = device
+        self.baud = baud
+        self.block_size = int(bsize)
+        self.simulate_errors = max(min(simerr, 1.0), 0.0);
+        self.connected = True
+        self.response_timeout = timeout
+
+        self.register(['ok', 'rs', 'ss', 'fe'], self.process_input)
+
+        self.worker_thread = threading.Thread(target=Protocol.receive_worker, args=(self,))
+        self.worker_thread.start()
+
+    def receive_worker(self):
+        while self.port.in_waiting:
+            self.port.reset_input_buffer()
+
+        def dispatch(data):
+            for tokens, callback in self.applications:
+                for token in tokens:
+                    if token == data[:len(token)]:
+                        callback((token, data[len(token):]))
+                        return
+
+        def reconnect():
+            print("Reconnecting..")
+            self.port.close()
+            for x in range(10):
+                try:
+                    if self.connected:
+                        self.port = serial.Serial(self.device, baudrate = self.baud, write_timeout = 0, timeout = 1)
+                        return
+                    else:
+                        print("Connection closed")
+                        return
+                except:
+                    time.sleep(1)
+            raise ConnectionLost()
+
+        while self.connected:
+            try:
+                data = self.port.readline().decode('utf8').rstrip()
+                if len(data):
+                    #print(data)
+                    dispatch(data)
+            except OSError:
+                reconnect()
+            except UnicodeDecodeError:
+                # dodgy client output or datastream corruption
+                self.port.reset_input_buffer()
+
+    def shutdown(self):
+        self.connected = False
+        self.worker_thread.join()
+        self.port.close()
+
+    def process_input(self, data):
+        #print(data)
+        self.responses.append(data)
+
+    def register(self, tokens, callback):
+        self.applications.append((tokens, callback))
+
+    def send(self, protocol, packet_type, data = bytearray()):
+        self.packet_transit = self.build_packet(protocol, packet_type, data)
+        self.packet_status = 0
+        self.transmit_attempt = 0
+
+        timeout = TimeOut(self.response_timeout * 20)
+        while self.packet_status == 0:
+            try:
+                if timeout.timedout():
+                    raise ConnectionLost()
+                self.transmit_packet(self.packet_transit)
+                self.await_response()
+            except ReadTimeout:
+                self.errors += 1
+                #print("Packetloss detected..")
+        self.packet_transit = None
+
+    def await_response(self):
+        timeout = TimeOut(self.response_timeout)
+        while not len(self.responses):
+            time.sleep(0.00001)
+            if timeout.timedout():
+                raise ReadTimeout()
+
+        while len(self.responses):
+            token, data = self.responses.popleft()
+            switch = {'ok' : self.response_ok, 'rs': self.response_resend, 'ss' : self.response_stream_sync, 'fe' : self.response_fatal_error}
+            switch[token](data)
+
+    def send_ascii(self, data, send_and_forget = False):
+        self.packet_transit = bytearray(data, "utf8") + b'\n'
+        self.packet_status = 0
+        self.transmit_attempt = 0
+
+        timeout = TimeOut(self.response_timeout * 20)
+        while self.packet_status == 0:
+            try:
+                if timeout.timedout():
+                    return
+                self.port.write(self.packet_transit)
+                if send_and_forget:
+                    self.packet_status = 1
+                else:
+                    self.await_response_ascii()
+            except ReadTimeout:
+                self.errors += 1
+                #print("Packetloss detected..")
+            except serial.serialutil.SerialException:
+                return
+        self.packet_transit = None
+
+    def await_response_ascii(self):
+        timeout = TimeOut(self.response_timeout)
+        while not len(self.responses):
+            time.sleep(0.00001)
+            if timeout.timedout():
+                raise ReadTimeout()
+        token, data = self.responses.popleft()
+        self.packet_status = 1
+
+    def corrupt_array(self, data):
+        rid = random.randint(0, len(data) - 1)
+        data[rid] ^= 0xAA
+        return data
+
+    def transmit_packet(self, packet):
+        packet = bytearray(packet)
+        if(self.simulate_errors > 0 and random.random() > (1.0 - self.simulate_errors)):
+            if random.random() > 0.9:
+                #random data drop
+                start = random.randint(0, len(packet))
+                end = start + random.randint(1, 10)
+                packet = packet[:start] + packet[end:]
+                #print("Dropping {0} bytes".format(end - start))
+            else:
+                #random corruption
+                packet = self.corrupt_array(packet)
+                #print("Single byte corruption")
+        self.port.write(packet)
+        self.transmit_attempt += 1
+
+    def build_packet(self, protocol, packet_type, data = bytearray()):
+        PACKET_TOKEN = 0xB5AD
+
+        if len(data) > self.max_block_size:
+            raise PayloadOverflow()
+
+        packet_buffer = bytearray()
+
+        packet_buffer += self.pack_int8(self.sync)                           # 8bit sync id
+        packet_buffer += self.pack_int4_2(protocol, packet_type)             # 4 bit protocol id, 4 bit packet type
+        packet_buffer += self.pack_int16(len(data))                          # 16bit packet length
+        packet_buffer += self.pack_int16(self.build_checksum(packet_buffer)) # 16bit header checksum
+
+        if len(data):
+            packet_buffer += data
+            packet_buffer += self.pack_int16(self.build_checksum(packet_buffer))
+
+        packet_buffer =  self.pack_int16(PACKET_TOKEN) + packet_buffer       # 16bit start token, not included in checksum
+        return packet_buffer
+
+    # checksum 16 fletchers
+    def checksum(self, cs, value):
+        cs_low = (((cs & 0xFF) + value) % 255);
+        return ((((cs >> 8) + cs_low) % 255) << 8) | cs_low;
+
+    def build_checksum(self, buffer):
+        cs = 0
+        for b in buffer:
+            cs = self.checksum(cs, b)
+        return cs
+
+    def pack_int32(self, value):
+        return value.to_bytes(4, byteorder='little')
+
+    def pack_int16(self, value):
+        return value.to_bytes(2, byteorder='little')
+
+    def pack_int8(self, value):
+        return value.to_bytes(1, byteorder='little')
+
+    def pack_int4_2(self, vh, vl):
+        value = ((vh & 0xF) << 4) | (vl & 0xF)
+        return value.to_bytes(1, byteorder='little')
+
+    def connect(self):
+        print("Connecting: Switching Marlin to Binary Protocol...")
+        self.send_ascii("M28B1")
+        self.send(0, 1)
+
+    def disconnect(self):
+        self.send(0, 2)
+        self.syncronised = False
+
+    def response_ok(self, data):
+        try:
+            packet_id = int(data);
+        except ValueError:
+            return
+        if packet_id != self.sync:
+            raise SycronisationError()
+        self.sync = (self.sync + 1) % 256
+        self.packet_status = 1
+
+    def response_resend(self, data):
+        packet_id = int(data);
+        self.errors += 1
+        if not self.syncronised:
+            print("Retrying syncronisation")
+        elif packet_id != self.sync:
+            raise SycronisationError()
+
+    def response_stream_sync(self, data):
+        sync, max_block_size, protocol_version = data.split(',')
+        self.sync = int(sync)
+        self.max_block_size = int(max_block_size)
+        self.block_size = self.max_block_size if self.max_block_size < self.block_size else self.block_size
+        self.protocol_version = protocol_version
+        self.packet_status = 1
+        self.syncronised = True
+        print("Connection synced [{0}], binary protocol version {1}, {2} byte payload buffer".format(self.sync, self.protocol_version, self.max_block_size))
+
+    def response_fatal_error(self, data):
+        raise FatalError()
+
+
+class FileTransferProtocol(object):
+    protocol_id = 1
+
+    class Packet(object):
+        QUERY = 0
+        OPEN  = 1
+        CLOSE = 2
+        WRITE = 3
+        ABORT = 4
+
+    responses = deque()
+    def __init__(self, protocol, timeout = None):
+        protocol.register(['PFT:success', 'PFT:version:', 'PFT:fail', 'PFT:busy', 'PFT:ioerror', 'PTF:invalid'], self.process_input)
+        self.protocol = protocol
+        self.response_timeout = timeout or protocol.response_timeout
+
+    def process_input(self, data):
+        #print(data)
+        self.responses.append(data)
+
+    def await_response(self, timeout = None):
+        timeout = TimeOut(timeout or self.response_timeout)
+        while not len(self.responses):
+            time.sleep(0.0001)
+            if timeout.timedout():
+                raise ReadTimeout()
+
+        return self.responses.popleft()
+
+    def connect(self):
+        self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.QUERY);
+
+        token, data = self.await_response()
+        if token != 'PFT:version:':
+            return False
+
+        self.version, _, compression = data.split(':')
+        if compression != 'none':
+            algorithm, window, lookahead = compression.split(',')
+            self.compression = {'algorithm': algorithm, 'window': int(window), 'lookahead': int(lookahead)}
+        else:
+            self.compression = {'algorithm': 'none'}
+
+        print("File Transfer version: {0}, compression: {1}".format(self.version, self.compression['algorithm']))
+
+    def open(self, filename, compression, dummy):
+        payload =  b'\1' if dummy else b'\0'          # dummy transfer
+        payload += b'\1' if compression else b'\0'    # payload compression
+        payload += bytearray(filename, 'utf8') + b'\0'# target filename + null terminator
+
+        timeout = TimeOut(5000)
+        token = None
+        self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.OPEN, payload);
+        while token != 'PFT:success' and not timeout.timedout():
+            try:
+                token, data = self.await_response(1000)
+                if token == 'PFT:success':
+                    print(filename,"opened")
+                    return
+                elif token == 'PFT:busy':
+                    print("Broken transfer detected, purging")
+                    self.abort()
+                    time.sleep(0.1)
+                    self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.OPEN, payload);
+                    timeout.reset()
+                elif token == 'PFT:fail':
+                    raise Exception("Can not open file on client")
+            except ReadTimeout:
+                pass
+        raise ReadTimeout()
+
+    def write(self, data):
+        self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.WRITE, data);
+
+    def close(self):
+        self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.CLOSE);
+        token, data = self.await_response(1000)
+        if token == 'PFT:success':
+            print("File closed")
+            return
+        elif token == 'PFT:ioerror':
+            print("Client storage device IO error")
+        elif token == 'PFT:invalid':
+            print("No open file")
+
+    def abort(self):
+        self.protocol.send(FileTransferProtocol.protocol_id, FileTransferProtocol.Packet.ABORT);
+        token, data = self.await_response()
+        if token == 'PFT:success':
+            print("Transfer Aborted")
+
+    def copy(self, filename, dest_filename, compression, dummy):
+        self.connect()
+
+        compression_support = heatshrink_exists and self.compression['algorithm'] == 'heatshrink' and compression
+        if compression and (not heatshrink_exists or not self.compression['algorithm'] == 'heatshrink'):
+            print("Compression not supported by client")
+        #compression_support = False
+
+        data = open(filename, "rb").read()
+        filesize = len(data)
+
+        self.open(dest_filename, compression_support, dummy)
+
+        block_size = self.protocol.block_size
+        if compression_support:
+            data = heatshrink.encode(data, window_sz2=self.compression['window'], lookahead_sz2=self.compression['lookahead'])
+
+        cratio = filesize / len(data)
+
+        blocks = math.floor((len(data) + block_size - 1) / block_size)
+        kibs = 0
+        dump_pctg = 0
+        start_time = millis()
+        for i in range(blocks):
+            start = block_size * i
+            end = start + block_size
+            self.write(data[start:end])
+            kibs = (( (i+1) * block_size) / 1024) / (millis() + 1 - start_time) * 1000
+            if (i / blocks) >= dump_pctg:
+                print("\r{0:2.2f}% {1:4.2f}KiB/s {2} Errors: {3}".format((i / blocks) * 100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors), end='')
+                dump_pctg += 0.1
+        print("\r{0:2.2f}% {1:4.2f}KiB/s {2} Errors: {3}".format(100, kibs, "[{0:4.2f}KiB/s]".format(kibs * cratio) if compression_support else "", self.protocol.errors)) # no one likes transfers finishing at 99.8%
+
+        self.close()
+        print("Transfer complete")
+
+
+class EchoProtocol(object):
+    def __init__(self, protocol):
+        protocol.register(['echo:'], self.process_input)
+        self.protocol = protocol
+
+    def process_input(self, data):
+        print(data)

buildroot/share/scripts/upload.py

@@ -0,0 +1,274 @@
+import argparse
+import sys
+import os
+import time
+import random
+import serial
+
+Import("env")
+
+# Needed (only) for compression, but there are problems with pip install heatshrink
+#try:
+#    import heatshrink
+#except ImportError:
+#    # Install heatshrink
+#    print("Installing 'heatshrink' python module...")
+#    env.Execute(env.subst("$PYTHONEXE -m pip install heatshrink"))
+#
+# Not tested: If it's safe to install python libraries in PIO python try:
+#    env.Execute(env.subst("$PYTHONEXE -m pip install https://github.com/p3p/pyheatshrink/releases/download/0.3.3/pyheatshrink-pip.zip"))
+
+import MarlinBinaryProtocol
+
+# Internal debug flag
+Debug = False
+
+#-----------------#
+# Upload Callback #
+#-----------------#
+def Upload(source, target, env):
+
+    #------------------#
+    # Marlin functions #
+    #------------------#
+    def _GetMarlinEnv(marlinEnv, feature):
+        if not marlinEnv: return None
+        return marlinEnv[feature] if feature in marlinEnv else None
+
+    #----------------#
+    # Port functions #
+    #----------------#
+    def _GetUploadPort(env):
+        if Debug: print('Autodetecting upload port...')
+        env.AutodetectUploadPort(env)
+        port = env.subst('$UPLOAD_PORT')
+        if not port:
+            raise Exception('Error detecting the upload port.')
+        if Debug: print('OK')
+        return port
+
+    #-------------------------#
+    # Simple serial functions #
+    #-------------------------#
+    def _Send(data):
+        if Debug: print(f'>> {data}')
+        strdata = bytearray(data, 'utf8') + b'\n'
+        port.write(strdata)
+        time.sleep(0.010)
+
+    def _Recv():
+        clean_responses = []
+        responses = port.readlines()
+        for Resp in responses:
+            # Test: suppress invaid chars (coming from debug info)
+            try:
+                clean_response = Resp.decode('utf8').rstrip().lstrip()
+                clean_responses.append(clean_response)
+            except:
+                pass
+            if Debug: print(f'<< {clean_response}')
+        return clean_responses
+
+    #------------------#
+    # SDCard functions #
+    #------------------#
+    def _CheckSDCard():
+        if Debug: print('Checking SD card...')
+        _Send('M21')
+        Responses = _Recv()
+        if len(Responses) < 1 or not any('SD card ok' in r for r in Responses):
+            raise Exception('Error accessing SD card')
+        if Debug: print('SD Card OK')
+        return True
+
+    #----------------#
+    # File functions #
+    #----------------#
+    def _GetFirmwareFiles():
+        if Debug: print('Get firmware files...')
+        _Send('M20 F')
+        Responses = _Recv()
+        if len(Responses) < 3 or not any('file list' in r for r in Responses):
+            raise Exception('Error getting firmware files')
+        if Debug: print('OK')
+        return Responses
+
+    def _FilterFirmwareFiles(FirmwareList):
+        Firmwares = []
+        for FWFile in FirmwareList:
+            if not '/' in FWFile and '.BIN' in FWFile:
+                idx = FWFile.index('.BIN')
+                Firmwares.append(FWFile[:idx+4])
+        return Firmwares
+
+    def _RemoveFirmwareFile(FirmwareFile):
+        _Send(f'M30 /{FirmwareFile}')
+        Responses = _Recv()
+        Removed = len(Responses) >= 1 and any('File deleted' in r for r in Responses)
+        if not Removed:
+            raise Exception(f"Firmware file '{FirmwareFile}' not removed")
+        return Removed
+
+
+    #---------------------#
+    # Callback Entrypoint #
+    #---------------------#
+    port = None
+    protocol = None
+    filetransfer = None
+
+    # Get Marlin evironment vars
+    MarlinEnv = env['MARLIN_FEATURES']
+    marlin_pioenv = _GetMarlinEnv(MarlinEnv, 'PIOENV')
+    marlin_motherboard = _GetMarlinEnv(MarlinEnv, 'MOTHERBOARD')
+    marlin_board_info_name = _GetMarlinEnv(MarlinEnv, 'BOARD_INFO_NAME')
+    marlin_board_custom_build_flags = _GetMarlinEnv(MarlinEnv, 'BOARD_CUSTOM_BUILD_FLAGS')
+    marlin_firmware_bin = _GetMarlinEnv(MarlinEnv, 'FIRMWARE_BIN')
+    marlin_custom_firmware_upload = _GetMarlinEnv(MarlinEnv, 'CUSTOM_FIRMWARE_UPLOAD') is not None
+    marlin_short_build_version = _GetMarlinEnv(MarlinEnv, 'SHORT_BUILD_VERSION')
+    marlin_string_config_h_author = _GetMarlinEnv(MarlinEnv, 'STRING_CONFIG_H_AUTHOR')
+
+    # Get firmware upload params
+    upload_firmware_source_name = str(source[0])    # Source firmware filename
+    upload_speed = env['UPLOAD_SPEED'] if 'UPLOAD_SPEED' in env else 115200
+                                                    # baud rate of serial connection
+    upload_port = _GetUploadPort(env)               # Serial port to use
+
+    # Set local upload params
+    upload_firmware_target_name = os.path.basename(upload_firmware_source_name)     # WARNING! Need rework on "binary_stream" to allow filename > 8.3
+                                                    # Target firmware filename
+    upload_timeout = 1000                           # Communication timout, lossy/slow connections need higher values
+    upload_blocksize = 512                          # Transfer block size. 512 = Autodetect
+    upload_compression = True                       # Enable compression
+    upload_error_ratio = 0                          # Simulated corruption ratio
+    upload_test = False                             # Benchmark the serial link without storing the file
+    upload_reset = True                             # Trigger a soft reset for firmware update after the upload
+
+    # Set local upload params based on board type to change script behavior
+    # "upload_delete_old_bins": delete all *.bin files in the root of SD Card
+    upload_delete_old_bins = marlin_motherboard in ['BOARD_CREALITY_V4',   'BOARD_CREALITY_V4210', 'BOARD_CREALITY_V423', 'BOARD_CREALITY_V427',
+                                                    'BOARD_CREALITY_V431', 'BOARD_CREALITY_V452',  'BOARD_CREALITY_V453', 'BOARD_CREALITY_V24S1']
+    try:
+
+        # Start upload job
+        print(f"Uploading firmware '{os.path.basename(upload_firmware_target_name)}' to '{marlin_motherboard}' via '{upload_port}'")
+
+        # Dump some debug info
+        if Debug:
+            print('Upload using:')
+            print('---- Marlin --------------------')
+            print(f' PIOENV                 : {marlin_pioenv}')
+            print(f' SHORT_BUILD_VERSION    : {marlin_short_build_version}')
+            print(f' STRING_CONFIG_H_AUTHOR : {marlin_string_config_h_author}')
+            print(f' MOTHERBOARD            : {marlin_motherboard}')
+            print(f' BOARD_INFO_NAME        : {marlin_board_info_name}')
+            print(f' CUSTOM_BUILD_FLAGS     : {marlin_board_custom_build_flags}')
+            print(f' FIRMWARE_BIN           : {marlin_firmware_bin}')
+            print(f' CUSTOM_FIRMWARE_UPLOAD : {marlin_custom_firmware_upload}')
+            print('---- Upload parameters ---------')
+            print(f' Source      : {upload_firmware_source_name}')
+            print(f' Target      : {upload_firmware_target_name}')
+            print(f' Port        : {upload_port} @ {upload_speed} baudrate')
+            print(f' Timeout     : {upload_timeout}')
+            print(f' Block size  : {upload_blocksize}')
+            print(f' Compression : {upload_compression}')
+            print(f' Error ratio : {upload_error_ratio}')
+            print(f' Test        : {upload_test}')
+            print(f' Reset       : {upload_reset}')
+            print('--------------------------------')
+
+        # Custom implementations based on board parameters
+
+        # Delete all *.bin files on the root of SD Card (if flagged)
+        if upload_delete_old_bins:
+            # CUSTOM_FIRMWARE_UPLOAD is needed for this feature
+            if not marlin_custom_firmware_upload:
+                raise Exception(f"CUSTOM_FIRMWARE_UPLOAD must be enabled in 'Configuration_adv.h' for '{marlin_motherboard}'")
+
+            # Generate a new 8.3 random filename
+            # This board remember the last firmware filename and doesn't allow to flash from that filename
+            upload_firmware_target_name = f"fw-{''.join(random.choices('ABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789', k=5))}.BIN"
+            print(f"Board {marlin_motherboard}: Overriding firmware filename to '{upload_firmware_target_name}'")
+
+            # Init serial port
+            port = serial.Serial(upload_port, baudrate = upload_speed, write_timeout = 0, timeout = 0.1)
+            port.reset_input_buffer()
+
+            # Check SD card status
+            _CheckSDCard()
+
+            # Get firmware files
+            FirmwareFiles = _GetFirmwareFiles()
+            if Debug:
+                for FirmwareFile in FirmwareFiles:
+                    print(f'Found: {FirmwareFile}')
+
+            # Get all 1st level firmware files (to remove)
+            OldFirmwareFiles = _FilterFirmwareFiles(FirmwareFiles[1:len(FirmwareFiles)-2])   # Skip header and footers of list
+            if len(OldFirmwareFiles) == 0:
+                print('No old firmware files to delete')
+            else:
+                print(f"Remove {len(OldFirmwareFiles)} old firmware file{'s' if len(OldFirmwareFiles) != 1 else ''}:")
+                for OldFirmwareFile in OldFirmwareFiles:
+                    print(f" -Removing- '{OldFirmwareFile}'...")
+                    print(' OK' if _RemoveFirmwareFile(OldFirmwareFile) else ' Error!')
+
+            # Close serial
+            port.close()
+
+            # Cleanup completed
+            if Debug: print('Cleanup completed')
+
+        # WARNING! The serial port must be closed here because the serial transfer that follow needs it!
+
+        # Upload firmware file
+        if Debug: print(f"Copy '{upload_firmware_source_name}' --> '{upload_firmware_target_name}'")
+        protocol = MarlinBinaryProtocol.Protocol(upload_port, upload_speed, upload_blocksize, float(upload_error_ratio), int(upload_timeout))
+        #echologger = MarlinBinaryProtocol.EchoProtocol(protocol)
+        protocol.connect()
+        filetransfer = MarlinBinaryProtocol.FileTransferProtocol(protocol)
+        filetransfer.copy(upload_firmware_source_name, upload_firmware_target_name, upload_compression, upload_test)
+        protocol.disconnect()
+
+        # Notify upload completed
+        protocol.send_ascii('M117 Firmware uploaded')
+
+        # Remount SD card
+        print('Wait for SD card release...')
+        time.sleep(1)
+        print('Remount SD card')
+        protocol.send_ascii('M21')
+
+        # Trigger firmware update
+        if upload_reset:
+            print('Trigger firmware update...')
+            protocol.send_ascii('M997', True)
+
+        protocol: protocol.shutdown()
+        print('Firmware update completed')
+
+    except KeyboardInterrupt:
+        if port: port.close()
+        if filetransfer: filetransfer.abort()
+        if protocol: protocol.shutdown()
+        raise
+
+    except serial.SerialException as se:
+        if port: port.close()
+        print(f'Serial excepion: {se}')
+        raise Exception(se)
+
+    except MarlinBinaryProtocol.FatalError:
+        if port: port.close()
+        if protocol: protocol.shutdown()
+        print('Too many retries, Abort')
+        raise
+
+    except:
+        if port: port.close()
+        if protocol: protocol.shutdown()
+        print('Firmware not updated')
+        raise
+
+# Attach custom upload callback
+env.Replace(UPLOADCMD=Upload)