core/homeassistant/components/octoprint/config_flow.py

222 lines
7.7 KiB
Python

"""Config flow for OctoPrint integration."""
import logging
from pyoctoprintapi import ApiError, OctoprintClient, OctoprintException
import voluptuous as vol
from yarl import URL
from homeassistant import config_entries, data_entry_flow, exceptions
from homeassistant.components import ssdp, zeroconf
from homeassistant.const import (
CONF_API_KEY,
CONF_HOST,
CONF_PATH,
CONF_PORT,
CONF_SSL,
CONF_USERNAME,
CONF_VERIFY_SSL,
)
from homeassistant.helpers.aiohttp_client import async_get_clientsession
import homeassistant.helpers.config_validation as cv
from .const import DOMAIN
_LOGGER = logging.getLogger(__name__)
def _schema_with_defaults(
username="", host="", port=80, path="/", ssl=False, verify_ssl=True
):
return vol.Schema(
{
vol.Required(CONF_USERNAME, default=username): str,
vol.Required(CONF_HOST, default=host): str,
vol.Required(CONF_PORT, default=port): cv.port,
vol.Required(CONF_PATH, default=path): str,
vol.Required(CONF_SSL, default=ssl): bool,
vol.Required(CONF_VERIFY_SSL, default=verify_ssl): bool,
},
extra=vol.ALLOW_EXTRA,
)
class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN):
"""Handle a config flow for OctoPrint."""
VERSION = 1
api_key_task = None
def __init__(self) -> None:
"""Handle a config flow for OctoPrint."""
self.discovery_schema = None
self._user_input = None
async def async_step_user(self, user_input=None):
"""Handle the initial step."""
# When coming back from the progress steps, the user_input is stored in the
# instance variable instead of being passed in
if user_input is None and self._user_input:
user_input = self._user_input
if user_input is None:
data = self.discovery_schema or _schema_with_defaults()
return self.async_show_form(step_id="user", data_schema=data)
if CONF_API_KEY in user_input:
errors = {}
try:
return await self._finish_config(user_input)
except data_entry_flow.AbortFlow as err:
raise err from None
except CannotConnect:
errors["base"] = "cannot_connect"
except Exception: # pylint: disable=broad-except
errors["base"] = "unknown"
if errors:
return self.async_show_form(
step_id="user",
errors=errors,
data_schema=_schema_with_defaults(
user_input.get(CONF_USERNAME),
user_input[CONF_HOST],
user_input[CONF_PORT],
user_input[CONF_PATH],
user_input[CONF_SSL],
user_input[CONF_VERIFY_SSL],
),
)
self.api_key_task = None
return await self.async_step_get_api_key(user_input)
async def async_step_get_api_key(self, user_input):
"""Get an Application Api Key."""
if not self.api_key_task:
self.api_key_task = self.hass.async_create_task(
self._async_get_auth_key(user_input)
)
return self.async_show_progress(
step_id="get_api_key", progress_action="get_api_key"
)
try:
await self.api_key_task
except OctoprintException as err:
_LOGGER.exception("Failed to get an application key: %s", err)
return self.async_show_progress_done(next_step_id="auth_failed")
except Exception as err: # pylint: disable=broad-except
_LOGGER.exception("Failed to get an application key : %s", err)
return self.async_show_progress_done(next_step_id="auth_failed")
# store this off here to pick back up in the user step
self._user_input = user_input
return self.async_show_progress_done(next_step_id="user")
async def _finish_config(self, user_input):
"""Finish the configuration setup."""
octoprint = self._get_octoprint_client(user_input)
octoprint.set_api_key(user_input[CONF_API_KEY])
try:
discovery = await octoprint.get_discovery_info()
except ApiError as err:
_LOGGER.error("Failed to connect to printer")
raise CannotConnect from err
await self.async_set_unique_id(discovery.upnp_uuid, raise_on_progress=False)
self._abort_if_unique_id_configured()
return self.async_create_entry(title=user_input[CONF_HOST], data=user_input)
async def async_step_auth_failed(self, user_input):
"""Handle api fetch failure."""
return self.async_abort(reason="auth_failed")
async def async_step_import(self, user_input):
"""Handle import."""
return await self.async_step_user(user_input)
async def async_step_zeroconf(
self, discovery_info: zeroconf.ZeroconfServiceInfo
) -> data_entry_flow.FlowResult:
"""Handle discovery flow."""
uuid = discovery_info.properties["uuid"]
await self.async_set_unique_id(uuid)
self._abort_if_unique_id_configured()
self.context.update(
{
"title_placeholders": {CONF_HOST: discovery_info.host},
"configuration_url": (
f"http://{discovery_info.host}:{discovery_info.port}"
f"{discovery_info.properties[CONF_PATH]}"
),
}
)
self.discovery_schema = _schema_with_defaults(
host=discovery_info.host,
port=discovery_info.port,
path=discovery_info.properties[CONF_PATH],
)
return await self.async_step_user()
async def async_step_ssdp(
self, discovery_info: ssdp.SsdpServiceInfo
) -> data_entry_flow.FlowResult:
"""Handle ssdp discovery flow."""
uuid = discovery_info.upnp["UDN"][5:]
await self.async_set_unique_id(uuid)
self._abort_if_unique_id_configured()
url = URL(discovery_info.upnp["presentationURL"])
self.context.update(
{
"title_placeholders": {CONF_HOST: url.host},
"configuration_url": discovery_info.upnp["presentationURL"],
}
)
self.discovery_schema = _schema_with_defaults(
host=url.host,
path=url.path,
port=url.port,
ssl=url.scheme == "https",
)
return await self.async_step_user()
async def _async_get_auth_key(self, user_input: dict):
"""Get application api key."""
octoprint = self._get_octoprint_client(user_input)
try:
user_input[CONF_API_KEY] = await octoprint.request_app_key(
"Home Assistant", user_input[CONF_USERNAME], 300
)
finally:
# Continue the flow after show progress when the task is done.
self.hass.async_create_task(
self.hass.config_entries.flow.async_configure(
flow_id=self.flow_id, user_input=user_input
)
)
def _get_octoprint_client(self, user_input: dict) -> OctoprintClient:
"""Build an octoprint client from the user_input."""
verify_ssl = user_input.get(CONF_VERIFY_SSL, True)
session = async_get_clientsession(self.hass, verify_ssl=verify_ssl)
return OctoprintClient(
user_input[CONF_HOST],
session,
user_input[CONF_PORT],
user_input[CONF_SSL],
user_input[CONF_PATH],
)
class CannotConnect(exceptions.HomeAssistantError):
"""Error to indicate we cannot connect."""