Button/winDeployScripts/software/collector/bluez_adv.py

#!/usr/bin/env python3
"""BlueZ LE advertisement registration (custom AD payload via D-Bus)."""
from __future__ import annotations

import logging
from typing import Optional

from config import COMPANY_ID, LOCAL_NAME, SERVICE_UUID_16

log = logging.getLogger("collector.bluez_adv")

dbus = None
try:
    import dbus as _dbus
    import dbus.mainloop.glib
    import dbus.service
    from gi.repository import GLib  # type: ignore

    dbus = _dbus
except Exception:  # pragma: no cover
    GLib = None
    dbus = None

BLUEZ = "org.bluez"
ADAPTER_IFACE = "org.bluez.Adapter1"
LE_ADVERTISING_MANAGER = "org.bluez.LEAdvertisingManager1"
LE_ADVERTISEMENT = "org.bluez.LEAdvertisement1"
PROPS = "org.freedesktop.DBus.Properties"
OBJECT_MANAGER = "org.freedesktop.DBus.ObjectManager"


def _ensure_mainloop():
    if dbus is None:
        raise RuntimeError("python3-dbus / PyGObject required")
    dbus.mainloop.glib.DBusGMainLoop(set_as_default=True)


def find_adapter(bus) -> Optional[str]:
    obj = bus.get_object(BLUEZ, "/")
    mgr = dbus.Interface(obj, OBJECT_MANAGER)
    for path, ifaces in mgr.GetManagedObjects().items():
        if LE_ADVERTISING_MANAGER in ifaces:
            return str(path)
    return None


def _make_advertisement_cls():
    class Advertisement(dbus.service.Object):
        def __init__(self, bus, index: int, mfr_payload: bytes):
            self.path = f"/org/buttontask/collector/advertisement{index}"
            self.mfr_payload = mfr_payload
            dbus.service.Object.__init__(self, bus, self.path)

        @dbus.service.method(LE_ADVERTISEMENT, in_signature="", out_signature="")
        def Release(self):
            log.debug("advertisement released %s", self.path)

        def get_properties(self):
            mfr = {
                dbus.UInt16(COMPANY_ID): dbus.Array(self.mfr_payload, signature="y")
            }
            return {
                LE_ADVERTISEMENT: {
                    "Type": "peripheral",
                    "ServiceUUIDs": dbus.Array(
                        [f"{SERVICE_UUID_16:04x}"], signature="s"
                    ),
                    "ManufacturerData": dbus.Dictionary(mfr, signature="qv"),
                    "LocalName": LOCAL_NAME,
                    "IncludeTxPower": dbus.Boolean(False),
                }
            }

        @dbus.service.method(PROPS, in_signature="s", out_signature="a{sv}")
        def GetAll(self, interface):
            if interface != LE_ADVERTISEMENT:
                raise dbus.exceptions.DBusException(
                    "org.freedesktop.DBus.Error.InvalidArgs"
                )
            return self.get_properties()[LE_ADVERTISEMENT]

    return Advertisement


class BlueZAdvertiser:
    def __init__(self):
        _ensure_mainloop()
        self.bus = dbus.SystemBus()
        self.adapter_path = find_adapter(self.bus)
        if not self.adapter_path:
            raise RuntimeError("no BlueZ LEAdvertisingManager1 adapter")
        self._adv = None
        self._index = 0
        self._AdvCls = _make_advertisement_cls()
        adapter = self.bus.get_object(BLUEZ, self.adapter_path)
        self._manager = dbus.Interface(adapter, LE_ADVERTISING_MANAGER)
        try:
            props = dbus.Interface(adapter, PROPS)
            props.Set(ADAPTER_IFACE, "Powered", dbus.Boolean(True))
        except Exception as e:
            log.warning("adapter power: %s", e)

    def unregister(self) -> None:
        if self._adv is None:
            return
        try:
            self._manager.UnregisterAdvertisement(self._adv.path)
        except Exception as e:
            log.debug("unregister adv: %s", e)
        try:
            self._adv.remove_from_connection()
        except Exception:
            pass
        self._adv = None

    def register(self, ad_bytes: bytes) -> None:
        """Register advertisement. `ad_bytes` may be full AD or raw mfr payload."""
        mfr = _extract_mfr_payload(ad_bytes)
        self.unregister()
        self._index += 1
        self._adv = self._AdvCls(self.bus, self._index, mfr)
        self._manager.RegisterAdvertisement(
            self._adv.path, dbus.Dictionary({}, signature="sv")
        )
        log.info("advertisement registered (mfr %d bytes)", len(mfr))


def _extract_mfr_payload(ad_bytes: bytes) -> bytes:
    if len(ad_bytes) >= 4 and ad_bytes[1] == 0x03:
        rest = ad_bytes[4:]
        if len(rest) >= 4 and rest[1] == 0xFF:
            return bytes(rest[4:])
        return bytes(rest)
    return bytes(ad_bytes)


def read_device_rssi(bus, device_path: str) -> Optional[int]:
    if dbus is None:
        return None
    try:
        obj = bus.get_object(BLUEZ, device_path)
        props = dbus.Interface(obj, PROPS)
        return int(props.Get("org.bluez.Device1", "RSSI"))
    except Exception:
        return None


def find_device_path(bus, addr: str) -> Optional[str]:
    if dbus is None:
        return None
    addr_u = (addr or "").upper()
    try:
        obj = bus.get_object(BLUEZ, "/")
        mgr = dbus.Interface(obj, OBJECT_MANAGER)
        for path, ifaces in mgr.GetManagedObjects().items():
            dev = ifaces.get("org.bluez.Device1")
            if not dev:
                continue
            if str(dev.get("Address", "")).upper() == addr_u:
                return str(path)
    except Exception as e:
        log.debug("find_device: %s", e)
    return None


def disconnect_device(bus, addr: str) -> None:
    path = find_device_path(bus, addr)
    if not path:
        return
    try:
        obj = bus.get_object(BLUEZ, path)
        dbus.Interface(obj, "org.bluez.Device1").Disconnect()
    except Exception as e:
        log.warning("disconnect %s: %s", addr, e)


def list_connected_addresses(bus) -> list:
    if dbus is None:
        return []
    out = []
    try:
        obj = bus.get_object(BLUEZ, "/")
        mgr = dbus.Interface(obj, OBJECT_MANAGER)
        for _path, ifaces in mgr.GetManagedObjects().items():
            dev = ifaces.get("org.bluez.Device1")
            if not dev:
                continue
            if bool(dev.get("Connected", False)):
                out.append(str(dev.get("Address", "")).upper())
    except Exception as e:
        log.debug("list_connected: %s", e)
    return out