From f3c1d8767b798c75763e817b05c8e4deff64fd0c Mon Sep 17 00:00:00 2001
From: Tyler Potyondy <potyondyt@gmail.com>
Date: Thu, 9 May 2024 09:56:46 -0700
Subject: [PATCH] Add jlink-serial-number argument to specify target

-adds flag `--jlink-serial-number XXXX` to specify target board
---
 README.md              |  6 ++++++
 tockloader/jlinkexe.py | 21 +++++++++++++++++++--
 tockloader/main.py     |  5 +++++
 3 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 095ff0e..d35ef5e 100644
--- a/README.md
+++ b/README.md
@@ -200,6 +200,12 @@ knowing the device type of the MCU on the board.
 - `speed`: The speed value to pass to JLink. Defaults to 1200.
 - `if`: The interface to pass to JLink.
 
+JLink supports selecting the target board based on the serial number of the
+JTAG. Tockloader exposes this functionality with the `jlink-serial-number`
+flag. 
+
+   tockloader [command] --jlink-serial-number [SERIAL_NUMBER]
+
 Tockloader can also do JTAG using OpenOCD. OpenOCD needs to know which config
 file to use.
 
diff --git a/tockloader/jlinkexe.py b/tockloader/jlinkexe.py
index 09f30f1..0655d99 100644
--- a/tockloader/jlinkexe.py
+++ b/tockloader/jlinkexe.py
@@ -36,6 +36,9 @@ def __init__(self, args):
             if platform.system() == "Windows":
                 self.jlink_cmd = "JLink"
 
+        # Obtain serial number if --jlink-serial-number argument provided
+        self.jlink_serial_number = getattr(self.args, "jlink_serial_number")
+
         # By default we assume that jlinkexe can be used to read any address on
         # this board, so we set `address_maximum` to None. In some cases,
         # however, particularly with external flash chips, jlinkexe may not be
@@ -196,6 +199,10 @@ def _run_jtag_commands(self, commands, binary, write=True):
                 jlink_file.name,
             )
 
+            # Append target selector if serial number provided.
+            if self.jlink_serial_number:
+                jlink_command += " -USB {}".format(self.jlink_serial_number)
+
             logging.debug('Running "{}".'.format(jlink_command))
 
             def print_output(subp):
@@ -456,9 +463,19 @@ def run_terminal(self):
             return
 
         logging.status("Starting JLinkExe JTAG connection.")
+
+        # Include serial number specifier if `--jlin-serial-number` flag provided.
+        jlink_serial_number_str = ""
+        if self.jlink_serial_number:
+            jlink_serial_number_str = ("-USB {}").format(self.jlink_serial_number)
+
         jtag_p = subprocess.Popen(
-            "{} -device {} -if {} -speed {} -autoconnect 1 -jtagconf -1,-1".format(
-                self.jlink_cmd, self.jlink_device, self.jlink_if, self.jlink_speed
+            "{} -device {} -if {} -speed {} -autoconnect 1 {} --jtagconf -1,-1".format(
+                self.jlink_cmd,
+                self.jlink_device,
+                self.jlink_if,
+                self.jlink_speed,
+                jlink_serial_number_str,
             ).split(),
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
diff --git a/tockloader/main.py b/tockloader/main.py
index 46055af..9b56761 100644
--- a/tockloader/main.py
+++ b/tockloader/main.py
@@ -778,6 +778,11 @@ def main():
     parent_channel.add_argument(
         "--jlink-if", help="The interface type to pass to JLinkExe."
     )
+    parent_channel.add_argument(
+        "--jlink-serial-number",
+        default=None,
+        help="Specify a specific JLink via serial number. Useful when multiple JLinks connected to same machine.",
+    )
     parent_channel.add_argument(
         "--openocd-board", help="The cfg file in OpenOCD `board` folder."
     )