From 749d1adaf55f574cbbc241ca9a10c56cc34b5bd0 Mon Sep 17 00:00:00 2001
From: Jan Beran <jan.beran@espressif.com>
Date: Fri, 20 Dec 2024 10:50:18 +0100
Subject: [PATCH] feat: print usb mode when output chip info

---
 esptool/__init__.py        | 11 +++++++----
 esptool/loader.py          | 31 ++++++++++++++++++++++++++++---
 esptool/targets/esp32h2.py |  3 +++
 test/test_esptool.py       | 16 ++++++++++++++++
 4 files changed, 54 insertions(+), 7 deletions(-)

diff --git a/esptool/__init__.py b/esptool/__init__.py
index 270c7c46f..1d786f374 100644
--- a/esptool/__init__.py
+++ b/esptool/__init__.py
@@ -829,11 +829,14 @@ def add_spi_flash_subparsers(
             )
 
         if esp.secure_download_mode:
-            print("Chip is %s in Secure Download Mode" % esp.CHIP_NAME)
+            print(f"Chip is {esp.CHIP_NAME} in Secure Download Mode")
         else:
-            print("Chip is %s" % (esp.get_chip_description()))
-            print("Features: %s" % ", ".join(esp.get_chip_features()))
-            print("Crystal is %dMHz" % esp.get_crystal_freq())
+            print(f"Chip is {esp.get_chip_description()}")
+            print(f"Features: {', '.join(esp.get_chip_features())}")
+            print(f"Crystal is {esp.get_crystal_freq()}MHz")
+            usb_mode = esp.get_usb_mode()
+            if usb_mode is not None:
+                print(f"USB mode: {usb_mode}")
             read_mac(esp, args)
 
         if not args.no_stub:
diff --git a/esptool/loader.py b/esptool/loader.py
index fabb5d82d..ca018c5ad 100644
--- a/esptool/loader.py
+++ b/esptool/loader.py
@@ -1047,9 +1047,34 @@ def get_uart_no(self):
         """
         Read the UARTDEV_BUF_NO register to get the number of the currently used console
         """
-        if self.cache["uart_no"] is None:
-            self.cache["uart_no"] = self.read_reg(self.UARTDEV_BUF_NO) & 0xFF
-        return self.cache["uart_no"]
+        # Some ESP chips do not have this register
+        try:
+            if self.cache["uart_no"] is None:
+                self.cache["uart_no"] = self.read_reg(self.UARTDEV_BUF_NO) & 0xFF
+            return self.cache["uart_no"]
+        except AttributeError:
+            return None
+
+    def uses_usb_jtag_serial(self):
+        """
+        Check if the chip uses USB JTAG/SERIAL mode.
+        """
+        return False
+
+    def uses_usb_otg(self):
+        """
+        Check if the chip uses USB OTG mode.
+        """
+        return False
+
+    def get_usb_mode(self):
+        """
+        Get the USB mode of the chip: USB-Serial/JTAG or USB-OTG. If the usb_mode is None, external USB-UART is used.
+        """
+        usb_jtag_serial = self.uses_usb_jtag_serial()
+        usb_otg = self.uses_usb_otg()
+
+        return "USB-Serial/JTAG" if usb_jtag_serial else "USB-OTG" if usb_otg else None
 
     @classmethod
     def parse_flash_size_arg(cls, arg):
diff --git a/esptool/targets/esp32h2.py b/esptool/targets/esp32h2.py
index be42a7b64..8857cdd46 100644
--- a/esptool/targets/esp32h2.py
+++ b/esptool/targets/esp32h2.py
@@ -24,6 +24,9 @@ class ESP32H2ROM(ESP32C6ROM):
     RTC_CNTL_SWD_WPROTECT_REG = DR_REG_LP_WDT_BASE + 0x0024  # LP_WDT_SWD_WPROTECT_REG
     RTC_CNTL_SWD_WKEY = 0x50D83AA1  # LP_WDT_SWD_WKEY, same as WDT key in this case
 
+    UARTDEV_BUF_NO = 0x4084FEFC  # Variable in ROM .bss which indicates the port in use
+    UARTDEV_BUF_NO_USB_JTAG_SERIAL = 3  # The above var when USB-JTAG/Serial is used
+
     FLASH_FREQUENCY = {
         "48m": 0xF,
         "24m": 0x0,
diff --git a/test/test_esptool.py b/test/test_esptool.py
index fdc9034e3..d5acf00cf 100755
--- a/test/test_esptool.py
+++ b/test/test_esptool.py
@@ -1353,6 +1353,22 @@ def test_auto_detect(self):
         self._check_output(output)
 
 
+class TestUSBMode(EsptoolTestCase):
+    @pytest.mark.quick_test
+    def test_usb_mode(self):
+        output = self.run_esptool("chip_id")
+        expected_usb_mode = (
+            "USB-OTG"
+            if os.environ.get("ESPTOOL_TEST_USB_OTG") == "1"
+            else "USB-Serial/JTAG"
+            if arg_preload_port
+            else None
+        )
+
+        if expected_usb_mode:
+            assert f"USB mode: {expected_usb_mode}" in output
+
+
 @pytest.mark.flaky(reruns=5)
 @pytest.mark.skipif(arg_preload_port is not False, reason="USB-to-UART bridge only")
 @pytest.mark.skipif(os.name == "nt", reason="Linux/MacOS only")