"""Websocket API for automation.""" import json import voluptuous as vol from homeassistant.components import websocket_api from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import HomeAssistantError from homeassistant.helpers.dispatcher import ( DATA_DISPATCHER, async_dispatcher_connect, async_dispatcher_send, ) from homeassistant.helpers.script import ( SCRIPT_BREAKPOINT_HIT, SCRIPT_DEBUG_CONTINUE_ALL, breakpoint_clear, breakpoint_clear_all, breakpoint_list, breakpoint_set, debug_continue, debug_step, debug_stop, ) from .trace import ( DATA_TRACE, get_debug_trace, get_debug_traces, get_debug_traces_for_automation, ) from .utils import TraceJSONEncoder # mypy: allow-untyped-calls, allow-untyped-defs @callback def async_setup(hass: HomeAssistant) -> None: """Set up the websocket API.""" websocket_api.async_register_command(hass, websocket_automation_trace_get) websocket_api.async_register_command(hass, websocket_automation_trace_list) websocket_api.async_register_command(hass, websocket_automation_trace_contexts) websocket_api.async_register_command(hass, websocket_automation_breakpoint_clear) websocket_api.async_register_command(hass, websocket_automation_breakpoint_list) websocket_api.async_register_command(hass, websocket_automation_breakpoint_set) websocket_api.async_register_command(hass, websocket_automation_debug_continue) websocket_api.async_register_command(hass, websocket_automation_debug_step) websocket_api.async_register_command(hass, websocket_automation_debug_stop) websocket_api.async_register_command(hass, websocket_subscribe_breakpoint_events) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/trace/get", vol.Required("automation_id"): str, vol.Required("run_id"): str, } ) def websocket_automation_trace_get(hass, connection, msg): """Get an automation trace.""" automation_id = msg["automation_id"] run_id = msg["run_id"] trace = get_debug_trace(hass, automation_id, run_id) message = websocket_api.messages.result_message(msg["id"], trace) connection.send_message(json.dumps(message, cls=TraceJSONEncoder, allow_nan=False)) @callback @websocket_api.require_admin @websocket_api.websocket_command( {vol.Required("type"): "automation/trace/list", vol.Optional("automation_id"): str} ) def websocket_automation_trace_list(hass, connection, msg): """Summarize automation traces.""" automation_id = msg.get("automation_id") if not automation_id: automation_traces = get_debug_traces(hass, summary=True) else: automation_traces = get_debug_traces_for_automation( hass, automation_id, summary=True ) connection.send_result(msg["id"], automation_traces) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/trace/contexts", vol.Optional("automation_id"): str, } ) def websocket_automation_trace_contexts(hass, connection, msg): """Retrieve contexts we have traces for.""" automation_id = msg.get("automation_id") if automation_id is not None: values = {automation_id: hass.data[DATA_TRACE].get(automation_id, {})} else: values = hass.data[DATA_TRACE] contexts = { trace.context.id: {"run_id": trace.run_id, "automation_id": automation_id} for automation_id, traces in values.items() for trace in traces.values() } connection.send_result(msg["id"], contexts) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/debug/breakpoint/set", vol.Required("automation_id"): str, vol.Required("node"): str, vol.Optional("run_id"): str, } ) def websocket_automation_breakpoint_set(hass, connection, msg): """Set breakpoint.""" automation_id = msg["automation_id"] node = msg["node"] run_id = msg.get("run_id") if ( SCRIPT_BREAKPOINT_HIT not in hass.data.get(DATA_DISPATCHER, {}) or not hass.data[DATA_DISPATCHER][SCRIPT_BREAKPOINT_HIT] ): raise HomeAssistantError("No breakpoint subscription") result = breakpoint_set(hass, automation_id, run_id, node) connection.send_result(msg["id"], result) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/debug/breakpoint/clear", vol.Required("automation_id"): str, vol.Required("node"): str, vol.Optional("run_id"): str, } ) def websocket_automation_breakpoint_clear(hass, connection, msg): """Clear breakpoint.""" automation_id = msg["automation_id"] node = msg["node"] run_id = msg.get("run_id") result = breakpoint_clear(hass, automation_id, run_id, node) connection.send_result(msg["id"], result) @callback @websocket_api.require_admin @websocket_api.websocket_command( {vol.Required("type"): "automation/debug/breakpoint/list"} ) def websocket_automation_breakpoint_list(hass, connection, msg): """List breakpoints.""" breakpoints = breakpoint_list(hass) for _breakpoint in breakpoints: _breakpoint["automation_id"] = _breakpoint.pop("unique_id") connection.send_result(msg["id"], breakpoints) @callback @websocket_api.require_admin @websocket_api.websocket_command( {vol.Required("type"): "automation/debug/breakpoint/subscribe"} ) def websocket_subscribe_breakpoint_events(hass, connection, msg): """Subscribe to breakpoint events.""" @callback def breakpoint_hit(automation_id, run_id, node): """Forward events to websocket.""" connection.send_message( websocket_api.event_message( msg["id"], { "automation_id": automation_id, "run_id": run_id, "node": node, }, ) ) remove_signal = async_dispatcher_connect( hass, SCRIPT_BREAKPOINT_HIT, breakpoint_hit ) @callback def unsub(): """Unsubscribe from breakpoint events.""" remove_signal() if ( SCRIPT_BREAKPOINT_HIT not in hass.data.get(DATA_DISPATCHER, {}) or not hass.data[DATA_DISPATCHER][SCRIPT_BREAKPOINT_HIT] ): breakpoint_clear_all(hass) async_dispatcher_send(hass, SCRIPT_DEBUG_CONTINUE_ALL) connection.subscriptions[msg["id"]] = unsub connection.send_message(websocket_api.result_message(msg["id"])) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/debug/continue", vol.Required("automation_id"): str, vol.Required("run_id"): str, } ) def websocket_automation_debug_continue(hass, connection, msg): """Resume execution of halted automation.""" automation_id = msg["automation_id"] run_id = msg["run_id"] result = debug_continue(hass, automation_id, run_id) connection.send_result(msg["id"], result) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/debug/step", vol.Required("automation_id"): str, vol.Required("run_id"): str, } ) def websocket_automation_debug_step(hass, connection, msg): """Single step a halted automation.""" automation_id = msg["automation_id"] run_id = msg["run_id"] result = debug_step(hass, automation_id, run_id) connection.send_result(msg["id"], result) @callback @websocket_api.require_admin @websocket_api.websocket_command( { vol.Required("type"): "automation/debug/stop", vol.Required("automation_id"): str, vol.Required("run_id"): str, } ) def websocket_automation_debug_stop(hass, connection, msg): """Stop a halted automation.""" automation_id = msg["automation_id"] run_id = msg["run_id"] result = debug_stop(hass, automation_id, run_id) connection.send_result(msg["id"], result)