core/homeassistant/components/auth/mfa_setup_flow.py

"""Helpers to setup multi-factor auth module."""
import logging

import voluptuous as vol

from homeassistant import data_entry_flow
from homeassistant.components import websocket_api
from homeassistant.core import callback, HomeAssistant

WS_TYPE_SETUP_MFA = "auth/setup_mfa"
SCHEMA_WS_SETUP_MFA = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend(
    {
        vol.Required("type"): WS_TYPE_SETUP_MFA,
        vol.Exclusive("mfa_module_id", "module_or_flow_id"): str,
        vol.Exclusive("flow_id", "module_or_flow_id"): str,
        vol.Optional("user_input"): object,
    }
)

WS_TYPE_DEPOSE_MFA = "auth/depose_mfa"
SCHEMA_WS_DEPOSE_MFA = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend(
    {vol.Required("type"): WS_TYPE_DEPOSE_MFA, vol.Required("mfa_module_id"): str}
)

DATA_SETUP_FLOW_MGR = "auth_mfa_setup_flow_manager"

_LOGGER = logging.getLogger(__name__)


async def async_setup(hass):
    """Init mfa setup flow manager."""

    async def _async_create_setup_flow(handler, context, data):
        """Create a setup flow. handler is a mfa module."""
        mfa_module = hass.auth.get_auth_mfa_module(handler)
        if mfa_module is None:
            raise ValueError("Mfa module {} is not found".format(handler))

        user_id = data.pop("user_id")
        return await mfa_module.async_setup_flow(user_id)

    async def _async_finish_setup_flow(flow, flow_result):
        _LOGGER.debug("flow_result: %s", flow_result)
        return flow_result

    hass.data[DATA_SETUP_FLOW_MGR] = data_entry_flow.FlowManager(
        hass, _async_create_setup_flow, _async_finish_setup_flow
    )

    hass.components.websocket_api.async_register_command(
        WS_TYPE_SETUP_MFA, websocket_setup_mfa, SCHEMA_WS_SETUP_MFA
    )

    hass.components.websocket_api.async_register_command(
        WS_TYPE_DEPOSE_MFA, websocket_depose_mfa, SCHEMA_WS_DEPOSE_MFA
    )


@callback
@websocket_api.ws_require_user(allow_system_user=False)
def websocket_setup_mfa(
    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg
):
    """Return a setup flow for mfa auth module."""

    async def async_setup_flow(msg):
        """Return a setup flow for mfa auth module."""
        flow_manager = hass.data[DATA_SETUP_FLOW_MGR]

        flow_id = msg.get("flow_id")
        if flow_id is not None:
            result = await flow_manager.async_configure(flow_id, msg.get("user_input"))
            connection.send_message(
                websocket_api.result_message(msg["id"], _prepare_result_json(result))
            )
            return

        mfa_module_id = msg.get("mfa_module_id")
        mfa_module = hass.auth.get_auth_mfa_module(mfa_module_id)
        if mfa_module is None:
            connection.send_message(
                websocket_api.error_message(
                    msg["id"],
                    "no_module",
                    "MFA module {} is not found".format(mfa_module_id),
                )
            )
            return

        result = await flow_manager.async_init(
            mfa_module_id, data={"user_id": connection.user.id}
        )

        connection.send_message(
            websocket_api.result_message(msg["id"], _prepare_result_json(result))
        )

    hass.async_create_task(async_setup_flow(msg))


@callback
@websocket_api.ws_require_user(allow_system_user=False)
def websocket_depose_mfa(
    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg
):
    """Remove user from mfa module."""

    async def async_depose(msg):
        """Remove user from mfa auth module."""
        mfa_module_id = msg["mfa_module_id"]
        try:
            await hass.auth.async_disable_user_mfa(
                connection.user, msg["mfa_module_id"]
            )
        except ValueError as err:
            connection.send_message(
                websocket_api.error_message(
                    msg["id"],
                    "disable_failed",
                    "Cannot disable MFA Module {}: {}".format(mfa_module_id, err),
                )
            )
            return

        connection.send_message(websocket_api.result_message(msg["id"], "done"))

    hass.async_create_task(async_depose(msg))


def _prepare_result_json(result):
    """Convert result to JSON."""
    if result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY:
        data = result.copy()
        return data

    if result["type"] != data_entry_flow.RESULT_TYPE_FORM:
        return result

    import voluptuous_serialize

    data = result.copy()

    schema = data["data_schema"]
    if schema is None:
        data["data_schema"] = []
    else:
        data["data_schema"] = voluptuous_serialize.convert(schema)

    return data