Add multi-factor auth module setup flow

homeassistant/auth/__init__.py

@@ -6,8 +6,6 @@
 
 import jwt
 
-import voluptuous as vol
-
 from homeassistant import data_entry_flow
 from homeassistant.core import callback, HomeAssistant
 from homeassistant.util import dt as dt_util
@@ -235,13 +233,6 @@ async def async_enable_user_mfa(self, user: models.User,
             raise ValueError('Unable find multi-factor auth module: {}'
                              .format(mfa_module_id))
 
-        if module.setup_schema is not None:
-            try:
-                # pylint: disable=not-callable
-                data = module.setup_schema(data)
-            except vol.Invalid as err:
-                raise ValueError('Data does not match schema: {}'.format(err))
-
         await module.async_setup_user(user.id, data)
 
     async def async_disable_user_mfa(self, user: models.User,

homeassistant/auth/mfa_modules/__init__.py

@@ -8,7 +8,7 @@
 import voluptuous as vol
 from voluptuous.humanize import humanize_error
 
-from homeassistant import requirements
+from homeassistant import requirements, data_entry_flow
 from homeassistant.const import CONF_ID, CONF_NAME, CONF_TYPE
 from homeassistant.core import HomeAssistant
 from homeassistant.util.decorator import Registry
@@ -64,15 +64,14 @@ def input_schema(self) -> vol.Schema:
         """Return a voluptuous schema to define mfa auth module's input."""
         raise NotImplementedError
 
-    @property
-    def setup_schema(self) -> Optional[vol.Schema]:
-        """Return a vol schema to validate mfa auth module's setup input.
+    async def async_setup_flow(self, user_id: str) -> 'SetupFlow':
+        """Return a data entry flow handler for setup module.
 
-        Optional
+        Mfa module should extend SetupFlow
         """
-        return None
+        raise NotImplementedError
 
-    async def async_setup_user(self, user_id: str, setup_data: Any) -> None:
+    async def async_setup_user(self, user_id: str, setup_data: Any) -> Any:
         """Set up user for mfa auth module."""
         raise NotImplementedError
 
@@ -90,6 +89,42 @@ async def async_validation(
         raise NotImplementedError
 
 
+class SetupFlow(data_entry_flow.FlowHandler):
+    """Handler for the setup flow."""
+
+    def __init__(self, auth_module: MultiFactorAuthModule,
+                 setup_schema: vol.Schema,
+                 user_id: str) -> None:
+        """Initialize the setup flow."""
+        self._auth_module = auth_module
+        self._setup_schema = setup_schema
+        self._user_id = user_id
+
+    async def async_step_init(
+            self, user_input: Optional[Dict[str, str]] = None) \
+            -> Dict[str, Any]:
+        """Handle the first step of setup flow.
+
+        Return self.async_show_form(step_id='init') if user_input == None.
+        Return self.async_create_entry(data={'result': result}) if finish.
+        """
+        errors = {}  # type: Dict[str, str]
+
+        if user_input:
+            result = await self._auth_module.async_setup_user(
+                self._user_id, user_input)
+            return self.async_create_entry(
+                title=self._auth_module.name,
+                data={'result': result}
+            )
+
+        return self.async_show_form(
+            step_id='init',
+            data_schema=self._setup_schema,
+            errors=errors
+        )
+
+
 async def auth_mfa_module_from_config(
         hass: HomeAssistant, config: Dict[str, Any]) \
         -> Optional[MultiFactorAuthModule]:

homeassistant/auth/mfa_modules/insecure_example.py

@@ -1,13 +1,13 @@
 """Example auth module."""
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 import voluptuous as vol
 
 from homeassistant.core import HomeAssistant
 
 from . import MultiFactorAuthModule, MULTI_FACTOR_AUTH_MODULES, \
-    MULTI_FACTOR_AUTH_MODULE_SCHEMA
+    MULTI_FACTOR_AUTH_MODULE_SCHEMA, SetupFlow
 
 CONFIG_SCHEMA = MULTI_FACTOR_AUTH_MODULE_SCHEMA.extend({
     vol.Required('data'): [vol.Schema({
@@ -36,11 +36,18 @@ def input_schema(self) -> vol.Schema:
         return vol.Schema({'pin': str})
 
     @property
-    def setup_schema(self) -> Optional[vol.Schema]:
+    def setup_schema(self) -> vol.Schema:
         """Validate async_setup_user input data."""
         return vol.Schema({'pin': str})
 
-    async def async_setup_user(self, user_id: str, setup_data: Any) -> None:
+    async def async_setup_flow(self, user_id: str) -> SetupFlow:
+        """Return a data entry flow handler for setup module.
+
+        Mfa module should extend SetupFlow
+        """
+        return SetupFlow(self, self.setup_schema, user_id)
+
+    async def async_setup_user(self, user_id: str, setup_data: Any) -> Any:
         """Set up user to use mfa module."""
         # data shall has been validate in caller
         pin = setup_data['pin']

homeassistant/components/auth/__init__.py

@@ -68,10 +68,12 @@
 from homeassistant.components.http.ban import log_invalid_auth
 from homeassistant.components.http.data_validator import RequestDataValidator
 from homeassistant.components.http.view import HomeAssistantView
-from homeassistant.core import callback
+from homeassistant.core import callback, HomeAssistant
 from homeassistant.util import dt as dt_util
+
 from . import indieauth
 from . import login_flow
+from . import mfa_setup_flow
 
 DOMAIN = 'auth'
 DEPENDENCIES = ['http']
@@ -100,6 +102,7 @@ async def async_setup(hass, config):
     )
 
     await login_flow.async_setup(hass, store_result)
+    await mfa_setup_flow.async_setup(hass)
 
     return True
 
@@ -315,21 +318,28 @@ def retrieve_result(client_id, result_type, code):
     return store_result, retrieve_result
 
 
+@websocket_api.ws_require_user()
 @callback
-def websocket_current_user(hass, connection, msg):
+def websocket_current_user(
+        hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg):
     """Return the current user."""
-    user = connection.request.get('hass_user')
-
-    if user is None:
-        connection.to_write.put_nowait(websocket_api.error_message(
-            msg['id'], 'no_user', 'Not authenticated as a user'))
-        return
-
-    connection.to_write.put_nowait(websocket_api.result_message(msg['id'], {
-        'id': user.id,
-        'name': user.name,
-        'is_owner': user.is_owner,
-        'credentials': [{'auth_provider_type': c.auth_provider_type,
-                         'auth_provider_id': c.auth_provider_id}
-                        for c in user.credentials]
-    }))
+    async def async_get_current_user(user):
+        """Get current user."""
+        enabled_modules = await hass.auth.async_get_enabled_mfa(user)
+
+        connection.send_message_outside(
+            websocket_api.result_message(msg['id'], {
+                'id': user.id,
+                'name': user.name,
+                'is_owner': user.is_owner,
+                'credentials': [{'auth_provider_type': c.auth_provider_type,
+                                 'auth_provider_id': c.auth_provider_id}
+                                for c in user.credentials],
+                'mfa_modules': [{
+                    'id': module.id,
+                    'name': module.name,
+                    'enabled': module.id in enabled_modules,
+                } for module in hass.auth.auth_mfa_modules],
+            }))
+
+    hass.async_create_task(async_get_current_user(connection.user))

homeassistant/components/auth/mfa_setup_flow.py

@@ -0,0 +1,134 @@
+"""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. hanlder 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_outside(
+                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_outside(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_outside(
+            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_outside(websocket_api.error_message(
+                msg['id'], 'disable_failed',
+                'Cannot disable MFA Module {}: {}'.format(
+                    mfa_module_id, err)))
+            return
+
+        connection.send_message_outside(
+            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