"""Helpers to setup multi-factor auth module.""" from __future__ import annotations import logging from typing import Any import voluptuous as vol import voluptuous_serialize from homeassistant import data_entry_flow from homeassistant.components import websocket_api from homeassistant.core import HomeAssistant, callback import homeassistant.helpers.config_validation as cv WS_TYPE_SETUP_MFA = "auth/setup_mfa" SCHEMA_WS_SETUP_MFA = vol.All( 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, } ), cv.has_at_least_one_key("mfa_module_id", "flow_id"), ) 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__) class MfaFlowManager(data_entry_flow.FlowManager): """Manage multi factor authentication flows.""" async def async_create_flow( # type: ignore[override] self, handler_key: str, *, context: dict[str, Any], data: dict[str, Any], ) -> data_entry_flow.FlowHandler: """Create a setup flow. handler is a mfa module.""" mfa_module = self.hass.auth.get_auth_mfa_module(handler_key) if mfa_module is None: raise ValueError(f"Mfa module {handler_key} is not found") user_id = data.pop("user_id") return await mfa_module.async_setup_flow(user_id) async def async_finish_flow( self, flow: data_entry_flow.FlowHandler, result: data_entry_flow.FlowResult ) -> data_entry_flow.FlowResult: """Complete an mfs setup flow.""" _LOGGER.debug("flow_result: %s", result) return result async def async_setup(hass: HomeAssistant) -> None: """Init mfa setup flow manager.""" hass.data[DATA_SETUP_FLOW_MGR] = MfaFlowManager(hass) websocket_api.async_register_command( hass, WS_TYPE_SETUP_MFA, websocket_setup_mfa, SCHEMA_WS_SETUP_MFA ) websocket_api.async_register_command( hass, 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: dict[str, Any] ) -> None: """Return a setup flow for mfa auth module.""" async def async_setup_flow(msg: dict[str, Any]) -> None: """Return a setup flow for mfa auth module.""" flow_manager: MfaFlowManager = hass.data[DATA_SETUP_FLOW_MGR] if (flow_id := msg.get("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["mfa_module_id"] if hass.auth.get_auth_mfa_module(mfa_module_id) is None: connection.send_message( websocket_api.error_message( msg["id"], "no_module", f"MFA module {mfa_module_id} is not found" ) ) 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: dict[str, Any] ) -> None: """Remove user from mfa module.""" async def async_depose(msg: dict[str, Any]) -> None: """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", f"Cannot disable MFA Module {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: data_entry_flow.FlowResult, ) -> data_entry_flow.FlowResult: """Convert result to JSON.""" if result["type"] == data_entry_flow.FlowResultType.CREATE_ENTRY: data = result.copy() return data if result["type"] != data_entry_flow.FlowResultType.FORM: return result data = result.copy() if (schema := data["data_schema"]) is None: data["data_schema"] = [] else: data["data_schema"] = voluptuous_serialize.convert(schema) return data