image loader improvements

flashloader/image-loader.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-from typing import List
+from typing import List, Tuple
 from spacepackets.ecss.defs import PusService
 from spacepackets.ecss.tm import PusTm
 from tmtccmd.com import ComInterface
@@ -60,6 +60,7 @@ class ActionId(enum.IntEnum):
 
 
 _LOGGER = logging.getLogger(__name__)
+SEQ_PROVIDER = SeqCountProvider(bit_width=14)
 
 
 @dataclasses.dataclass
@@ -70,7 +71,174 @@ class LoadableSegment:
     data: bytes
 
 
-SEQ_PROVIDER = SeqCountProvider(bit_width=14)
+class Target(enum.Enum):
+    BOOTLOADER = 0
+    APP_A = 1
+    APP_B = 1
+
+
+class ImageLoader:
+    def __init__(self, com_if: ComInterface, verificator: PusVerificator) -> None:
+        self.com_if = com_if
+        self.verificator = verificator
+
+    def handle_ping_cmd(self):
+        _LOGGER.info("Sending ping command")
+        ping_tc = PusTc(
+            apid=0x00,
+            service=PusService.S17_TEST,
+            subservice=1,
+            seq_count=SEQ_PROVIDER.get_and_increment(),
+            app_data=bytes(PING_PAYLOAD_SIZE),
+        )
+        self.verificator.add_tc(ping_tc)
+        self.com_if.send(bytes(ping_tc.pack()))
+
+        data_available = self.com_if.data_available(0.4)
+        if not data_available:
+            _LOGGER.warning("no ping reply received")
+        for reply in self.com_if.receive():
+            result = self.verificator.add_tm(
+                Service1Tm.from_tm(PusTm.unpack(reply, 0), UnpackParams(0))
+            )
+            if result is not None and result.completed:
+                _LOGGER.info("received ping completion reply")
+
+    def handle_corruption_cmd(self, target: Target):
+
+        if target == Target.BOOTLOADER:
+            _LOGGER.error("can not corrupt bootloader")
+        if target == Target.APP_A:
+            self.send_tc(
+                PusTc(
+                    apid=0,
+                    service=ACTION_SERVICE,
+                    subservice=ActionId.CORRUPT_APP_A,
+                ),
+            )
+        if target == Target.APP_B:
+            self.send_tc(
+                PusTc(
+                    apid=0,
+                    service=ACTION_SERVICE,
+                    subservice=ActionId.CORRUPT_APP_B,
+                ),
+            )
+
+    def handle_flash_cmd(self, target: Target, file_path: Path) -> int:
+        loadable_segments = []
+        _LOGGER.info("Parsing ELF file for loadable sections")
+        total_size = 0
+        loadable_segments, total_size = create_loadable_segments(target, file_path)
+        segments_info_str(target, loadable_segments, total_size, file_path)
+        result = self._perform_flashing_algorithm(loadable_segments)
+        if result != 0:
+            return result
+        self._crc_and_app_size_postprocessing(target, total_size, loadable_segments)
+        return 0
+
+    def _perform_flashing_algorithm(
+        self,
+        loadable_segments: List[LoadableSegment],
+    ) -> int:
+        # Perform the flashing algorithm.
+        for segment in loadable_segments:
+            segment_end = segment.offset + segment.size
+            current_addr = segment.offset
+            pos_in_segment = 0
+            while pos_in_segment < segment.size:
+                next_chunk_size = min(segment_end - current_addr, CHUNK_SIZE)
+                data = segment.data[pos_in_segment : pos_in_segment + next_chunk_size]
+                next_packet = pack_memory_write_command(current_addr, data)
+                _LOGGER.info(
+                    f"Sending memory write command for address {current_addr:#08x} and data with "
+                    f"length {len(data)}"
+                )
+                self.verificator.add_tc(next_packet)
+                self.com_if.send(bytes(next_packet.pack()))
+                current_addr += next_chunk_size
+                pos_in_segment += next_chunk_size
+                start_time = time.time()
+                while True:
+                    if time.time() - start_time > 1.0:
+                        _LOGGER.error("Timeout while waiting for reply")
+                        return -1
+                    data_available = self.com_if.data_available(0.1)
+                    done = False
+                    if not data_available:
+                        continue
+                    replies = self.com_if.receive()
+                    for reply in replies:
+                        tm = PusTm.unpack(reply, 0)
+                        if tm.service != 1:
+                            continue
+                        service_1_tm = Service1Tm.from_tm(tm, UnpackParams(0))
+                        check_result = self.verificator.add_tm(service_1_tm)
+                        # We could send after we have received the step reply, but that can
+                        # somehow lead to overrun errors. I think it's okay to do it like
+                        # this as long as the flash loader only uses polling..
+                        if (
+                            check_result is not None
+                            and check_result.status.completed == StatusField.SUCCESS
+                        ):
+                            done = True
+
+                        # This is an optimized variant, but I think the small delay is not an issue.
+                        """
+                        if (
+                            check_result is not None
+                            and check_result.status.step == StatusField.SUCCESS
+                            and len(check_result.status.step_list) == 1
+                        ):
+                            done = True
+                        """
+                    self.verificator.remove_completed_entries()
+                    if done:
+                        break
+        return 0
+
+    def _crc_and_app_size_postprocessing(
+        self,
+        target: Target,
+        total_size: int,
+        loadable_segments: List[LoadableSegment],
+    ):
+        if target == Target.BOOTLOADER:
+            _LOGGER.info("Blanking the bootloader checksum")
+            # Blank the checksum. For the bootloader, the bootloader will calculate the
+            # checksum itself on the initial run.
+            checksum_write_packet = pack_memory_write_command(
+                BOOTLOADER_CRC_ADDR, bytes([0x00, 0x00, 0x00, 0x00])
+            )
+            self.send_tc(checksum_write_packet)
+        else:
+            crc_addr = None
+            size_addr = None
+            if target == Target.APP_A:
+                crc_addr = APP_A_CRC_ADDR
+                size_addr = APP_A_SIZE_ADDR
+            elif target == Target.APP_B:
+                crc_addr = APP_B_CRC_ADDR
+                size_addr = APP_B_SIZE_ADDR
+            assert crc_addr is not None
+            assert size_addr is not None
+            _LOGGER.info(f"Writing app size {total_size} at address {size_addr:#08x}")
+            size_write_packet = pack_memory_write_command(
+                size_addr, struct.pack("!I", total_size)
+            )
+            self.com_if.send(bytes(size_write_packet.pack()))
+            time.sleep(0.2)
+            crc_calc = PredefinedCrc("crc-32")
+            for segment in loadable_segments:
+                crc_calc.update(segment.data)
+            checksum = crc_calc.digest()
+            _LOGGER.info(
+                f"Writing checksum 0x[{checksum.hex(sep=',')}] at address {crc_addr:#08x}"
+            )
+            self.send_tc(pack_memory_write_command(crc_addr, checksum))
+
+    def send_tc(self, tc: PusTc):
+        self.com_if.send(bytes(tc.pack()))
 
 
 def main() -> int:
@@ -110,74 +278,56 @@ def main() -> int:
     verificator = PusVerificator()
     com_if = SerialCobsComIF(serial_cfg)
     com_if.open()
+    target = None
+    if args.target == "bl":
+        target = Target.BOOTLOADER
+    elif args.target == "a":
+        target = Target.APP_A
+    elif args.target == "b":
+        target = Target.APP_B
+    image_loader = ImageLoader(com_if, verificator)
     file_path = None
+    result = -1
     if args.ping:
-        _LOGGER.info("Sending ping command")
+        image_loader.handle_ping_cmd()
-        ping_tc = PusTc(
+        result = 0
-            apid=0x00,
+    if target:
-            service=PusService.S17_TEST,
-            subservice=1,
-            seq_count=SEQ_PROVIDER.get_and_increment(),
-            app_data=bytes(PING_PAYLOAD_SIZE),
-        )
-        verificator.add_tc(ping_tc)
-        com_if.send(ping_tc.pack())
-
-        data_available = com_if.data_available(0.4)
-        if not data_available:
-            _LOGGER.warning("no ping reply received")
-        for reply in com_if.receive():
-            result = verificator.add_tm(
-                Service1Tm.from_tm(PusTm.unpack(reply, 0), UnpackParams(0))
-            )
-            if result is not None and result.completed:
-                _LOGGER.info("received ping completion reply")
-        if not args.target:
-            return 0
-    if args.target:
         if not args.corrupt:
             if not args.path:
                 _LOGGER.error("App Path needs to be specified for the flash process")
-                return -1
             file_path = Path(args.path)
             if not file_path.exists():
                 _LOGGER.error("File does not exist")
-                return -1
     if args.corrupt:
-        if not args.target:
+        if target:
+            image_loader.handle_corruption_cmd(target)
+            result = 0
+        else:
             _LOGGER.error("target for corruption command required")
-            return -1
-        if args.target == "bl":
-            _LOGGER.error("can not corrupt bootloader")
-        if args.target == "a":
-            packet = PusTc(
-                apid=0,
-                service=ACTION_SERVICE,
-                subservice=ActionId.CORRUPT_APP_A,
-            )
-            com_if.send(packet.pack())
-        if args.target == "b":
-            packet = PusTc(
-                apid=0,
-                service=ACTION_SERVICE,
-                subservice=ActionId.CORRUPT_APP_B,
-            )
-            com_if.send(packet.pack())
     else:
         assert file_path is not None
+        assert target is not None
+        result = image_loader.handle_flash_cmd(target, file_path)
+
+    com_if.close()
+    return result
+
+
+def create_loadable_segments(
+    target: Target, file_path: Path
+) -> Tuple[List[LoadableSegment], int]:
     loadable_segments = []
-        _LOGGER.info("Parsing ELF file for loadable sections")
     total_size = 0
     with open(file_path, "rb") as app_file:
         elf_file = ELFFile(app_file)
 
-            for (idx, segment) in enumerate(elf_file.iter_segments("PT_LOAD")):
+        for idx, segment in enumerate(elf_file.iter_segments("PT_LOAD")):
             if segment.header.p_filesz == 0:
                 continue
             # Basic validity checks of the base addresses.
             if idx == 0:
                 if (
-                        args.target == "bl"
+                    target == Target.BOOTLOADER
                     and segment.header.p_paddr != BOOTLOADER_START_ADDR
                 ):
                     raise ValueError(
@@ -185,7 +335,7 @@ def main() -> int:
                         f"bootloader, expected {BOOTLOADER_START_ADDR}"
                     )
                 if (
-                        args.target == "a"
+                    target == Target.APP_A
                     and segment.header.p_paddr != APP_A_START_ADDR
                 ):
                     raise ValueError(
@@ -193,7 +343,7 @@ def main() -> int:
                         f"App A, expected {APP_A_START_ADDR}"
                     )
                 if (
-                        args.target == "b"
+                    target == Target.APP_B
                     and segment.header.p_paddr != APP_B_START_ADDR
                 ):
                     raise ValueError(
@@ -222,150 +372,39 @@ def main() -> int:
                 )
             )
             total_size += segment.header.p_filesz
-            context_str = None
+    return loadable_segments, total_size
 
+
+def segments_info_str(
+    target: Target,
+    loadable_segments: List[LoadableSegment],
+    total_size: int,
+    file_path: Path,
+):
     # Set context string and perform basic sanity checks.
-            if args.target == "bl":
+    if target == Target.BOOTLOADER:
         if total_size > BOOTLOADER_MAX_SIZE:
             _LOGGER.error(
                 f"provided bootloader app larger than allowed {total_size} bytes"
             )
             return -1
         context_str = "Bootloader"
-            elif args.target == "a":
+    elif target == Target.APP_A:
         if total_size > APP_A_MAX_SIZE:
-                    _LOGGER.error(
+            _LOGGER.error(f"provided App A larger than allowed {total_size} bytes")
-                        f"provided App A larger than allowed {total_size} bytes"
-                    )
             return -1
         context_str = "App Slot A"
-            elif args.target == "b":
+    elif target == Target.APP_B:
         if total_size > APP_B_MAX_SIZE:
-                    _LOGGER.error(
+            _LOGGER.error(f"provided App B larger than allowed {total_size} bytes")
-                        f"provided App B larger than allowed {total_size} bytes"
-                    )
             return -1
         context_str = "App Slot B"
-            _LOGGER.info(
+    _LOGGER.info(f"Flashing {context_str} with image {file_path} (size {total_size})")
-                f"Flashing {context_str} with image {file_path} (size {total_size})"
-            )
     for idx, segment in enumerate(loadable_segments):
         _LOGGER.info(
             f"Loadable section {idx} {segment.name} with offset {segment.offset:#08x} and "
             f"size {segment.size}"
         )
-            result = perform_flashing_algorithm(com_if, loadable_segments, verificator)
-            if result != 0:
-                return result
-            crc_and_app_size_postprocessing(
-                com_if, args.target, total_size, loadable_segments
-            )
-    com_if.close()
-    return 0
-
-
-def perform_flashing_algorithm(
-    com_if: ComInterface,
-    loadable_segments: List[LoadableSegment],
-    verificator: PusVerificator,
-) -> int:
-    # Perform the flashing algorithm.
-    for segment in loadable_segments:
-        segment_end = segment.offset + segment.size
-        current_addr = segment.offset
-        pos_in_segment = 0
-        while pos_in_segment < segment.size:
-            next_chunk_size = min(segment_end - current_addr, CHUNK_SIZE)
-            data = segment.data[pos_in_segment : pos_in_segment + next_chunk_size]
-            next_packet = pack_memory_write_command(current_addr, data)
-            _LOGGER.info(
-                f"Sending memory write command for address {current_addr:#08x} and data with "
-                f"length {len(data)}"
-            )
-            verificator.add_tc(next_packet)
-            com_if.send(next_packet.pack())
-            current_addr += next_chunk_size
-            pos_in_segment += next_chunk_size
-            start_time = time.time()
-            while True:
-                if time.time() - start_time > 1.0:
-                    _LOGGER.error("Timeout while waiting for reply")
-                    return -1
-                data_available = com_if.data_available(0.1)
-                done = False
-                if not data_available:
-                    continue
-                replies = com_if.receive()
-                for reply in replies:
-                    tm = PusTm.unpack(reply, 0)
-                    if tm.service != 1:
-                        continue
-                    service_1_tm = Service1Tm.from_tm(tm, UnpackParams(0))
-                    check_result = verificator.add_tm(service_1_tm)
-                    # We could send after we have received the step reply, but that can
-                    # somehow lead to overrun errors. I think it's okay to do it like
-                    # this as long as the flash loader only uses polling..
-                    if (
-                        check_result is not None
-                        and check_result.status.completed == StatusField.SUCCESS
-                    ):
-                        done = True
-
-                    # This is an optimized variant, but I think the small delay is not an issue.
-                    """
-                    if (
-                        check_result is not None
-                        and check_result.status.step == StatusField.SUCCESS
-                        and len(check_result.status.step_list) == 1
-                    ):
-                        done = True
-                    """
-                verificator.remove_completed_entries()
-                if done:
-                    break
-    return 0
-
-
-def crc_and_app_size_postprocessing(
-    com_if: ComInterface,
-    target: str,
-    total_size: int,
-    loadable_segments: List[LoadableSegment],
-):
-    if target == "bl":
-        _LOGGER.info("Blanking the bootloader checksum")
-        # Blank the checksum. For the bootloader, the bootloader will calculate the
-        # checksum itself on the initial run.
-        checksum_write_packet = pack_memory_write_command(
-            BOOTLOADER_CRC_ADDR, bytes([0x00, 0x00, 0x00, 0x00])
-        )
-        com_if.send(checksum_write_packet.pack())
-    else:
-        crc_addr = None
-        size_addr = None
-        if target == "a":
-            crc_addr = APP_A_CRC_ADDR
-            size_addr = APP_A_SIZE_ADDR
-        elif target == "b":
-            crc_addr = APP_B_CRC_ADDR
-            size_addr = APP_B_SIZE_ADDR
-        assert crc_addr is not None
-        assert size_addr is not None
-        _LOGGER.info(f"Writing app size {total_size} at address {size_addr:#08x}")
-        size_write_packet = pack_memory_write_command(
-            size_addr, struct.pack("!I", total_size)
-        )
-        com_if.send(size_write_packet.pack())
-        time.sleep(0.2)
-        crc_calc = PredefinedCrc("crc-32")
-        for segment in loadable_segments:
-            crc_calc.update(segment.data)
-        checksum = crc_calc.digest()
-        _LOGGER.info(
-            f"Writing checksum 0x[{checksum.hex(sep=',')}] at address {crc_addr:#08x}"
-        )
-        checksum_write_packet = pack_memory_write_command(crc_addr, checksum)
-        com_if.send(checksum_write_packet.pack())
 
 
 def pack_memory_write_command(addr: int, data: bytes) -> PusTc:
@@ -381,7 +420,7 @@ def pack_memory_write_command(addr: int, data: bytes) -> PusTc:
         service=MEMORY_SERVICE,
         subservice=RAW_MEMORY_WRITE_SUBSERVICE,
         seq_count=SEQ_PROVIDER.get_and_increment(),
-        app_data=app_data,
+        app_data=bytes(app_data),
     )