core/homeassistant/components/mqtt/config_flow.py

146 lines
4.7 KiB
Python

"""Config flow for MQTT."""
from collections import OrderedDict
import queue
import voluptuous as vol
from homeassistant import config_entries
from homeassistant.const import (
CONF_HOST, CONF_PASSWORD, CONF_PORT, CONF_PROTOCOL, CONF_USERNAME)
from .const import CONF_BROKER, CONF_DISCOVERY, DEFAULT_DISCOVERY
@config_entries.HANDLERS.register('mqtt')
class FlowHandler(config_entries.ConfigFlow):
"""Handle a config flow."""
VERSION = 1
CONNECTION_CLASS = config_entries.CONN_CLASS_LOCAL_PUSH
_hassio_discovery = None
async def async_step_user(self, user_input=None):
"""Handle a flow initialized by the user."""
if self._async_current_entries():
return self.async_abort(reason='single_instance_allowed')
return await self.async_step_broker()
async def async_step_broker(self, user_input=None):
"""Confirm the setup."""
errors = {}
if user_input is not None:
can_connect = await self.hass.async_add_executor_job(
try_connection, user_input[CONF_BROKER], user_input[CONF_PORT],
user_input.get(CONF_USERNAME), user_input.get(CONF_PASSWORD))
if can_connect:
return self.async_create_entry(
title=user_input[CONF_BROKER], data=user_input)
errors['base'] = 'cannot_connect'
fields = OrderedDict()
fields[vol.Required(CONF_BROKER)] = str
fields[vol.Required(CONF_PORT, default=1883)] = vol.Coerce(int)
fields[vol.Optional(CONF_USERNAME)] = str
fields[vol.Optional(CONF_PASSWORD)] = str
fields[vol.Optional(CONF_DISCOVERY, default=DEFAULT_DISCOVERY)] = bool
return self.async_show_form(
step_id='broker', data_schema=vol.Schema(fields), errors=errors)
async def async_step_import(self, user_input):
"""Import a config entry.
Special type of import, we're not actually going to store any data.
Instead, we're going to rely on the values that are in config file.
"""
if self._async_current_entries():
return self.async_abort(reason='single_instance_allowed')
return self.async_create_entry(title='configuration.yaml', data={})
async def async_step_hassio(self, user_input=None):
"""Receive a Hass.io discovery."""
if self._async_current_entries():
return self.async_abort(reason='single_instance_allowed')
self._hassio_discovery = user_input
return await self.async_step_hassio_confirm()
async def async_step_hassio_confirm(self, user_input=None):
"""Confirm a Hass.io discovery."""
errors = {}
if user_input is not None:
data = self._hassio_discovery
can_connect = await self.hass.async_add_executor_job(
try_connection,
data[CONF_HOST],
data[CONF_PORT],
data.get(CONF_USERNAME),
data.get(CONF_PASSWORD),
data.get(CONF_PROTOCOL)
)
if can_connect:
return self.async_create_entry(
title=data['addon'], data={
CONF_BROKER: data[CONF_HOST],
CONF_PORT: data[CONF_PORT],
CONF_USERNAME: data.get(CONF_USERNAME),
CONF_PASSWORD: data.get(CONF_PASSWORD),
CONF_PROTOCOL: data.get(CONF_PROTOCOL),
CONF_DISCOVERY: user_input[CONF_DISCOVERY],
})
errors['base'] = 'cannot_connect'
return self.async_show_form(
step_id='hassio_confirm',
description_placeholders={
'addon': self._hassio_discovery['addon']
},
data_schema=vol.Schema({
vol.Optional(CONF_DISCOVERY, default=DEFAULT_DISCOVERY): bool
}),
errors=errors,
)
def try_connection(broker, port, username, password, protocol='3.1'):
"""Test if we can connect to an MQTT broker."""
import paho.mqtt.client as mqtt
if protocol == '3.1':
proto = mqtt.MQTTv31
else:
proto = mqtt.MQTTv311
client = mqtt.Client(protocol=proto)
if username and password:
client.username_pw_set(username, password)
result = queue.Queue(maxsize=1)
def on_connect(client_, userdata, flags, result_code):
"""Handle connection result."""
result.put(result_code == mqtt.CONNACK_ACCEPTED)
client.on_connect = on_connect
client.connect_async(broker, port)
client.loop_start()
try:
return result.get(timeout=5)
except queue.Empty:
return False
finally:
client.disconnect()
client.loop_stop()