database: transition to get_device_shared_secret

Switch from `get_shared_secret` to `get_device_shared_secret`, which
requires additional authentication information already exposed by the
embedded devices.

Both `gateway` and `native_bt` required updates to run an explicit RPC
to get authentication information, instead of just reading the command
characteristic.

Signed-off-by: Jordan Yates <jordan@embeint.com>

Author: JordanYates

src/infuse_iot/database.py

@@ -11,8 +11,9 @@
 from cryptography.hazmat.primitives.asymmetric import x25519
 
 from infuse_iot.api_client import Client
-from infuse_iot.api_client.api.key import get_shared_secret
-from infuse_iot.api_client.models import Key
+from infuse_iot.api_client.api.key import get_device_shared_secret
+from infuse_iot.api_client.models.get_device_shared_secret_body import GetDeviceSharedSecretBody
+from infuse_iot.api_client.models.security_state import SecurityState
 from infuse_iot.credentials import get_api_key, load_network
 from infuse_iot.epacket.interface import Address as InterfaceAddress
 from infuse_iot.util.crypto import hkdf_derive
@@ -156,7 +157,14 @@ def _from_cache(self, infuse_id: int, device_pub_key: bytes) -> bytes | None:
             return state["shared_key"]
 
     def observe_security_state(
-        self, infuse_id: int, cloud_pub_key: bytes, device_pub_key: bytes, network_id: int
+        self,
+        infuse_id: int,
+        cloud_pub_key: bytes,
+        device_pub_key: bytes,
+        network_id: int,
+        challenge: bytes,
+        challenge_resp_type: int,
+        challenge_resp: bytes,
     ) -> None:
         """Update device state based on security_state response"""
         if infuse_id not in self.devices:
@@ -174,8 +182,16 @@ def observe_security_state(
 
         client = Client(base_url="https://api.infuse-iot.com").with_headers({"x-api-key": f"Bearer {get_api_key()}"})
         with client as client:
-            body = Key(base64.b64encode(device_pub_key).decode("utf-8"))
-            response = get_shared_secret.sync(client=client, body=body)
+            security_state = SecurityState(
+                base64.b64encode(cloud_pub_key).decode("utf-8"),
+                base64.b64encode(device_pub_key).decode("utf-8"),
+                network_id,
+                base64.b64encode(challenge).decode("utf-8"),
+                challenge_resp_type,
+                base64.b64encode(challenge_resp).decode("utf-8"),
+            )
+            body = GetDeviceSharedSecretBody(f"{infuse_id:016x}", security_state)
+            response = get_device_shared_secret.sync(client=client, body=body)
             if response is not None:
                 key = base64.b64decode(response.key)
                 self.devices[infuse_id].shared_key = key
@@ -208,6 +224,12 @@ def has_public_key(self, infuse_id: int) -> bool:
             return False
         return self.devices[infuse_id].device_public_key is not None
 
+    def has_shared_key(self, infuse_id: int) -> bool:
+        """Does the database have the shared key for this device?"""
+        if infuse_id not in self.devices:
+            return False
+        return self.devices[infuse_id].shared_key is not None
+
     def has_network_id(self, infuse_id: int) -> bool:
         """Does the database know the network ID for this device?"""
         if infuse_id not in self.devices:

src/infuse_iot/tools/gateway.py

@@ -75,13 +75,18 @@ def notification_broadcast(self, notification: ClientNotification):
         if self.server:
             self.server.broadcast(notification)
 
-    def query_device_key(self, cb_event: threading.Event | None = None):
+    def query_device_key(self, infuse_id: int, cb_event: threading.Event | None = None):
         def security_state_done(pkt: PacketReceived, _rc: int, response: bytes, challenge):
-            cloud_key = response[:32]
-            device_key = response[32:64]
-            network_id = int.from_bytes(response[64:68], "little")
-
-            self.ddb.observe_security_state(pkt.route[0].infuse_id, cloud_key, device_key, network_id)
+            decoded = defs.security_state.response.vla_from_buffer_copy(response)
+            self.ddb.observe_security_state(
+                infuse_id,
+                bytes(decoded.cloud_public_key),
+                bytes(decoded.device_public_key),
+                decoded.network_id,
+                challenge,
+                decoded.challenge_response_type,
+                bytes(decoded.challenge_response),
+            )
             if cb_event is not None:
                 cb_event.set()
 
@@ -91,7 +96,6 @@ def public_keys_done(pkt: PacketReceived, rc: int, response: bytes, _):
             decoded = defs.security_public_keys.response.vla_from_buffer_copy(response)
             for key in decoded.public_keys:
                 if key.id == defs.rpc_enum_key_id.SECONDARY_REMOTE_PUBLIC_KEY:
-                    infuse_id = pkt.route[0].infuse_id
                     self.ddb.observe_secondary_remote_public_key(infuse_id, bytes(key.key))
 
         def run_cmd_pkt(cmd_pkt: PacketOutputRouted):
@@ -105,15 +109,15 @@ def run_cmd_pkt(cmd_pkt: PacketOutputRouted):
 
         # Run security_state RPC
         challenge = random.randbytes(16)
-        cmd_pkt = self.rpc.generate(
-            defs.security_state.COMMAND_ID, challenge, Auth.NETWORK, security_state_done, challenge
+        cmd_pkt = self.rpc.generate_addressed(
+            infuse_id, defs.security_state.COMMAND_ID, challenge, Auth.NETWORK, security_state_done, challenge
         )
         run_cmd_pkt(cmd_pkt)
 
         if self.ddb.has_local_root:
             # Query other public keys from the device
-            cmd_pkt = self.rpc.generate(
-                defs.security_public_keys.COMMAND_ID, b"\x00", Auth.NETWORK, public_keys_done, None
+            cmd_pkt = self.rpc.generate_addressed(
+                infuse_id, defs.security_public_keys.COMMAND_ID, b"\x00", Auth.NETWORK, public_keys_done, None
             )
             run_cmd_pkt(cmd_pkt)
 
@@ -193,7 +197,7 @@ def _handle_serial_frame(self, frame: bytearray):
                     else:
                         Console.log_info(f"Dropping {len(frame)} byte packet...")
                 else:
-                    self._common.query_device_key(None)
+                    self._common.query_device_key(self._common.ddb.gateway, None)
                     Console.log_info(f"Dropping {len(frame)} byte packet to query device key...")
                 return
             except cryptography.exceptions.InvalidTag as e:
@@ -213,7 +217,8 @@ def _handle_serial_frame(self, frame: bytearray):
                     self._handle_memfault_pkt(pkt)
                 # Proactively requery keys
                 elif pkt.ptype == InfuseType.KEY_IDS:
-                    self._common.query_device_key(None)
+                    assert self._common.ddb.gateway is not None
+                    self._common.query_device_key(self._common.ddb.gateway, None)
 
                 # Forward to clients
                 notification = ClientNotificationEpacketReceived(pkt)
@@ -264,9 +269,9 @@ def _handle_epacket_send(self, req: GatewayRequestEpacketSend):
 
         # Do we have the device public keys we need?
         for hop in routed.route:
-            if hop.auth == Auth.DEVICE and not self._common.ddb.has_public_key(hop.infuse_id):
+            if hop.auth == Auth.DEVICE and not self._common.ddb.has_shared_key(hop.infuse_id):
                 cb_event = threading.Event()
-                self._common.query_device_key(cb_event)
+                self._common.query_device_key(hop.infuse_id, cb_event)
 
         # Encode and encrypt payload
         encrypted = routed.to_serial(self._common.ddb)
@@ -353,7 +358,7 @@ def _handle_conn_request(self, req: GatewayRequestConnectionRequest):
             subs,
             0,
         )
-        cmd = self._common.rpc.generate(
+        cmd = self._common.rpc.generate_serial(
             defs.bt_connect_infuse.COMMAND_ID, bytes(connect_args), Auth.DEVICE, self._bt_connect_cb, None
         )
         encrypted = cmd.to_serial(self._common.ddb)
@@ -379,7 +384,9 @@ def _handle_conn_release(self, req: GatewayRequestConnectionRelease):
             self._connected.pop(req.infuse_id)
 
         disconnect_args = defs.bt_disconnect.request(state.bt_addr.to_rpc_struct())
-        cmd = self._common.rpc.generate(defs.bt_disconnect.COMMAND_ID, bytes(disconnect_args), Auth.DEVICE, None, None)
+        cmd = self._common.rpc.generate_serial(
+            defs.bt_disconnect.COMMAND_ID, bytes(disconnect_args), Auth.DEVICE, None, None
+        )
         encrypted = cmd.to_serial(self._common.ddb)
         Console.log_tx(cmd.ptype, len(encrypted))
         self._common.port.write(encrypted)

src/infuse_iot/tools/native_bt.py

@@ -9,6 +9,7 @@
 import asyncio
 import ctypes
 import json
+import random
 from typing import Any
 
 from bleak import BleakClient, BleakScanner
@@ -17,6 +18,7 @@
 from bleak.backends.scanner import AdvertisementData
 from cryptography.exceptions import InvalidTag
 
+import infuse_iot.definitions.rpc as defs
 from infuse_iot.commands import InfuseCommand
 from infuse_iot.common import InfuseBluetoothUUID, InfuseType
 from infuse_iot.database import DeviceDatabase, UnknownNetworkError
@@ -46,6 +48,7 @@
 )
 from infuse_iot.util.argparse import BtLeAddress, ValidFile, add_server_port_parser
 from infuse_iot.util.console import Console
+from infuse_iot.util.local_rpc_server import LocalRpcServer
 
 
 class InfuseGattReadResponse(ctypes.LittleEndianStructure):
@@ -66,6 +69,7 @@ def __init__(self, database: DeviceDatabase, server: LocalServer, bleak_mapping:
         self._mapping = bleak_mapping
         self._queues: dict[int, asyncio.Queue] = {}
         self._tasks: dict[int, asyncio.Task] = {}
+        self._rpc = LocalRpcServer(database, native_bt=True)
 
     def wrapped_broadcast(self, notifcation: ClientNotification):
         try:
@@ -98,35 +102,84 @@ def notification_handler(self, _characteristic: BleakGATTCharacteristic, data: b
             bytes(decr),
         )
         Console.log_rx(pkt.ptype, len(data))
+        # Handle any local RPC responses
+        self._rpc.handle(pkt)
+        # Forward to clients
         self.wrapped_broadcast(ClientNotificationEpacketReceived(pkt))
 
     async def create_connection_internal(
         self, request: GatewayRequestConnectionRequest, dev: BLEDevice, queue: asyncio.Queue
     ):
-        Console.log_info(f"{dev}: Initiating connection")
+        command_notify_enabled = False
+        Console.log_info(f"{request.infuse_id:016x}: Initiating connection")
         async with BleakClient(dev, timeout=request.timeout_ms / 1000) as client:
             # Modified from bleak example code
             if client._backend.__class__.__name__ == "BleakClientBlueZDBus":
                 await client._backend._acquire_mtu()  # type: ignore
 
-            security_info = await client.read_gatt_char(InfuseBluetoothUUID.COMMAND_CHAR)
-            resp = InfuseGattReadResponse.from_buffer_copy(security_info)
-            self._db.observe_security_state(
-                request.infuse_id,
-                bytes(resp.cloud_public_key),
-                bytes(resp.device_public_key),
-                resp.network_id,
-            )
+            Console.log_info(f"{request.infuse_id:016x}: Connected (MTU {client.mtu_size})")
+
+            have_shared_key = self._db.has_shared_key(request.infuse_id)
+            if have_shared_key:
+                # Read the current keys back to confirm they haven't changed
+                security_info = await client.read_gatt_char(InfuseBluetoothUUID.COMMAND_CHAR)
+                resp = InfuseGattReadResponse.from_buffer_copy(security_info)
+                key_id = self._db.get_device_key_id(bytes(resp.cloud_public_key), bytes(resp.device_public_key))
+                if self._db.devices[request.infuse_id].device_key_id != key_id:
+                    # Keys mismatch, invaidate the shared key
+                    Console.log_info(f"{dev}: Key mismatch, re-running derivation")
+                    have_shared_key = False
+
+            if not have_shared_key:
+                # Always need the command characteristic to get the response
+                await client.start_notify(InfuseBluetoothUUID.COMMAND_CHAR, self.notification_handler)
+                command_notify_enabled = True
+
+                security_state_received = asyncio.Event()
+
+                def security_state_done(pkt: PacketReceived, _rc: int, response: bytes, challenge):
+                    decoded = defs.security_state.response.vla_from_buffer_copy(response)
+                    self._db.observe_security_state(
+                        request.infuse_id,
+                        bytes(decoded.cloud_public_key),
+                        bytes(decoded.device_public_key),
+                        decoded.network_id,
+                        challenge,
+                        decoded.challenge_response_type,
+                        bytes(decoded.challenge_response),
+                    )
+                    security_state_received.set()
+
+                # Construct the Security State RPC command
+                challenge = random.randbytes(16)
+                ss_pkt = self._rpc.generate_addressed(
+                    request.infuse_id,
+                    defs.security_state.COMMAND_ID,
+                    challenge,
+                    Auth.NETWORK,
+                    security_state_done,
+                    challenge,
+                )
 
-            if request.data_types & request.DataType.COMMAND:
+                # Encrypt command and write to remote
+                encr = CtypeBtGattFrame.encrypt(self._db, request.infuse_id, ss_pkt.ptype, Auth.NETWORK, ss_pkt.payload)
+                Console.log_tx(ss_pkt.ptype, len(encr))
+                await client.write_gatt_char(InfuseBluetoothUUID.COMMAND_CHAR, encr, response=False)
+
+                # Wait for a response
+                await asyncio.wait_for(security_state_received.wait(), timeout=request.timeout_ms / 1000)
+
+                # Disable the command characteristic if not requested
+                if not (request.data_types & request.DataType.COMMAND):
+                    await client.stop_notify(InfuseBluetoothUUID.COMMAND_CHAR)
+
+            if (request.data_types & request.DataType.COMMAND) and not command_notify_enabled:
                 await client.start_notify(InfuseBluetoothUUID.COMMAND_CHAR, self.notification_handler)
             if request.data_types & request.DataType.DATA:
                 await client.start_notify(InfuseBluetoothUUID.DATA_CHAR, self.notification_handler)
             if request.data_types & request.DataType.LOGGING:
                 await client.start_notify(InfuseBluetoothUUID.LOGGING_CHAR, self.notification_handler)
 
-            Console.log_info(f"{dev}: Connected (MTU {client.mtu_size})")
-
             self.wrapped_broadcast(
                 ClientNotificationConnectionCreated(
                     request.infuse_id,

src/infuse_iot/util/local_rpc_server.py

@@ -26,12 +26,13 @@
 class LocalRpcServer:
     """Basic class supporting locally generated commands"""
 
-    def __init__(self, database: DeviceDatabase):
+    def __init__(self, database: DeviceDatabase, native_bt: bool = False):
         self._cnt = random.randint(0, 2**31)
         self._ddb = database
+        self._native_bt = native_bt
         self._queued: dict[int, tuple[RpcCallback | None, typing.Any]] = {}
 
-    def generate(
+    def generate_serial(
         self, command: int, args: bytes, auth: Auth, cb: RpcCallback | None, cb_ctx: typing.Any
     ) -> PacketOutputRouted:
         """Generate RPC packet from arguments"""
@@ -47,6 +48,20 @@ def generate(
         self._cnt += 1
         return cmd_pkt
 
+    def generate_remote_bt_direct(
+        self, remote: int, command: int, args: bytes, auth: Auth, cb: RpcCallback | None, cb_ctx: typing.Any
+    ) -> PacketOutputRouted:
+        """Generate RPC packet for Bluetooth remote from arguments, without intermediate hops"""
+        cmd_bytes = bytes(rpc.RequestHeader(self._cnt, command)) + args
+        cmd_pkt = PacketOutputRouted(
+            [HopOutput(remote, interface.ID.BT_CENTRAL, auth)],
+            InfuseType.RPC_CMD,
+            cmd_bytes,
+        )
+        self._queued[self._cnt] = (cb, cb_ctx)
+        self._cnt += 1
+        return cmd_pkt
+
     def generate_remote_bt(
         self, remote: int, command: int, args: bytes, auth: Auth, cb: RpcCallback | None, cb_ctx: typing.Any
     ) -> PacketOutputRouted:
@@ -64,6 +79,16 @@ def generate_remote_bt(
             cmd_bytes,
         )
 
+    def generate_addressed(
+        self, address: int, command: int, args: bytes, auth: Auth, cb: RpcCallback | None, cb_ctx: typing.Any
+    ) -> PacketOutputRouted:
+        """Generate RPC packet for explicit address"""
+        if self._native_bt:
+            return self.generate_remote_bt_direct(address, command, args, auth, cb, cb_ctx)
+        if address == self._ddb.gateway:
+            return self.generate_serial(command, args, auth, cb, cb_ctx)
+        return self.generate_remote_bt(address, command, args, auth, cb, cb_ctx)
+
     def handle(self, pkt: PacketReceived):
         """Handle received packets"""
         # Only care about RPC responses