From a9a68d164e9bdf58289efd6d45eba6f91b58c0af Mon Sep 17 00:00:00 2001
From: Mark Qvist <mark@unsigned.io>
Date: Mon, 30 Sep 2024 13:50:06 +0200
Subject: [PATCH] Only allow BLE modes with MITM mitigations enabled, and
 require bonding

---
 Bluetooth.h | 90 ++++++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 78 insertions(+), 12 deletions(-)

diff --git a/Bluetooth.h b/Bluetooth.h
index 7170934..0c05a2b 100644
--- a/Bluetooth.h
+++ b/Bluetooth.h
@@ -170,6 +170,8 @@ char bt_devname[11];
 
   #elif HAS_BLE == true
     BLESecurity *ble_security = new BLESecurity();
+    bool ble_authenticated = false;
+    uint32_t pairing_pin = 0;
 
     void bt_flush() { if (bt_state == BT_STATE_CONNECTED) { SerialBT.flush(); } }
 
@@ -181,6 +183,7 @@ char bt_devname[11];
     }
 
     void bt_passkey_notify_callback(uint32_t passkey) {
+      // Serial.printf("Got passkey notification: %d\n", passkey);
       bt_ssp_pin = passkey;
       bt_state = BT_STATE_PAIRING;
       bt_allow_pairing = true;
@@ -188,33 +191,100 @@ char bt_devname[11];
       kiss_indicate_btpin();
     }
 
+    bool bt_confirm_pin_callback(uint32_t pin) {
+      // Serial.printf("Confirm PIN callback: %d\n", pin);
+      return true;
+    }
+
+    void bt_debond_all() {
+      // Serial.println("Debonding all");
+      int dev_num = esp_ble_get_bond_device_num();
+      esp_ble_bond_dev_t *dev_list = (esp_ble_bond_dev_t *)malloc(sizeof(esp_ble_bond_dev_t) * dev_num);
+      esp_ble_get_bond_device_list(&dev_num, dev_list);
+      for (int i = 0; i < dev_num; i++) { esp_ble_remove_bond_device(dev_list[i].bd_addr); }
+      free(dev_list);
+    }
+
+    void bt_update_passkey() {
+      // Serial.println("Updating passkey");
+      pairing_pin = random(899999)+100000;
+      bt_ssp_pin = pairing_pin;
+    }
+
     uint32_t bt_passkey_callback() {
-      uint32_t pairing_pin = random(899999)+100000;
+      // Serial.println("API passkey request");
+      if (pairing_pin == 0) { bt_update_passkey(); }
       return pairing_pin;
     }
 
+    bool bt_client_authenticated() {
+      return ble_authenticated;
+    }
+
+    void bt_security_setup() {
+      uint32_t passkey = bt_passkey_callback();
+
+      // Serial.printf("Executing BT security setup, passkey is %d\n", passkey);
+
+      uint8_t key_size = 16;
+      uint8_t init_key = ESP_BLE_ENC_KEY_MASK | ESP_BLE_ID_KEY_MASK;
+      uint8_t rsp_key = ESP_BLE_ENC_KEY_MASK | ESP_BLE_ID_KEY_MASK;
+
+      esp_ble_auth_req_t auth_req = ESP_LE_AUTH_REQ_SC_MITM_BOND;
+      uint8_t auth_option = ESP_BLE_ONLY_ACCEPT_SPECIFIED_AUTH_ENABLE;
+      uint8_t oob_support = ESP_BLE_OOB_DISABLE;
+
+      esp_ble_io_cap_t iocap = ESP_IO_CAP_OUT;
+
+      esp_ble_gap_set_security_param(ESP_BLE_SM_SET_STATIC_PASSKEY, &passkey, sizeof(uint32_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_AUTHEN_REQ_MODE, &auth_req, sizeof(uint8_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_IOCAP_MODE, &iocap, sizeof(uint8_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_MAX_KEY_SIZE, &key_size, sizeof(uint8_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_ONLY_ACCEPT_SPECIFIED_SEC_AUTH, &auth_option, sizeof(uint8_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_OOB_SUPPORT, &oob_support, sizeof(uint8_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_SET_INIT_KEY, &init_key, sizeof(uint8_t));
+      esp_ble_gap_set_security_param(ESP_BLE_SM_SET_RSP_KEY, &rsp_key, sizeof(uint8_t));
+    }
+
     bool bt_security_request_callback() {
-      return true;
+      if (bt_allow_pairing) {
+          // Serial.println("Accepting security request");
+          return true;
+        } else {
+          // Serial.println("Rejecting security request");
+          return false;
+        }
     }
 
     void bt_authentication_complete_callback(esp_ble_auth_cmpl_t auth_result) {
       if (auth_result.success == true) {
+        // Serial.println("Authentication success");
+        ble_authenticated = true;
         bt_state = BT_STATE_CONNECTED;
       } else {
+        // Serial.println("Authentication fail");
+        ble_authenticated = false;
         bt_state = BT_STATE_ON;
+        bt_security_setup();
       }
       bt_allow_pairing = false;
       bt_ssp_pin = 0;
     }
 
-    void bt_connect_callback(uint16_t conn_handle) {
+    void bt_connect_callback(BLEServer *server) {
+      // uint16_t conn_id = server->getConnId();
+      // Serial.printf("Connected: %d\n", conn_id);
       display_unblank();
+      ble_authenticated = false;
       bt_state = BT_STATE_CONNECTED;
       cable_state = CABLE_STATE_DISCONNECTED;
     }
 
-    void bt_disconnect_callback(uint16_t conn_handle, uint8_t reason) {
+    void bt_disconnect_callback(BLEServer *server) {
+      // uint16_t conn_id = server->getConnId();
+      // Serial.printf("Disconnected: %d\n", conn_id);
       display_unblank();
+      ble_authenticated = false;
       bt_state = BT_STATE_ON;
     }
 
@@ -239,6 +309,8 @@ char bt_devname[11];
               sprintf(bt_devname, "RNode %02X%02X", bt_dh[14], bt_dh[15]);
               free(data);
 
+              bt_security_setup();
+
               bt_ready = true;
               return true;
 
@@ -280,14 +352,8 @@ char bt_devname[11];
       display_unblank();
       if (bt_state == BT_STATE_OFF) bt_start();
 
-      int dev_num = esp_ble_get_bond_device_num();
-      esp_ble_bond_dev_t *dev_list = (esp_ble_bond_dev_t *)malloc(sizeof(esp_ble_bond_dev_t) * dev_num);
-      esp_ble_get_bond_device_list(&dev_num, dev_list);
-      for (int i = 0; i < dev_num; i++) { esp_ble_remove_bond_device(dev_list[i].bd_addr); }
-      free(dev_list);
-
-      ble_security->setStaticPIN(bt_passkey_callback());
-      ble_security->setAuthenticationMode(ESP_LE_AUTH_REQ_SC_MITM_BOND);
+      bt_debond_all();
+      bt_update_passkey();
 
       bt_allow_pairing = true;
       bt_pairing_started = millis();