[ambz] Use delayed marker message for UART data readout

ltchiptool/soc/ambz/util/ambzcode.py

@@ -21,6 +21,41 @@ def rom_console() -> bytes:
             b"\x05\x22\x00\x00"  # RtlConsolRom()
         )
 
+    @staticmethod
+    def print_greeting(delay: float, data: bytes) -> bytes:
+        # ldr r0, ms_delay
+        # ldr r3, DelayMs
+        # blx r3
+        # adr r0, message_data
+        # ldr r1, message_size
+        # ldr r3, xmodem_uart_putdata
+        # blx r3
+        # b next
+        # DelayMs: .word 0x346C+1
+        # xmodem_uart_putdata: .word 0xEC48+1
+        # ms_delay: .word 1000
+        # message_size: .word 16
+        # message_data: .word 0,0,0,0
+        # next:
+        ms_delay = int(delay * 1000)
+        message_size = len(data)
+        if message_size > 16:
+            raise ValueError("Message must be 16 bytes or shorter")
+        message_data = data.ljust(16, b"\x00")
+        return (
+            (
+                b"\x05\x48\x03\x4b"
+                b"\x98\x47\x06\xa0"
+                b"\x04\x49\x02\x4b"
+                b"\x98\x47\x0f\xe0"
+                b"\x6d\x34\x00\x00"  # DelayMs()
+                b"\x49\xec\x00\x00"  # xmodem_uart_putdata()
+            )
+            + inttole32(ms_delay)
+            + inttole32(message_size)
+            + message_data
+        )
+
     @staticmethod
     def download_mode() -> bytes:
         # movs r0, #2
@@ -166,14 +201,15 @@ def read_data_md5(address: int, length: int, offset: int = 0) -> bytes:
 
     @staticmethod
     def print_data(length: int, address: int = AMBZ_DATA_ADDRESS) -> bytes:
-        # ldr r0, read_data
-        # movs r1, #length
+        # ldr r0, address
+        # ldr r1, length
         # ldr r3, xmodem_uart_putdata
         # blx r3
         # b next
         # movs r0, r0
         # xmodem_uart_putdata: .word 0xEC48+1
-        # read_data: .word 0x10003000
+        # address: .word 0x10003000
+        # length: .word 4
         # next:
         return (
             (

ltchiptool/soc/ambz/util/ambztool.py

@@ -5,16 +5,17 @@
 from enum import IntEnum
 from hashlib import md5
 from io import BytesIO
-from logging import debug, warning
+from logging import debug, info, warning
 from time import sleep, time
 from typing import IO, Callable, Generator, Optional
 
 import click
 from xmodem import XMODEM
 
+from ltchiptool.util.cli import DevicePortParamType
 from ltchiptool.util.intbin import align_down, align_up, inttole16, inttole24, inttole32
 from ltchiptool.util.logging import LoggingHandler, verbose
-from ltchiptool.util.misc import retry_generator
+from ltchiptool.util.misc import retry_generator, sizeof
 from ltchiptool.util.serialtool import SerialToolBase
 from ltchiptool.util.streams import StreamHook
 
@@ -29,6 +30,7 @@
 AMBZ_DIAG_BAUDRATE = 115200
 AMBZ_FLASH_ADDRESS = 0x8000000
 AMBZ_RAM_ADDRESS = 0x10002000
+AMBZ_GREETING_TEXT = b"AmbZTool_Marker!"
 
 AMBZ_CHIP_TYPE = {
     0xE0: "RTL8710BL",  # ???
@@ -109,7 +111,6 @@ def read(self, io: IO[bytes], n: int) -> bytes:
 
 
 class AmbZTool(SerialToolBase):
-    crc_speed_bps: int = 1500000
     xm_send_code: Optional[bytes] = None
     xm_fake_ack: bool = False
 
@@ -120,6 +121,7 @@ def __init__(
         link_timeout: float = 10.0,
         read_timeout: float = 0.6,
         retry_count: int = 10,
+        quiet_timeout: float = 10.0,
     ):
         super().__init__(port, baudrate, link_timeout, read_timeout, retry_count)
         LoggingHandler.get().attach(logging.getLogger("xmodem.XMODEM"))
@@ -128,6 +130,7 @@ def __init__(
             putc=self.xm_putc,
             mode="xmodem1k",
         )
+        self.quiet_timeout = quiet_timeout
 
     #############################
     # Xmodem serial port access #
@@ -187,7 +190,7 @@ def link(self, disconnect: bool = False) -> None:
     def quiet_handshake(self) -> None:
         self.flush()
         self.push_timeout(0.1)
-        end = time() + self.link_timeout
+        end = time() + self.quiet_timeout
         while time() < end:
             self.write(ACK)
             # discard everything from Loud-Handshake
@@ -422,12 +425,14 @@ def ram_boot_read(
         prev_baudrate = self.s.baudrate
         self.change_baudrate(AMBZ_DIAG_BAUDRATE)
 
+        # find actual response using a marker message
+        # wait before printing to let previous bytes through
+        code = AmbZCode.print_greeting(delay=0.4, data=AMBZ_GREETING_TEXT) + code
         # go back into download mode after we're done
         code = code + AmbZCode.download_mode()
 
         # messages printed by the ROM
-        # DiagPrintf changes "\n" to "\n\r"
-        msg_pre = b"\rclose xModem Transfer ...\r\n\r"
+        msg_pre = AMBZ_GREETING_TEXT
         msg_post = b"UARTIMG_Download"
         # send RAM code, exit download mode (changes baudrate to 115200)
         self.ram_boot(code=code, callback=callback, keep_baudrate=True)
@@ -447,10 +452,10 @@ def ram_boot_read(
         if msg_pre in resp:
             resp = resp.partition(msg_pre)[2]
         elif msg_pre[-7:] in resp:
-            warning(f"Expected message not found: {resp!r}")
+            warning(f"Partial marker message found: {resp!r}")
             resp = resp.partition(msg_pre[-7:])[2]
         else:
-            raise RuntimeError(f"Response unreadable: {resp!r}")
+            raise RuntimeError(f"Marker message not found: {resp!r}")
 
         if msg_post in resp:
             resp = resp.partition(msg_post)[0]
@@ -467,8 +472,34 @@ def ram_boot_read(
 @click.command(
     help="AmebaZ flashing tool",
 )
-def cli():
-    raise NotImplementedError()
+@click.option(
+    "-d",
+    "--device",
+    help="Target device port (default: auto detect)",
+    type=DevicePortParamType(),
+    default=(),
+)
+def cli(device: str):
+    amb = AmbZTool(port=device, baudrate=AMBZ_ROM_BAUDRATE)
+    info("Linking...")
+    amb.link()
+
+    chip_info = amb.ram_boot_read(
+        AmbZCode.read_chip_id(offset=0)
+        + AmbZCode.read_flash_id(offset=1)
+        + AmbZCode.print_data(length=4)
+    )
+    info(f"Received chip info: {chip_info.hex()}")
+    chip_id = chip_info[0]
+    size_id = chip_info[3]
+    info("Chip type: " + AMBZ_CHIP_TYPE.get(chip_id, f"Unknown 0x{chip_id:02X}"))
+    if 0x14 <= size_id <= 0x19:
+        info("Flash size: " + sizeof(1 << size_id))
+    else:
+        warning(f"Couldn't process flash ID: got {chip_info!r}")
+
+    info("Disconnecting...")
+    amb.link(disconnect=True)
 
 
 if __name__ == "__main__":