From b89f9030a0ae1a49f1d7d284de88163b9c6bcd33 Mon Sep 17 00:00:00 2001
From: Alan Rosenthal <alanrosenthal@google.com>
Date: Wed, 23 Nov 2022 13:06:31 -0500
Subject: [PATCH] Added class CharacteristicDeclaration, gatt_server getters

* Converted CharacteristicDeclaration implementation to class
* Added ability to get a gatt_server attribute by service UUID, characteristics UUID, descriptor UUID
---
 bumble/gatt.py        | 17 +++++++++++
 bumble/gatt_server.py | 71 +++++++++++++++++++++++++++++++++++++------
 tests/gatt_test.py    | 32 +++++++++++++++++++
 3 files changed, 110 insertions(+), 10 deletions(-)

diff --git a/bumble/gatt.py b/bumble/gatt.py
index 889eaa44..d2aaf571 100644
--- a/bumble/gatt.py
+++ b/bumble/gatt.py
@@ -271,6 +271,23 @@ class Characteristic(Attribute):
         return f'Characteristic(handle=0x{self.handle:04X}, end=0x{self.end_group_handle:04X}, uuid={self.uuid}, properties={Characteristic.properties_as_string(self.properties)})'
 
 
+# -----------------------------------------------------------------------------
+class CharacteristicDeclaration(Attribute):
+    '''
+    See Vol 3, Part G - 3.3.1 CHARACTERISTIC DECLARATION
+    '''
+    def __init__(self, characteristic, value_handle):
+        declaration_bytes = struct.pack(
+            '<BH',
+            characteristic.properties,
+            value_handle
+        ) + characteristic.uuid.to_pdu_bytes()
+        super().__init__(GATT_CHARACTERISTIC_ATTRIBUTE_TYPE, Attribute.READABLE, declaration_bytes)
+        self.characteristic = characteristic
+
+    def __str__(self):
+        return f'CharacteristicDeclaration(handle=0x{self.handle:04X}, value_handle=0x{self.value_handle:04X}, uuid={self.uuid}, properties={Characteristic.properties_as_string(self.properties)})'
+
 # -----------------------------------------------------------------------------
 class CharacteristicValue:
     '''
diff --git a/bumble/gatt_server.py b/bumble/gatt_server.py
index 5b6fea31..8580ea80 100644
--- a/bumble/gatt_server.py
+++ b/bumble/gatt_server.py
@@ -26,6 +26,7 @@
 import asyncio
 import logging
 from collections import defaultdict
+from typing import Tuple, Optional
 from pyee import EventEmitter
 from colors import color
 
@@ -79,6 +80,63 @@ class Server(EventEmitter):
                 return attribute
         return None
 
+    def get_service_attribute(self, service_uuid: UUID) -> Optional[Service]:
+        return next(
+            (
+                attribute
+                for attribute in self.attributes
+                if attribute.type == GATT_PRIMARY_SERVICE_ATTRIBUTE_TYPE
+                and attribute.uuid == service_uuid
+            ),
+            None,
+        )
+
+    def get_characteristic_attributes(
+        self, service_uuid: UUID, characteristic_uuid: UUID
+    ) -> Optional[Tuple[CharacteristicDeclaration, Characteristic]]:
+        service_handle = self.get_service_attribute(service_uuid)
+        if not service_handle:
+            return None
+
+        return next(
+            (
+                (attribute, self.get_attribute(attribute.characteristic.handle))
+                for attribute in map(
+                    self.get_attribute,
+                    range(service_handle.handle, service_handle.end_group_handle + 1),
+                )
+                if attribute.type == GATT_CHARACTERISTIC_ATTRIBUTE_TYPE
+                and attribute.characteristic.uuid == characteristic_uuid
+            ),
+            None,
+        )
+
+    def get_descriptor_attribute(
+        self, service_uuid: UUID, characteristic_uuid: UUID, descriptor_uuid: UUID
+    ) -> Optional[Descriptor]:
+        characteristics = self.get_characteristic_attributes(
+            service_uuid, characteristic_uuid
+        )
+        if not characteristics:
+            return None
+
+        (_, characteristic_value) = characteristics
+
+        return next(
+            (
+                attribute
+                for attribute in map(
+                    self.get_attribute,
+                    range(
+                        characteristic_value.handle + 1,
+                        characteristic_value.end_group_handle + 1,
+                    ),
+                )
+                if attribute.type == descriptor_uuid
+            ),
+            None,
+        )
+
     def add_attribute(self, attribute):
         # Assign a handle to this attribute
         attribute.handle = self.next_handle()
@@ -95,16 +153,9 @@ class Server(EventEmitter):
 
         # Add all characteristics
         for characteristic in service.characteristics:
-            # Add a Characteristic Declaration (Vol 3, Part G - 3.3.1 Characteristic Declaration)
-            declaration_bytes = struct.pack(
-                '<BH',
-                characteristic.properties,
-                self.next_handle() + 1,  # The value will be the next attribute after this declaration
-            ) + characteristic.uuid.to_pdu_bytes()
-            characteristic_declaration = Attribute(
-                GATT_CHARACTERISTIC_ATTRIBUTE_TYPE,
-                Attribute.READABLE,
-                declaration_bytes
+            # Add a Characteristic Declaration
+            characteristic_declaration = CharacteristicDeclaration(
+                characteristic, self.next_handle() + 1
             )
             self.add_attribute(characteristic_declaration)
 
diff --git a/tests/gatt_test.py b/tests/gatt_test.py
index 927aeee3..8adbfb48 100644
--- a/tests/gatt_test.py
+++ b/tests/gatt_test.py
@@ -28,6 +28,7 @@ from bumble.device import Device, Peer
 from bumble.host import Host
 from bumble.gatt import (
     GATT_BATTERY_LEVEL_CHARACTERISTIC,
+    GATT_CLIENT_CHARACTERISTIC_CONFIGURATION_DESCRIPTOR,
     CharacteristicAdapter,
     DelegatedCharacteristicAdapter,
     PackedCharacteristicAdapter,
@@ -226,6 +227,37 @@ async def test_characteristic_encoding():
     assert last_change is None
 
 
+# -----------------------------------------------------------------------------
+@pytest.mark.asyncio
+async def test_attribute_getters():
+    [client, server] = LinkedDevices().devices[:2]
+
+    characteristic_uuid = UUID('FDB159DB-036C-49E3-B3DB-6325AC750806')
+    characteristic = Characteristic(
+        characteristic_uuid,
+        Characteristic.READ | Characteristic.WRITE | Characteristic.NOTIFY,
+        Characteristic.READABLE | Characteristic.WRITEABLE,
+        bytes([123])
+    )
+
+    service_uuid = UUID('3A657F47-D34F-46B3-B1EC-698E29B6B829')
+    service = Service(service_uuid, [characteristic])
+    server.add_service(service)
+
+    service_attr = server.gatt_server.get_service_attribute(service_uuid)
+    assert service_attr
+
+    (char_decl_attr, char_value_attr) = server.gatt_server.get_characteristic_attributes(service_uuid, characteristic_uuid)
+    assert char_decl_attr and char_value_attr
+
+    desc_attr = server.gatt_server.get_descriptor_attribute(service_uuid, characteristic_uuid, GATT_CLIENT_CHARACTERISTIC_CONFIGURATION_DESCRIPTOR)
+    assert desc_attr
+
+    # assert all handles are in expected order
+    assert service_attr.handle < char_decl_attr.handle < char_value_attr.handle < desc_attr.handle == service_attr.end_group_handle
+    # assert characteristic declarations attribute is followed by characteristic value attribute
+    assert char_decl_attr.handle + 1 == char_value_attr.handle
+
 # -----------------------------------------------------------------------------
 def test_CharacteristicAdapter():
     # Check that the CharacteristicAdapter base class is transparent