Add IRK support to ble_rssi (#6422)

Co-authored-by: Jesse Hills <3060199+jesserockz@users.noreply.github.com>

esphome/components/ble_presence/ble_presence_device.h

@@ -6,16 +6,6 @@
 
 #ifdef USE_ESP32
 
-#ifdef USE_ARDUINO
-#include "mbedtls/aes.h"
-#include "mbedtls/base64.h"
-#endif
-
-#ifdef USE_ESP_IDF
-#define MBEDTLS_AES_ALT
-#include <aes_alt.h>
-#endif
-
 namespace esphome {
 namespace ble_presence {
 
@@ -72,7 +62,7 @@ class BLEPresenceDevice : public binary_sensor::BinarySensorInitiallyOff,
         }
         break;
       case MATCH_BY_IRK:
-        if (resolve_irk_(device.address_uint64(), this->irk_)) {
+        if (device.resolve_irk(this->irk_)) {
           this->set_found_(true);
           return true;
         }
@@ -142,43 +132,6 @@ class BLEPresenceDevice : public binary_sensor::BinarySensorInitiallyOff,
   bool check_ibeacon_minor_{false};
   bool check_minimum_rssi_{false};
 
-  bool resolve_irk_(uint64_t addr64, const uint8_t *irk) {
-    uint8_t ecb_key[16];
-    uint8_t ecb_plaintext[16];
-    uint8_t ecb_ciphertext[16];
-
-    memcpy(&ecb_key, irk, 16);
-    memset(&ecb_plaintext, 0, 16);
-
-    ecb_plaintext[13] = (addr64 >> 40) & 0xff;
-    ecb_plaintext[14] = (addr64 >> 32) & 0xff;
-    ecb_plaintext[15] = (addr64 >> 24) & 0xff;
-
-    mbedtls_aes_context ctx = {0, 0, {0}};
-    mbedtls_aes_init(&ctx);
-
-    if (mbedtls_aes_setkey_enc(&ctx, ecb_key, 128) != 0) {
-      mbedtls_aes_free(&ctx);
-      return false;
-    }
-
-    if (mbedtls_aes_crypt_ecb(&ctx,
-#ifdef USE_ARDUINO
-                              MBEDTLS_AES_ENCRYPT,
-#elif defined(USE_ESP_IDF)
-                              ESP_AES_ENCRYPT,
-#endif
-                              ecb_plaintext, ecb_ciphertext) != 0) {
-      mbedtls_aes_free(&ctx);
-      return false;
-    }
-
-    mbedtls_aes_free(&ctx);
-
-    return ecb_ciphertext[15] == (addr64 & 0xff) && ecb_ciphertext[14] == ((addr64 >> 8) & 0xff) &&
-           ecb_ciphertext[13] == ((addr64 >> 16) & 0xff);
-  }
-
   bool found_{false};
   uint32_t last_seen_{};
   uint32_t timeout_{};

esphome/components/ble_rssi/ble_rssi_sensor.h

@@ -15,6 +15,10 @@ class BLERSSISensor : public sensor::Sensor, public esp32_ble_tracker::ESPBTDevi
     this->match_by_ = MATCH_BY_MAC_ADDRESS;
     this->address_ = address;
   }
+  void set_irk(uint8_t *irk) {
+    this->match_by_ = MATCH_BY_IRK;
+    this->irk_ = irk;
+  }
   void set_service_uuid16(uint16_t uuid) {
     this->match_by_ = MATCH_BY_SERVICE_UUID;
     this->uuid_ = esp32_ble_tracker::ESPBTUUID::from_uint16(uuid);
@@ -53,6 +57,13 @@ class BLERSSISensor : public sensor::Sensor, public esp32_ble_tracker::ESPBTDevi
           return true;
         }
         break;
+      case MATCH_BY_IRK:
+        if (device.resolve_irk(this->irk_)) {
+          this->publish_state(device.get_rssi());
+          this->found_ = true;
+          return true;
+        }
+        break;
       case MATCH_BY_SERVICE_UUID:
         for (auto uuid : device.get_service_uuids()) {
           if (this->uuid_ == uuid) {
@@ -91,12 +102,13 @@ class BLERSSISensor : public sensor::Sensor, public esp32_ble_tracker::ESPBTDevi
   float get_setup_priority() const override { return setup_priority::DATA; }
 
  protected:
-  enum MatchType { MATCH_BY_MAC_ADDRESS, MATCH_BY_SERVICE_UUID, MATCH_BY_IBEACON_UUID };
+  enum MatchType { MATCH_BY_MAC_ADDRESS, MATCH_BY_IRK, MATCH_BY_SERVICE_UUID, MATCH_BY_IBEACON_UUID };
   MatchType match_by_;
 
   bool found_{false};
 
   uint64_t address_;
+  uint8_t *irk_;
 
   esp32_ble_tracker::ESPBTUUID uuid_;
 

esphome/components/ble_rssi/sensor.py

@@ -12,6 +12,8 @@ from esphome.const import (
     UNIT_DECIBEL_MILLIWATT,
 )
 
+CONF_IRK = "irk"
+
 DEPENDENCIES = ["esp32_ble_tracker"]
 
 ble_rssi_ns = cg.esphome_ns.namespace("ble_rssi")
@@ -39,6 +41,7 @@ CONFIG_SCHEMA = cv.All(
     .extend(
         {
             cv.Optional(CONF_MAC_ADDRESS): cv.mac_address,
+            cv.Optional(CONF_IRK): cv.uuid,
             cv.Optional(CONF_SERVICE_UUID): esp32_ble_tracker.bt_uuid,
             cv.Optional(CONF_IBEACON_MAJOR): cv.uint16_t,
             cv.Optional(CONF_IBEACON_MINOR): cv.uint16_t,
@@ -47,7 +50,9 @@ CONFIG_SCHEMA = cv.All(
     )
     .extend(esp32_ble_tracker.ESP_BLE_DEVICE_SCHEMA)
     .extend(cv.COMPONENT_SCHEMA),
-    cv.has_exactly_one_key(CONF_MAC_ADDRESS, CONF_SERVICE_UUID, CONF_IBEACON_UUID),
+    cv.has_exactly_one_key(
+        CONF_MAC_ADDRESS, CONF_IRK, CONF_SERVICE_UUID, CONF_IBEACON_UUID
+    ),
     _validate,
 )
 
@@ -60,6 +65,10 @@ async def to_code(config):
     if mac_address := config.get(CONF_MAC_ADDRESS):
         cg.add(var.set_address(mac_address.as_hex))
 
+    if irk := config.get(CONF_IRK):
+        irk = esp32_ble_tracker.as_hex_array(str(irk))
+        cg.add(var.set_irk(irk))
+
     if service_uuid := config.get(CONF_SERVICE_UUID):
         if len(service_uuid) == len(esp32_ble_tracker.bt_uuid16_format):
             cg.add(var.set_service_uuid16(esp32_ble_tracker.as_hex(service_uuid)))

esphome/components/esp32_ble_tracker/esp32_ble_tracker.cpp

@@ -25,6 +25,9 @@
 #include <esp32-hal-bt.h>
 #endif
 
+#define MBEDTLS_AES_ALT
+#include <aes_alt.h>
+
 // bt_trace.h
 #undef TAG
 
@@ -692,6 +695,39 @@ void ESP32BLETracker::print_bt_device_info(const ESPBTDevice &device) {
   }
 }
 
+bool ESPBTDevice::resolve_irk(const uint8_t *irk) const {
+  uint8_t ecb_key[16];
+  uint8_t ecb_plaintext[16];
+  uint8_t ecb_ciphertext[16];
+
+  uint64_t addr64 = esp32_ble::ble_addr_to_uint64(this->address_);
+
+  memcpy(&ecb_key, irk, 16);
+  memset(&ecb_plaintext, 0, 16);
+
+  ecb_plaintext[13] = (addr64 >> 40) & 0xff;
+  ecb_plaintext[14] = (addr64 >> 32) & 0xff;
+  ecb_plaintext[15] = (addr64 >> 24) & 0xff;
+
+  mbedtls_aes_context ctx = {0, 0, {0}};
+  mbedtls_aes_init(&ctx);
+
+  if (mbedtls_aes_setkey_enc(&ctx, ecb_key, 128) != 0) {
+    mbedtls_aes_free(&ctx);
+    return false;
+  }
+
+  if (mbedtls_aes_crypt_ecb(&ctx, ESP_AES_ENCRYPT, ecb_plaintext, ecb_ciphertext) != 0) {
+    mbedtls_aes_free(&ctx);
+    return false;
+  }
+
+  mbedtls_aes_free(&ctx);
+
+  return ecb_ciphertext[15] == (addr64 & 0xff) && ecb_ciphertext[14] == ((addr64 >> 8) & 0xff) &&
+         ecb_ciphertext[13] == ((addr64 >> 16) & 0xff);
+}
+
 }  // namespace esp32_ble_tracker
 }  // namespace esphome
 

esphome/components/esp32_ble_tracker/esp32_ble_tracker.h

@@ -86,6 +86,8 @@ class ESPBTDevice {
 
   const esp_ble_gap_cb_param_t::ble_scan_result_evt_param &get_scan_result() const { return scan_result_; }
 
+  bool resolve_irk(const uint8_t *irk) const;
+
   optional<ESPBLEiBeacon> get_ibeacon() const {
     for (auto &it : this->manufacturer_datas_) {
       auto res = ESPBLEiBeacon::from_manufacturer_data(it);

tests/components/ble_rssi/common.yaml

@@ -16,3 +16,6 @@ sensor:
   - platform: ble_rssi
     service_uuid: 11223344-5566-7788-99aa-bbccddeeff00
     name: BLE Test iBeacon UUID
+  - platform: ble_rssi
+    irk: 1234567890abcdef1234567890abcdef
+    name: "BLE Tracker with Identity Resolving Key"