From d0b814fd97c8a5beef5a948c8606abbfe0596ea6 Mon Sep 17 00:00:00 2001
From: Martin Milata <martin@martinmilata.cz>
Date: Wed, 9 Apr 2025 01:20:13 +0200
Subject: [PATCH] refactor(core): ble: expose connection state in micropython

[no changelog]
---
 .../upymod/modtrezorio/modtrezorio-ble.h      | 25 +++++++++++++------
 core/mocks/generated/trezorio/ble.pyi         | 10 ++++++--
 core/src/apps/homescreen/device_menu.py       |  3 ++-
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/core/embed/upymod/modtrezorio/modtrezorio-ble.h b/core/embed/upymod/modtrezorio/modtrezorio-ble.h
index 58e357631e..593ac56db1 100644
--- a/core/embed/upymod/modtrezorio/modtrezorio-ble.h
+++ b/core/embed/upymod/modtrezorio/modtrezorio-ble.h
@@ -180,17 +180,26 @@ STATIC mp_obj_t mod_trezorio_BLE_peer_count(void) {
 STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorio_BLE_peer_count_obj,
                                  mod_trezorio_BLE_peer_count);
 
-/// def is_connected() -> bool:
+/// def connection_state() -> int:
 ///     """
-///     TODO: this should really return struct or enum
+///     Returns current connection state as flags:
+///
+///     0x01 state known
+///     0x02 connectable
+///     0x04 connected
+///     0x08 pairing
+///     0x10 pairing request
 ///     """
-STATIC mp_obj_t mod_trezorio_BLE_is_connected(void) {
+STATIC mp_obj_t mod_trezorio_BLE_connection_state(void) {
   ble_state_t state;
   ble_get_state(&state);
-  return MP_OBJ_NEW_SMALL_INT(state.connected);
+  mp_int_t flags = (state.state_known << 0) | (state.connectable << 1) |
+                   (state.connected << 2) | (state.pairing << 3) |
+                   (state.pairing_requested << 4);
+  return MP_OBJ_NEW_SMALL_INT(flags);
 }
-STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorio_BLE_is_connected_obj,
-                                 mod_trezorio_BLE_is_connected);
+STATIC MP_DEFINE_CONST_FUN_OBJ_0(mod_trezorio_BLE_connection_state_obj,
+                                 mod_trezorio_BLE_connection_state);
 
 const size_t CODE_LEN = 6;
 static bool encode_pairing_code(mp_obj_t obj, uint8_t *outbuf) {
@@ -365,8 +374,8 @@ STATIC const mp_rom_map_elem_t mod_trezorio_BLE_globals_table[] = {
      MP_ROM_PTR(&mod_trezorio_BLE_disconnect_obj)},
     {MP_ROM_QSTR(MP_QSTR_peer_count),
      MP_ROM_PTR(&mod_trezorio_BLE_peer_count_obj)},
-    {MP_ROM_QSTR(MP_QSTR_is_connected),
-     MP_ROM_PTR(&mod_trezorio_BLE_is_connected_obj)},
+    {MP_ROM_QSTR(MP_QSTR_connection_state),
+     MP_ROM_PTR(&mod_trezorio_BLE_connection_state_obj)},
     {MP_ROM_QSTR(MP_QSTR_allow_pairing),
      MP_ROM_PTR(&mod_trezorio_BLE_allow_pairing_obj)},
     {MP_ROM_QSTR(MP_QSTR_reject_pairing),
diff --git a/core/mocks/generated/trezorio/ble.pyi b/core/mocks/generated/trezorio/ble.pyi
index 808f06b9a1..edae01d981 100644
--- a/core/mocks/generated/trezorio/ble.pyi
+++ b/core/mocks/generated/trezorio/ble.pyi
@@ -52,9 +52,15 @@ def peer_count() -> int:
 
 
 # upymod/modtrezorio/modtrezorio-ble.h
-def is_connected() -> bool:
+def connection_state() -> int:
     """
-    TODO: this should really return struct or enum
+    Returns current connection state as flags:
+
+    0x01 state known
+    0x02 connectable
+    0x04 connected
+    0x08 pairing
+    0x10 pairing request
     """
 
 
diff --git a/core/src/apps/homescreen/device_menu.py b/core/src/apps/homescreen/device_menu.py
index d92cf635c0..aa05669d7f 100644
--- a/core/src/apps/homescreen/device_menu.py
+++ b/core/src/apps/homescreen/device_menu.py
@@ -10,10 +10,11 @@ async def handle_device_menu() -> None:
     # MOCK DATA
     failed_backup = True
     battery_percentage = 22
-    paired_devices = ["Trezor Suite"] if ble.is_connected() else []
+    paired_devices = ["Trezor Suite"] if (ble.connection_state() & 0x04) else []
     # ###
     firmware_version = ".".join(map(str, utils.VERSION))
     device_name = storage.device.get_label() or "Trezor"
+    print(bin(ble.connection_state()))
 
     menu_result = await raise_if_not_confirmed(
         trezorui_api.show_device_menu(