bumble 0.0.211__py3-none-any.whl → 0.0.213__py3-none-any.whl

bumble/host.py

@@ -26,10 +26,7 @@ from typing import (
     Any,
     Awaitable,
     Callable,
-    Deque,
-    Dict,
     Optional,
-    Set,
     cast,
     TYPE_CHECKING,
 )
@@ -41,7 +38,6 @@ from bumble.snoop import Snooper
 from bumble import drivers
 from bumble import hci
 from bumble.core import (
-    PhysicalTransport,
     PhysicalTransport,
     ConnectionPHY,
     ConnectionParameters,
@@ -75,6 +71,11 @@ class DataPacketQueue(utils.EventEmitter):
 
     max_packet_size: int
 
+    class PerConnectionState:
+        def __init__(self) -> None:
+            self.in_flight = 0
+            self.drained = asyncio.Event()
+
     def __init__(
         self,
         max_packet_size: int,
@@ -85,11 +86,16 @@ class DataPacketQueue(utils.EventEmitter):
         self.max_packet_size = max_packet_size
         self.max_in_flight = max_in_flight
         self._in_flight = 0  # Total number of packets in flight across all connections
-        self._in_flight_per_connection: dict[int, int] = collections.defaultdict(
-            int
-        )  # Number of packets in flight per connection
+        self._connection_state: dict[int, DataPacketQueue.PerConnectionState] = (
+            collections.defaultdict(DataPacketQueue.PerConnectionState)
+        )
+        self._drained_per_connection: dict[int, asyncio.Event] = (
+            collections.defaultdict(asyncio.Event)
+        )
         self._send = send
-        self._packets: Deque[tuple[hci.HCI_Packet, int]] = collections.deque()
+        self._packets: collections.deque[tuple[hci.HCI_Packet, int]] = (
+            collections.deque()
+        )
         self._queued = 0
         self._completed = 0
 
@@ -137,36 +143,40 @@ class DataPacketQueue(utils.EventEmitter):
             self._completed += flushed_count
             self._packets = collections.deque(packets_to_keep)
 
-        if connection_handle in self._in_flight_per_connection:
-            in_flight = self._in_flight_per_connection[connection_handle]
+        if connection_state := self._connection_state.pop(connection_handle, None):
+            in_flight = connection_state.in_flight
             self._completed += in_flight
             self._in_flight -= in_flight
-            del self._in_flight_per_connection[connection_handle]
+            connection_state.drained.set()
 
     def _check_queue(self) -> None:
         while self._packets and self._in_flight < self.max_in_flight:
             packet, connection_handle = self._packets.pop()
             self._send(packet)
             self._in_flight += 1
-            self._in_flight_per_connection[connection_handle] += 1
+            connection_state = self._connection_state[connection_handle]
+            connection_state.in_flight += 1
+            connection_state.drained.clear()
 
     def on_packets_completed(self, packet_count: int, connection_handle: int) -> None:
         """Mark one or more packets associated with a connection as completed."""
-        if connection_handle not in self._in_flight_per_connection:
+        if connection_handle not in self._connection_state:
             logger.warning(
                 f'received completion for unknown connection {connection_handle}'
             )
             return
 
-        in_flight_for_connection = self._in_flight_per_connection[connection_handle]
-        if packet_count <= in_flight_for_connection:
-            self._in_flight_per_connection[connection_handle] -= packet_count
+        connection_state = self._connection_state[connection_handle]
+        if packet_count <= connection_state.in_flight:
+            connection_state.in_flight -= packet_count
         else:
             logger.warning(
                 f'{packet_count} completed for {connection_handle} '
-                f'but only {in_flight_for_connection} in flight'
+                f'but only {connection_state.in_flight} in flight'
             )
-            self._in_flight_per_connection[connection_handle] = 0
+            connection_state.in_flight = 0
+        if connection_state.in_flight == 0:
+            connection_state.drained.set()
 
         if packet_count <= self._in_flight:
             self._in_flight -= packet_count
@@ -181,6 +191,13 @@ class DataPacketQueue(utils.EventEmitter):
         self._check_queue()
         self.emit('flow')
 
+    async def drain(self, connection_handle: int) -> None:
+        """Wait until there are no pending packets for a connection."""
+        if not (connection_state := self._connection_state.get(connection_handle)):
+            raise ValueError('no such connection')
+
+        await connection_state.drained.wait()
+
 
 # -----------------------------------------------------------------------------
 class Connection:
@@ -234,16 +251,16 @@ class IsoLink:
 
 # -----------------------------------------------------------------------------
 class Host(utils.EventEmitter):
-    connections: Dict[int, Connection]
-    cis_links: Dict[int, IsoLink]
-    bis_links: Dict[int, IsoLink]
-    sco_links: Dict[int, ScoLink]
+    connections: dict[int, Connection]
+    cis_links: dict[int, IsoLink]
+    bis_links: dict[int, IsoLink]
+    sco_links: dict[int, ScoLink]
     bigs: dict[int, set[int]]
     acl_packet_queue: Optional[DataPacketQueue] = None
     le_acl_packet_queue: Optional[DataPacketQueue] = None
     iso_packet_queue: Optional[DataPacketQueue] = None
     hci_sink: Optional[TransportSink] = None
-    hci_metadata: Dict[str, Any]
+    hci_metadata: dict[str, Any]
     long_term_key_provider: Optional[
         Callable[[int, bytes, int], Awaitable[Optional[bytes]]]
     ]
@@ -435,6 +452,14 @@ class Host(utils.EventEmitter):
                 )
             )
         )
+        if self.supports_command(hci.HCI_SET_EVENT_MASK_PAGE_2_COMMAND):
+            await self.send_command(
+                hci.HCI_Set_Event_Mask_Page_2_Command(
+                    event_mask_page_2=hci.HCI_Set_Event_Mask_Page_2_Command.mask(
+                        [hci.HCI_ENCRYPTION_CHANGE_V2_EVENT]
+                    )
+                )
+            )
 
         if (
             self.local_version is not None
@@ -456,6 +481,7 @@ class Host(utils.EventEmitter):
                     hci.HCI_LE_READ_LOCAL_P_256_PUBLIC_KEY_COMPLETE_EVENT,
                     hci.HCI_LE_GENERATE_DHKEY_COMPLETE_EVENT,
                     hci.HCI_LE_ENHANCED_CONNECTION_COMPLETE_EVENT,
+                    hci.HCI_LE_ENHANCED_CONNECTION_COMPLETE_V2_EVENT,
                     hci.HCI_LE_DIRECTED_ADVERTISING_REPORT_EVENT,
                     hci.HCI_LE_PHY_UPDATE_COMPLETE_EVENT,
                     hci.HCI_LE_EXTENDED_ADVERTISING_REPORT_EVENT,
@@ -804,7 +830,7 @@ class Host(utils.EventEmitter):
         ) != 0
 
     @property
-    def supported_commands(self) -> Set[int]:
+    def supported_commands(self) -> set[int]:
         return set(
             op_code
             for op_code, mask in hci.HCI_SUPPORTED_COMMANDS_MASKS.items()
@@ -827,8 +853,8 @@ class Host(utils.EventEmitter):
     def on_packet(self, packet: bytes) -> None:
         try:
             hci_packet = hci.HCI_Packet.from_bytes(packet)
-        except Exception as error:
-            logger.warning(f'!!! error parsing packet from bytes: {error}')
+        except Exception:
+            logger.exception('!!! error parsing packet from bytes')
             return
 
         if self.ready or (
@@ -1118,11 +1144,19 @@ class Host(utils.EventEmitter):
         else:
             self.emit('connection_phy_update_failure', connection.handle, event.status)
 
-    def on_hci_le_advertising_report_event(self, event):
+    def on_hci_le_advertising_report_event(
+        self,
+        event: (
+            hci.HCI_LE_Advertising_Report_Event
+            | hci.HCI_LE_Extended_Advertising_Report_Event
+        ),
+    ):
         for report in event.reports:
             self.emit('advertising_report', report)
 
-    def on_hci_le_extended_advertising_report_event(self, event):
+    def on_hci_le_extended_advertising_report_event(
+        self, event: hci.HCI_LE_Extended_Advertising_Report_Event
+    ):
         self.on_hci_le_advertising_report_event(event)
 
     def on_hci_le_advertising_set_terminated_event(self, event):
@@ -1253,7 +1287,24 @@ class Host(utils.EventEmitter):
             self.cis_links[event.connection_handle] = IsoLink(
                 handle=event.connection_handle, packet_queue=self.iso_packet_queue
             )
-            self.emit('cis_establishment', event.connection_handle)
+            self.emit(
+                'cis_establishment',
+                event.connection_handle,
+                event.cig_sync_delay,
+                event.cis_sync_delay,
+                event.transport_latency_c_to_p,
+                event.transport_latency_p_to_c,
+                event.phy_c_to_p,
+                event.phy_p_to_c,
+                event.nse,
+                event.bn_c_to_p,
+                event.bn_p_to_c,
+                event.ft_c_to_p,
+                event.ft_p_to_c,
+                event.max_pdu_c_to_p,
+                event.max_pdu_p_to_c,
+                event.iso_interval,
+            )
         else:
             self.emit(
                 'cis_establishment_failure', event.connection_handle, event.status
@@ -1341,6 +1392,15 @@ class Host(utils.EventEmitter):
     def on_hci_synchronous_connection_changed_event(self, event):
         pass
 
+    def on_hci_mode_change_event(self, event: hci.HCI_Mode_Change_Event):
+        self.emit(
+            'mode_change',
+            event.connection_handle,
+            event.status,
+            event.current_mode,
+            event.interval,
+        )
+
     def on_hci_role_change_event(self, event):
         if event.status == hci.HCI_SUCCESS:
             logger.debug(
@@ -1356,6 +1416,10 @@ class Host(utils.EventEmitter):
             self.emit('role_change_failure', event.bd_addr, event.status)
 
     def on_hci_le_data_length_change_event(self, event):
+        if (connection := self.connections.get(event.connection_handle)) is None:
+            logger.warning('!!! DATA LENGTH CHANGE: unknown handle')
+            return
+
         self.emit(
             'connection_data_length_change',
             event.connection_handle,
@@ -1376,13 +1440,30 @@ class Host(utils.EventEmitter):
                 event.status,
             )
 
-    def on_hci_encryption_change_event(self, event):
+    def on_hci_encryption_change_event(self, event: hci.HCI_Encryption_Change_Event):
         # Notify the client
         if event.status == hci.HCI_SUCCESS:
             self.emit(
                 'connection_encryption_change',
                 event.connection_handle,
                 event.encryption_enabled,
+                0,
+            )
+        else:
+            self.emit(
+                'connection_encryption_failure', event.connection_handle, event.status
+            )
+
+    def on_hci_encryption_change_v2_event(
+        self, event: hci.HCI_Encryption_Change_V2_Event
+    ):
+        # Notify the client
+        if event.status == hci.HCI_SUCCESS:
+            self.emit(
+                'connection_encryption_change',
+                event.connection_handle,
+                event.encryption_enabled,
+                event.encryption_key_size,
             )
         else:
             self.emit(
@@ -1496,13 +1577,15 @@ class Host(utils.EventEmitter):
         self.emit('inquiry_complete')
 
     def on_hci_inquiry_result_with_rssi_event(self, event):
-        for response in event.responses:
+        for bd_addr, class_of_device, rssi in zip(
+            event.bd_addr, event.class_of_device, event.rssi
+        ):
             self.emit(
                 'inquiry_result',
-                response.bd_addr,
-                response.class_of_device,
+                bd_addr,
+                class_of_device,
                 b'',
-                response.rssi,
+                rssi,
             )
 
     def on_hci_extended_inquiry_result_event(self, event):

bumble/keys.py

@@ -22,14 +22,15 @@
 # -----------------------------------------------------------------------------
 from __future__ import annotations
 import asyncio
+import dataclasses
 import logging
 import os
 import json
-from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Optional, Any
 from typing_extensions import Self
 
 from bumble.colors import color
-from bumble.hci import Address
+from bumble import hci
 
 if TYPE_CHECKING:
     from bumble.device import Device
@@ -42,16 +43,17 @@ logger = logging.getLogger(__name__)
 
 
 # -----------------------------------------------------------------------------
+@dataclasses.dataclass
 class PairingKeys:
+    @dataclasses.dataclass
     class Key:
-        def __init__(self, value, authenticated=False, ediv=None, rand=None):
-            self.value = value
-            self.authenticated = authenticated
-            self.ediv = ediv
-            self.rand = rand
+        value: bytes
+        authenticated: bool = False
+        ediv: Optional[int] = None
+        rand: Optional[bytes] = None
 
         @classmethod
-        def from_dict(cls, key_dict):
+        def from_dict(cls, key_dict: dict[str, Any]) -> PairingKeys.Key:
             value = bytes.fromhex(key_dict['value'])
             authenticated = key_dict.get('authenticated', False)
             ediv = key_dict.get('ediv')
@@ -61,7 +63,7 @@ class PairingKeys:
 
             return cls(value, authenticated, ediv, rand)
 
-        def to_dict(self):
+        def to_dict(self) -> dict[str, Any]:
             key_dict = {'value': self.value.hex(), 'authenticated': self.authenticated}
             if self.ediv is not None:
                 key_dict['ediv'] = self.ediv
@@ -70,39 +72,42 @@ class PairingKeys:
 
             return key_dict
 
-    def __init__(self):
-        self.address_type = None
-        self.ltk = None
-        self.ltk_central = None
-        self.ltk_peripheral = None
-        self.irk = None
-        self.csrk = None
-        self.link_key = None  # Classic
-
-    @staticmethod
-    def key_from_dict(keys_dict, key_name):
+    address_type: Optional[hci.AddressType] = None
+    ltk: Optional[Key] = None
+    ltk_central: Optional[Key] = None
+    ltk_peripheral: Optional[Key] = None
+    irk: Optional[Key] = None
+    csrk: Optional[Key] = None
+    link_key: Optional[Key] = None  # Classic
+    link_key_type: Optional[int] = None  # Classic
+
+    @classmethod
+    def key_from_dict(cls, keys_dict: dict[str, Any], key_name: str) -> Optional[Key]:
         key_dict = keys_dict.get(key_name)
         if key_dict is None:
             return None
 
         return PairingKeys.Key.from_dict(key_dict)
 
-    @staticmethod
-    def from_dict(keys_dict):
-        keys = PairingKeys()
-
-        keys.address_type = keys_dict.get('address_type')
-        keys.ltk = PairingKeys.key_from_dict(keys_dict, 'ltk')
-        keys.ltk_central = PairingKeys.key_from_dict(keys_dict, 'ltk_central')
-        keys.ltk_peripheral = PairingKeys.key_from_dict(keys_dict, 'ltk_peripheral')
-        keys.irk = PairingKeys.key_from_dict(keys_dict, 'irk')
-        keys.csrk = PairingKeys.key_from_dict(keys_dict, 'csrk')
-        keys.link_key = PairingKeys.key_from_dict(keys_dict, 'link_key')
-
-        return keys
-
-    def to_dict(self):
-        keys = {}
+    @classmethod
+    def from_dict(cls, keys_dict: dict[str, Any]) -> PairingKeys:
+        return PairingKeys(
+            address_type=(
+                hci.AddressType(t)
+                if (t := keys_dict.get('address_type')) is not None
+                else None
+            ),
+            ltk=PairingKeys.key_from_dict(keys_dict, 'ltk'),
+            ltk_central=PairingKeys.key_from_dict(keys_dict, 'ltk_central'),
+            ltk_peripheral=PairingKeys.key_from_dict(keys_dict, 'ltk_peripheral'),
+            irk=PairingKeys.key_from_dict(keys_dict, 'irk'),
+            csrk=PairingKeys.key_from_dict(keys_dict, 'csrk'),
+            link_key=PairingKeys.key_from_dict(keys_dict, 'link_key'),
+            link_key_type=keys_dict.get('link_key_type'),
+        )
+
+    def to_dict(self) -> dict[str, Any]:
+        keys: dict[str, Any] = {}
 
         if self.address_type is not None:
             keys['address_type'] = self.address_type
@@ -125,9 +130,12 @@ class PairingKeys:
         if self.link_key is not None:
             keys['link_key'] = self.link_key.to_dict()
 
+        if self.link_key_type is not None:
+            keys['link_key_type'] = self.link_key_type
+
         return keys
 
-    def print(self, prefix=''):
+    def print(self, prefix: str = '') -> None:
         keys_dict = self.to_dict()
         for container_property, value in keys_dict.items():
             if isinstance(value, dict):
@@ -149,27 +157,35 @@ class KeyStore:
     async def get(self, _name: str) -> Optional[PairingKeys]:
         return None
 
-    async def get_all(self) -> List[Tuple[str, PairingKeys]]:
+    async def get_all(self) -> list[tuple[str, PairingKeys]]:
         return []
 
     async def delete_all(self) -> None:
         all_keys = await self.get_all()
         await asyncio.gather(*(self.delete(name) for (name, _) in all_keys))
 
-    async def get_resolving_keys(self):
+    async def get_resolving_keys(self) -> list[tuple[bytes, hci.Address]]:
         all_keys = await self.get_all()
         resolving_keys = []
         for name, keys in all_keys:
             if keys.irk is not None:
-                if keys.address_type is None:
-                    address_type = Address.RANDOM_DEVICE_ADDRESS
-                else:
-                    address_type = keys.address_type
-                resolving_keys.append((keys.irk.value, Address(name, address_type)))
+                resolving_keys.append(
+                    (
+                        keys.irk.value,
+                        hci.Address(
+                            name,
+                            (
+                                keys.address_type
+                                if keys.address_type is not None
+                                else hci.Address.RANDOM_DEVICE_ADDRESS
+                            ),
+                        ),
+                    )
+                )
 
         return resolving_keys
 
-    async def print(self, prefix=''):
+    async def print(self, prefix: str = '') -> None:
         entries = await self.get_all()
         separator = ''
         for name, keys in entries:
@@ -177,8 +193,8 @@ class KeyStore:
             keys.print(prefix=prefix + '  ')
             separator = '\n'
 
-    @staticmethod
-    def create_for_device(device: Device) -> KeyStore:
+    @classmethod
+    def create_for_device(cls, device: Device) -> KeyStore:
         if device.config.keystore is None:
             return MemoryKeyStore()
 
@@ -256,7 +272,7 @@ class JsonKeyStore(KeyStore):
 
     @classmethod
     def from_device(
-        cls: Type[Self], device: Device, filename: Optional[str] = None
+        cls: type[Self], device: Device, filename: Optional[str] = None
     ) -> Self:
         if not filename:
             # Extract the filename from the config if there is one
@@ -266,9 +282,9 @@ class JsonKeyStore(KeyStore):
                     filename = params[0]
 
         # Use a namespace based on the device address
-        if device.public_address not in (Address.ANY, Address.ANY_RANDOM):
+        if device.public_address not in (hci.Address.ANY, hci.Address.ANY_RANDOM):
             namespace = str(device.public_address)
-        elif device.random_address != Address.ANY_RANDOM:
+        elif device.random_address != hci.Address.ANY_RANDOM:
             namespace = str(device.random_address)
         else:
             namespace = JsonKeyStore.DEFAULT_NAMESPACE
@@ -340,7 +356,7 @@ class JsonKeyStore(KeyStore):
 
 # -----------------------------------------------------------------------------
 class MemoryKeyStore(KeyStore):
-    all_keys: Dict[str, PairingKeys]
+    all_keys: dict[str, PairingKeys]
 
     def __init__(self) -> None:
         self.all_keys = {}
@@ -355,5 +371,5 @@ class MemoryKeyStore(KeyStore):
     async def get(self, name: str) -> Optional[PairingKeys]:
         return self.all_keys.get(name)
 
-    async def get_all(self) -> List[Tuple[str, PairingKeys]]:
+    async def get_all(self) -> list[tuple[str, PairingKeys]]:
         return list(self.all_keys.items())