Add dynamic subscription for ZHA add device page (#22164)

* add ws subscription for zha gateway messages
* add debug mode
* only relay certain logs
* add missing require admin
* add devices command
* add area_id
* fix manufacturer code

homeassistant/components/zha/api.py

@@ -10,13 +10,15 @@ import logging
 import voluptuous as vol
 
 from homeassistant.components import websocket_api
+from homeassistant.core import callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.device_registry import async_get_registry
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from .core.const import (
     DOMAIN, ATTR_CLUSTER_ID, ATTR_CLUSTER_TYPE, ATTR_ATTRIBUTE, ATTR_VALUE,
     ATTR_MANUFACTURER, ATTR_COMMAND, ATTR_COMMAND_TYPE, ATTR_ARGS, IN, OUT,
     CLIENT_COMMANDS, SERVER_COMMANDS, SERVER, NAME, ATTR_ENDPOINT_ID,
-    DATA_ZHA_GATEWAY, DATA_ZHA)
+    DATA_ZHA_GATEWAY, DATA_ZHA, MFG_CLUSTER_ID_START)
 from .core.helpers import get_matched_clusters, async_is_bindable_target
 
 _LOGGER = logging.getLogger(__name__)
@@ -74,6 +76,38 @@ SERVICE_SCHEMAS = {
 }
 
 
+@websocket_api.require_admin
+@websocket_api.async_response
+@websocket_api.websocket_command({
+    vol.Required('type'): 'zha/devices/permit'
+})
+async def websocket_permit_devices(hass, connection, msg):
+    """Permit ZHA zigbee devices."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+
+    async def forward_messages(data):
+        """Forward events to websocket."""
+        connection.send_message(websocket_api.event_message(msg['id'], data))
+
+    remove_dispatcher_function = async_dispatcher_connect(
+        hass,
+        "zha_gateway_message",
+        forward_messages
+    )
+
+    @callback
+    def async_cleanup() -> None:
+        """Remove signal listener and turn off debug mode."""
+        zha_gateway.async_disable_debug_mode()
+        remove_dispatcher_function()
+
+    connection.subscriptions[msg['id']] = async_cleanup
+    zha_gateway.async_enable_debug_mode()
+    await zha_gateway.application_controller.permit(60)
+
+    connection.send_result(msg['id'])
+
+
 @websocket_api.require_admin
 @websocket_api.async_response
 @websocket_api.websocket_command({
@@ -86,22 +120,33 @@ async def websocket_get_devices(hass, connection, msg):
 
     devices = []
     for device in zha_gateway.devices.values():
-        ret_device = {}
+        devices.append(
-        ret_device.update(device.device_info)
+            async_get_device_info(
-        ret_device['entities'] = [{
+                hass, device, ha_device_registry=ha_device_registry
-            'entity_id': entity_ref.reference_id,
+            )
-            NAME: entity_ref.device_info[NAME]
+        )
-        } for entity_ref in zha_gateway.device_registry[device.ieee]]
+    connection.send_result(msg[ID], devices)
 
+
+@callback
+def async_get_device_info(hass, device, ha_device_registry=None):
+    """Get ZHA device."""
+    zha_gateway = hass.data[DATA_ZHA][DATA_ZHA_GATEWAY]
+    ret_device = {}
+    ret_device.update(device.device_info)
+    ret_device['entities'] = [{
+        'entity_id': entity_ref.reference_id,
+        NAME: entity_ref.device_info[NAME]
+    } for entity_ref in zha_gateway.device_registry[device.ieee]]
+
+    if ha_device_registry is not None:
         reg_device = ha_device_registry.async_get_device(
             {(DOMAIN, str(device.ieee))}, set())
         if reg_device is not None:
             ret_device['user_given_name'] = reg_device.name_by_user
             ret_device['device_reg_id'] = reg_device.id
-
+            ret_device['area_id'] = reg_device.area_id
-        devices.append(ret_device)
+    return ret_device
-
-    connection.send_result(msg[ID], devices)
 
 
 @websocket_api.require_admin
@@ -265,7 +310,10 @@ async def websocket_read_zigbee_cluster_attributes(hass, connection, msg):
     cluster_id = msg[ATTR_CLUSTER_ID]
     cluster_type = msg[ATTR_CLUSTER_TYPE]
     attribute = msg[ATTR_ATTRIBUTE]
-    manufacturer = msg.get(ATTR_MANUFACTURER) or None
+    manufacturer = None
+    #  only use manufacturer code for manufacturer clusters
+    if cluster_id >= MFG_CLUSTER_ID_START:
+        manufacturer = msg.get(ATTR_MANUFACTURER) or None
     zha_device = zha_gateway.get_device(ieee)
     success = failure = None
     if zha_device is not None:
@@ -428,7 +476,10 @@ def async_load_api(hass):
         cluster_type = service.data.get(ATTR_CLUSTER_TYPE)
         attribute = service.data.get(ATTR_ATTRIBUTE)
         value = service.data.get(ATTR_VALUE)
-        manufacturer = service.data.get(ATTR_MANUFACTURER) or None
+        manufacturer = None
+        #  only use manufacturer code for manufacturer clusters
+        if cluster_id >= MFG_CLUSTER_ID_START:
+            manufacturer = service.data.get(ATTR_MANUFACTURER) or None
         zha_device = zha_gateway.get_device(ieee)
         response = None
         if zha_device is not None:
@@ -466,7 +517,10 @@ def async_load_api(hass):
         command = service.data.get(ATTR_COMMAND)
         command_type = service.data.get(ATTR_COMMAND_TYPE)
         args = service.data.get(ATTR_ARGS)
-        manufacturer = service.data.get(ATTR_MANUFACTURER) or None
+        manufacturer = None
+        #  only use manufacturer code for manufacturer clusters
+        if cluster_id >= MFG_CLUSTER_ID_START:
+            manufacturer = service.data.get(ATTR_MANUFACTURER) or None
         zha_device = zha_gateway.get_device(ieee)
         response = None
         if zha_device is not None:
@@ -497,6 +551,7 @@ def async_load_api(hass):
             SERVICE_ISSUE_ZIGBEE_CLUSTER_COMMAND
         ])
 
+    websocket_api.async_register_command(hass, websocket_permit_devices)
     websocket_api.async_register_command(hass, websocket_get_devices)
     websocket_api.async_register_command(hass, websocket_reconfigure_node)
     websocket_api.async_register_command(hass, websocket_device_clusters)

homeassistant/components/zha/core/const.py

@@ -1,5 +1,6 @@
 """All constants related to the ZHA component."""
 import enum
+import logging
 
 from homeassistant.components.binary_sensor import DOMAIN as BINARY_SENSOR
 from homeassistant.components.fan import DOMAIN as FAN
@@ -106,6 +107,34 @@ QUIRK_CLASS = 'quirk_class'
 MANUFACTURER_CODE = 'manufacturer_code'
 POWER_SOURCE = 'power_source'
 
+BELLOWS = 'bellows'
+ZHA = 'homeassistant.components.zha'
+ZIGPY = 'zigpy'
+ZIGPY_XBEE = 'zigpy_xbee'
+ZIGPY_DECONZ = 'zigpy_deconz'
+ORIGINAL = 'original'
+CURRENT = 'current'
+DEBUG_LEVELS = {
+    BELLOWS: logging.DEBUG,
+    ZHA: logging.DEBUG,
+    ZIGPY: logging.DEBUG,
+    ZIGPY_XBEE: logging.DEBUG,
+    ZIGPY_DECONZ: logging.DEBUG,
+}
+ADD_DEVICE_RELAY_LOGGERS = [ZHA, ZIGPY]
+TYPE = 'type'
+NWK = 'nwk'
+SIGNATURE = 'signature'
+RAW_INIT = 'raw_device_initialized'
+ZHA_GW_MSG = 'zha_gateway_message'
+DEVICE_REMOVED = 'device_removed'
+DEVICE_INFO = 'device_info'
+DEVICE_FULL_INIT = 'device_fully_initialized'
+DEVICE_JOINED = 'device_joined'
+LOG_OUTPUT = 'log_output'
+LOG_ENTRY = 'log_entry'
+MFG_CLUSTER_ID_START = 0xfc00
+
 
 class RadioType(enum.Enum):
     """Possible options for radio type."""

homeassistant/components/zha/core/gateway.py

@@ -11,6 +11,8 @@ import itertools
 import logging
 import os
 
+import traceback
+from homeassistant.components.system_log import LogEntry, _figure_out_source
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.entity_component import EntityComponent
@@ -18,7 +20,11 @@ from .const import (
     DATA_ZHA, DATA_ZHA_CORE_COMPONENT, DOMAIN, SIGNAL_REMOVE, DATA_ZHA_GATEWAY,
     CONF_USB_PATH, CONF_BAUDRATE, DEFAULT_BAUDRATE, CONF_RADIO_TYPE,
     DATA_ZHA_RADIO, CONF_DATABASE, DEFAULT_DATABASE_NAME, DATA_ZHA_BRIDGE_ID,
-    RADIO, CONTROLLER, RADIO_DESCRIPTION
+    RADIO, CONTROLLER, RADIO_DESCRIPTION, BELLOWS, ZHA, ZIGPY, ZIGPY_XBEE,
+    ZIGPY_DECONZ, ORIGINAL, CURRENT, DEBUG_LEVELS, ADD_DEVICE_RELAY_LOGGERS,
+    TYPE, NWK, IEEE, MODEL, SIGNATURE, ATTR_MANUFACTURER, RAW_INIT,
+    ZHA_GW_MSG, DEVICE_REMOVED, DEVICE_INFO, DEVICE_FULL_INIT, DEVICE_JOINED,
+    LOG_OUTPUT, LOG_ENTRY
 )
 from .device import ZHADevice, DeviceStatus
 from .channels import (
@@ -32,6 +38,7 @@ from .discovery import (
 from .store import async_get_registry
 from .patches import apply_application_controller_patch
 from .registries import RADIO_TYPES
+from ..api import async_get_device_info
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -54,6 +61,12 @@ class ZHAGateway:
         self.radio_description = None
         hass.data[DATA_ZHA][DATA_ZHA_CORE_COMPONENT] = self._component
         hass.data[DATA_ZHA][DATA_ZHA_GATEWAY] = self
+        self._log_levels = {
+            ORIGINAL: async_capture_log_levels(),
+            CURRENT: async_capture_log_levels()
+        }
+        self.debug_enabled = False
+        self._log_relay_handler = LogRelayHandler(hass, self)
 
     async def async_initialize(self, config_entry):
         """Initialize controller and connect radio."""
@@ -94,13 +107,37 @@ class ZHAGateway:
         At this point, no information about the device is known other than its
         address
         """
-        # Wait for device_initialized, instead
+        async_dispatcher_send(
-        pass
+            self._hass,
+            ZHA_GW_MSG,
+            {
+                TYPE: DEVICE_JOINED,
+                NWK: device.nwk,
+                IEEE: str(device.ieee)
+            }
+        )
 
     def raw_device_initialized(self, device):
         """Handle a device initialization without quirks loaded."""
-        # Wait for device_initialized, instead
+        endpoint_ids = device.endpoints.keys()
-        pass
+        ept_id = next((ept_id for ept_id in endpoint_ids if ept_id != 0), None)
+        manufacturer = 'Unknown'
+        model = 'Unknown'
+        if ept_id is not None:
+            manufacturer = device.endpoints[ept_id].manufacturer
+            model = device.endpoints[ept_id].model
+        async_dispatcher_send(
+            self._hass,
+            ZHA_GW_MSG,
+            {
+                TYPE: RAW_INIT,
+                NWK: device.nwk,
+                IEEE: str(device.ieee),
+                MODEL: model,
+                ATTR_MANUFACTURER: manufacturer,
+                SIGNATURE: device.get_signature()
+            }
+        )
 
     def device_initialized(self, device):
         """Handle device joined and basic information discovered."""
@@ -116,11 +153,21 @@ class ZHAGateway:
         device = self._devices.pop(device.ieee, None)
         self._device_registry.pop(device.ieee, None)
         if device is not None:
+            device_info = async_get_device_info(self._hass, device)
             self._hass.async_create_task(device.async_unsub_dispatcher())
             async_dispatcher_send(
                 self._hass,
                 "{}_{}".format(SIGNAL_REMOVE, str(device.ieee))
             )
+            if device_info is not None:
+                async_dispatcher_send(
+                    self._hass,
+                    ZHA_GW_MSG,
+                    {
+                        TYPE: DEVICE_REMOVED,
+                        DEVICE_INFO: device_info
+                    }
+                )
 
     def get_device(self, ieee_str):
         """Return ZHADevice for given ieee."""
@@ -157,6 +204,28 @@ class ZHAGateway:
             )
         )
 
+    @callback
+    def async_enable_debug_mode(self):
+        """Enable debug mode for ZHA."""
+        self._log_levels[ORIGINAL] = async_capture_log_levels()
+        async_set_logger_levels(DEBUG_LEVELS)
+        self._log_levels[CURRENT] = async_capture_log_levels()
+
+        for logger_name in ADD_DEVICE_RELAY_LOGGERS:
+            logging.getLogger(logger_name).addHandler(self._log_relay_handler)
+
+        self.debug_enabled = True
+
+    @callback
+    def async_disable_debug_mode(self):
+        """Disable debug mode for ZHA."""
+        async_set_logger_levels(self._log_levels[ORIGINAL])
+        self._log_levels[CURRENT] = async_capture_log_levels()
+        for logger_name in ADD_DEVICE_RELAY_LOGGERS:
+            logging.getLogger(logger_name).removeHandler(
+                self._log_relay_handler)
+        self.debug_enabled = False
+
     @callback
     def _async_get_or_create_device(self, zigpy_device, is_new_join):
         """Get or create a ZHA device."""
@@ -231,3 +300,64 @@ class ZHAGateway:
 
         device_entity = async_create_device_entity(zha_device)
         await self._component.async_add_entities([device_entity])
+
+        if is_new_join:
+            device_info = async_get_device_info(self._hass, zha_device)
+            async_dispatcher_send(
+                self._hass,
+                ZHA_GW_MSG,
+                {
+                    TYPE: DEVICE_FULL_INIT,
+                    DEVICE_INFO: device_info
+                }
+            )
+
+
+@callback
+def async_capture_log_levels():
+    """Capture current logger levels for ZHA."""
+    return {
+        BELLOWS: logging.getLogger(BELLOWS).getEffectiveLevel(),
+        ZHA: logging.getLogger(ZHA).getEffectiveLevel(),
+        ZIGPY: logging.getLogger(ZIGPY).getEffectiveLevel(),
+        ZIGPY_XBEE: logging.getLogger(ZIGPY_XBEE).getEffectiveLevel(),
+        ZIGPY_DECONZ: logging.getLogger(ZIGPY_DECONZ).getEffectiveLevel(),
+    }
+
+
+@callback
+def async_set_logger_levels(levels):
+    """Set logger levels for ZHA."""
+    logging.getLogger(BELLOWS).setLevel(levels[BELLOWS])
+    logging.getLogger(ZHA).setLevel(levels[ZHA])
+    logging.getLogger(ZIGPY).setLevel(levels[ZIGPY])
+    logging.getLogger(ZIGPY_XBEE).setLevel(levels[ZIGPY_XBEE])
+    logging.getLogger(ZIGPY_DECONZ).setLevel(levels[ZIGPY_DECONZ])
+
+
+class LogRelayHandler(logging.Handler):
+    """Log handler for error messages."""
+
+    def __init__(self, hass, gateway):
+        """Initialize a new LogErrorHandler."""
+        super().__init__()
+        self.hass = hass
+        self.gateway = gateway
+
+    def emit(self, record):
+        """Relay log message via dispatcher."""
+        stack = []
+        if record.levelno >= logging.WARN:
+            if not record.exc_info:
+                stack = [f for f, _, _, _ in traceback.extract_stack()]
+
+        entry = LogEntry(record, stack,
+                         _figure_out_source(record, stack, self.hass))
+        async_dispatcher_send(
+            self.hass,
+            ZHA_GW_MSG,
+            {
+                TYPE: LOG_OUTPUT,
+                LOG_ENTRY: entry.to_dict()
+            }
+        )