From c25f6f4fd3a3275049e77059c5ab1e5bd3684dd1 Mon Sep 17 00:00:00 2001
From: "Brian S. Stephan" <bss@incorporeal.org>
Date: Sat, 8 Jul 2023 23:48:47 -0500
Subject: [PATCH] add command for writing to the BOOTSEL Pico

---
 gp2040ce_bintools/pico.py | 35 +++++++++++++++++++++++++++++++++++
 tests/test_pico.py        | 18 ++++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/gp2040ce_bintools/pico.py b/gp2040ce_bintools/pico.py
index ed0f9e2..6d86892 100644
--- a/gp2040ce_bintools/pico.py
+++ b/gp2040ce_bintools/pico.py
@@ -27,6 +27,7 @@ PICO_COMMANDS = {
     'REBOOT': 0x2,
     'ERASE': 0x3,
     'READ': 0x4,
+    'WRITE': 0x5,
     'EXIT_XIP': 0x6,
 }
 
@@ -89,6 +90,7 @@ def erase(out_end: usb.core.Endpoint, in_end: usb.core.Endpoint, location: int,
         location: memory address of where to start erasing from
         size: number of bytes to erase
     """
+    logger.debug("clearing %s bytes starting at %s", size, hex(location))
     # set up the data
     pico_token = 1
     command_size = 8
@@ -173,3 +175,36 @@ def reboot(out_end: usb.core.Endpoint) -> None:
                               PICO_MAGIC, pico_token, PICO_COMMANDS['REBOOT'], command_size, transfer_len,
                               boot_start, boot_end, boot_delay_ms))
     # we don't even bother reading here because it may have already rebooted
+
+
+def write(out_end: usb.core.Endpoint, in_end: usb.core.Endpoint, location: int, content: bytes) -> None:
+    """Write content to a Pico in BOOTSEL, starting from the specified location.
+
+    This also prepares the USB device for writing, so it expects to be able to grab
+    exclusive access.
+
+    Args:
+        out_endpoint: the out direction USB endpoint to write to
+        in_endpoint: the in direction USB endpoint to read from
+        location: memory address of where to start reading from
+        content: the data to write
+    """
+    # set up the data
+    command_size = 8
+
+    exclusive_access(out_end, in_end, is_exclusive=True)
+    exit_xip(out_end, in_end)
+    erase(out_end, in_end, location, len(content))
+    exit_xip(out_end, in_end)
+    pico_token = 1
+    logger.debug("writing %s bytes to %s", len(content), location)
+    payload = struct.pack(PICOBOOT_CMD_STRUCT + PICOBOOT_CMD_READ_SUFFIX_STRUCT,
+                          PICO_MAGIC, pico_token, PICO_COMMANDS['WRITE'], command_size, len(content),
+                          location, len(content))
+    logger.debug("WRITE: %s", payload)
+    out_end.write(payload)
+    logger.debug("actually writing bytes now...")
+    out_end.write(content)
+    res = in_end.read(256)
+    logger.debug("res: %s", res)
+    exclusive_access(out_end, in_end, is_exclusive=False)
diff --git a/tests/test_pico.py b/tests/test_pico.py
index 481d076..caf3b54 100644
--- a/tests/test_pico.py
+++ b/tests/test_pico.py
@@ -126,3 +126,21 @@ def test_reboot():
 
     payload = struct.pack('<LLBBxxLLLL4x', 0x431fd10b, 1, 0x2, 12, 0, 0, 0x20042000, 500)
     end_out.write.assert_called_with(payload)
+
+
+def test_write():
+    """Test that we can write to a board in BOOTSEL mode."""
+    end_out, end_in = mock.MagicMock(), mock.MagicMock()
+    _ = pico.write(end_out, end_in, 0x101FE000, b'\x00\x01\x02\x03')
+
+    expected_writes = [
+        mock.call(struct.pack('<LLBBxxLL12x', 0x431fd10b, 1, 0x1, 1, 0, 1)),
+        mock.call(struct.pack('<LLBBxxL16x', 0x431fd10b, 1, 0x6, 0, 0)),
+        mock.call(struct.pack('<LLBBxxLLL8x', 0x431fd10b, 1, 0x3, 8, 0, 0x101FE000, 4)),
+        mock.call(struct.pack('<LLBBxxL16x', 0x431fd10b, 1, 0x6, 0, 0)),
+        mock.call(struct.pack('<LLBBxxLLL8x', 0x431fd10b, 1, 0x5, 8, 4, 0x101FE000, 4)),
+        mock.call(b'\x00\x01\x02\x03'),
+        mock.call(struct.pack('<LLBBxxLL12x', 0x431fd10b, 1, 0x1, 1, 0, 0)),
+    ]
+    end_out.write.assert_has_calls(expected_writes)
+    assert end_in.read.call_count == 6