Add multi-factor auth module setup flow (#16141)

* Add mfa setup flow

* Lint

* Address code review comment

* Fix unit test

* Add assertion for WS response ordering

* Missed a return

* Remove setup_schema from MFA base class

* Move auth.util.validate_current_user -> webscoket_api.ws_require_user

homeassistant/auth/__init__.py

@@ -6,8 +6,6 @@ from typing import Any, Dict, List, Optional, Tuple, cast
 
 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 @@ class AuthManager:
             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 @@ from typing import Any, Dict, Optional
 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 @@ class MultiFactorAuthModule:
         """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 @@ class MultiFactorAuthModule:
         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 @@ class InsecureExampleModule(MultiFactorAuthModule):
         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 import websocket_api
 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 _create_auth_code_store():
     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')
+    async def async_get_current_user(user):
+        """Get current user."""
+        enabled_modules = await hass.auth.async_get_enabled_mfa(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.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],
+            }))
 
-    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]
-    }))
+    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

homeassistant/components/websocket_api.py

@@ -18,7 +18,7 @@ from voluptuous.humanize import humanize_error
 from homeassistant.const import (
     MATCH_ALL, EVENT_TIME_CHANGED, EVENT_HOMEASSISTANT_STOP,
     __version__)
-from homeassistant.core import Context, callback
+from homeassistant.core import Context, callback, HomeAssistant
 from homeassistant.loader import bind_hass
 from homeassistant.helpers.json import JSONEncoder
 from homeassistant.helpers import config_validation as cv
@@ -576,3 +576,59 @@ def handle_ping(hass, connection, msg):
     Async friendly.
     """
     connection.to_write.put_nowait(pong_message(msg['id']))
+
+
+def ws_require_user(
+        only_owner=False, only_system_user=False, allow_system_user=True,
+        only_active_user=True, only_inactive_user=False):
+    """Decorate function validating login user exist in current WS connection.
+
+    Will write out error message if not authenticated.
+    """
+    def validator(func):
+        """Decorate func."""
+        @wraps(func)
+        def check_current_user(hass: HomeAssistant,
+                               connection: ActiveConnection,
+                               msg):
+            """Check current user."""
+            def output_error(message_id, message):
+                """Output error message."""
+                connection.send_message_outside(error_message(
+                    msg['id'], message_id, message))
+
+            if connection.user is None:
+                output_error('no_user', 'Not authenticated as a user')
+                return
+
+            if only_owner and not connection.user.is_owner:
+                output_error('only_owner', 'Only allowed as owner')
+                return
+
+            if (only_system_user and
+                    not connection.user.system_generated):
+                output_error('only_system_user',
+                             'Only allowed as system user')
+                return
+
+            if (not allow_system_user
+                    and connection.user.system_generated):
+                output_error('not_system_user', 'Not allowed as system user')
+                return
+
+            if (only_active_user and
+                    not connection.user.is_active):
+                output_error('only_active_user',
+                             'Only allowed as active user')
+                return
+
+            if only_inactive_user and connection.user.is_active:
+                output_error('only_inactive_user',
+                             'Not allowed as active user')
+                return
+
+            return func(hass, connection, msg)
+
+        return check_current_user
+
+    return validator

tests/auth/mfa_modules/test_insecure_example.py

@@ -125,3 +125,21 @@ async def test_login(hass):
         result['flow_id'], {'pin': '123456'})
     assert result['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
     assert result['data'].id == 'mock-user'
+
+
+async def test_setup_flow(hass):
+    """Test validating pin."""
+    auth_module = await auth_mfa_module_from_config(hass, {
+        'type': 'insecure_example',
+        'data': [{'user_id': 'test-user', 'pin': '123456'}]
+    })
+
+    flow = await auth_module.async_setup_flow('new-user')
+
+    result = await flow.async_step_init()
+    assert result['type'] == data_entry_flow.RESULT_TYPE_FORM
+
+    result = await flow.async_step_init({'pin': 'abcdefg'})
+    assert result['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert auth_module._data[1]['user_id'] == 'new-user'
+    assert auth_module._data[1]['pin'] == 'abcdefg'

tests/components/auth/test_mfa_setup_flow.py

@@ -0,0 +1,99 @@
+"""Tests for the mfa setup flow."""
+from homeassistant import data_entry_flow
+from homeassistant.auth import auth_manager_from_config
+from homeassistant.components.auth import mfa_setup_flow
+from homeassistant.setup import async_setup_component
+
+from tests.common import MockUser, CLIENT_ID, ensure_auth_manager_loaded
+
+
+async def test_ws_setup_depose_mfa(hass, hass_ws_client):
+    """Test set up mfa module for current user."""
+    hass.auth = await auth_manager_from_config(
+        hass, provider_configs=[{
+            'type': 'insecure_example',
+            'users': [{
+                'username': 'test-user',
+                'password': 'test-pass',
+                'name': 'Test Name',
+            }]
+        }], module_configs=[{
+            'type': 'insecure_example',
+            'id': 'example_module',
+            'data': [{'user_id': 'mock-user', 'pin': '123456'}]
+        }])
+    ensure_auth_manager_loaded(hass.auth)
+    await async_setup_component(hass, 'auth', {'http': {}})
+
+    user = MockUser(id='mock-user').add_to_hass(hass)
+    cred = await hass.auth.auth_providers[0].async_get_or_create_credentials(
+        {'username': 'test-user'})
+    await hass.auth.async_link_user(user, cred)
+    refresh_token = await hass.auth.async_create_refresh_token(user, CLIENT_ID)
+    access_token = hass.auth.async_create_access_token(refresh_token)
+
+    client = await hass_ws_client(hass, access_token)
+
+    await client.send_json({
+        'id': 10,
+        'type': mfa_setup_flow.WS_TYPE_SETUP_MFA,
+    })
+
+    result = await client.receive_json()
+    assert result['id'] == 10
+    assert result['success'] is False
+    assert result['error']['code'] == 'no_module'
+
+    await client.send_json({
+        'id': 11,
+        'type': mfa_setup_flow.WS_TYPE_SETUP_MFA,
+        'mfa_module_id': 'example_module',
+    })
+
+    result = await client.receive_json()
+    assert result['id'] == 11
+    assert result['success']
+
+    flow = result['result']
+    assert flow['type'] == data_entry_flow.RESULT_TYPE_FORM
+    assert flow['handler'] == 'example_module'
+    assert flow['step_id'] == 'init'
+    assert flow['data_schema'][0] == {'type': 'string', 'name': 'pin'}
+
+    await client.send_json({
+        'id': 12,
+        'type': mfa_setup_flow.WS_TYPE_SETUP_MFA,
+        'flow_id': flow['flow_id'],
+        'user_input': {'pin': '654321'},
+    })
+
+    result = await client.receive_json()
+    assert result['id'] == 12
+    assert result['success']
+
+    flow = result['result']
+    assert flow['type'] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
+    assert flow['handler'] == 'example_module'
+    assert flow['data']['result'] is None
+
+    await client.send_json({
+        'id': 13,
+        'type': mfa_setup_flow.WS_TYPE_DEPOSE_MFA,
+        'mfa_module_id': 'invalid_id',
+    })
+
+    result = await client.receive_json()
+    assert result['id'] == 13
+    assert result['success'] is False
+    assert result['error']['code'] == 'disable_failed'
+
+    await client.send_json({
+        'id': 14,
+        'type': mfa_setup_flow.WS_TYPE_DEPOSE_MFA,
+        'mfa_module_id': 'example_module',
+    })
+
+    result = await client.receive_json()
+    assert result['id'] == 14
+    assert result['success']
+    assert result['result'] == 'done'