From 92e6abc3945dbad70202fa45779af7f68e5b4fe5 Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Wed, 22 Aug 2018 12:05:55 -0700 Subject: [PATCH 1/8] Add mfa setup flow --- homeassistant/auth/mfa_modules/__init__.py | 43 ++++- .../auth/mfa_modules/insecure_example.py | 2 +- homeassistant/components/auth/__init__.py | 36 +++-- .../components/auth/mfa_setup_flow.py | 153 ++++++++++++++++++ .../auth/mfa_modules/test_insecure_example.py | 18 +++ tests/components/auth/test_mfa_setup_flow.py | 94 +++++++++++ 6 files changed, 333 insertions(+), 13 deletions(-) create mode 100644 homeassistant/components/auth/mfa_setup_flow.py create mode 100644 tests/components/auth/test_mfa_setup_flow.py diff --git a/homeassistant/auth/mfa_modules/__init__.py b/homeassistant/auth/mfa_modules/__init__.py index d0707c4a7452f..d8223e8e4eb5b 100644 --- a/homeassistant/auth/mfa_modules/__init__.py +++ b/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 @@ -72,7 +72,14 @@ def setup_schema(self) -> Optional[vol.Schema]: """ return None - 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, user_id) + + async def async_setup_user(self, user_id: str, setup_data: Any) -> Any: """Set up user for mfa auth module.""" raise NotImplementedError @@ -90,6 +97,38 @@ async def async_validation( raise NotImplementedError +class SetupFlow(data_entry_flow.FlowHandler): + """Handler for the setup flow.""" + + def __init__(self, auth_module: MultiFactorAuthModule, + user_id: str) -> None: + """Initialize the setup flow.""" + self._auth_module = auth_module + self._user_id = user_id + + async def async_step_init(self, user_input=None): + """Handle the first step of setup flow. + + Return self.async_show_form(step_id='init') if user_input == None. + Return await self.async_finish(flow_result) if finish. + """ + errors = {} + + 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._auth_module.setup_schema, + errors=errors + ) + + async def auth_mfa_module_from_config( hass: HomeAssistant, config: Dict[str, Any]) \ -> Optional[MultiFactorAuthModule]: diff --git a/homeassistant/auth/mfa_modules/insecure_example.py b/homeassistant/auth/mfa_modules/insecure_example.py index 59b3f64d2e052..e870fa2ee2785 100644 --- a/homeassistant/auth/mfa_modules/insecure_example.py +++ b/homeassistant/auth/mfa_modules/insecure_example.py @@ -40,7 +40,7 @@ def setup_schema(self) -> Optional[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_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'] diff --git a/homeassistant/components/auth/__init__.py b/homeassistant/components/auth/__init__.py index 4251b23e51475..5f7544181dd05 100644 --- a/homeassistant/components/auth/__init__.py +++ b/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 @@ -316,7 +319,7 @@ def retrieve_result(client_id, result_type, code): @callback -def websocket_current_user(hass, connection, msg): +def websocket_current_user(hass: HomeAssistant, connection, msg): """Return the current user.""" user = connection.request.get('hass_user') @@ -325,11 +328,24 @@ def websocket_current_user(hass, connection, msg): 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.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], + '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(user)) diff --git a/homeassistant/components/auth/mfa_setup_flow.py b/homeassistant/components/auth/mfa_setup_flow.py new file mode 100644 index 0000000000000..62a9a75a9977d --- /dev/null +++ b/homeassistant/components/auth/mfa_setup_flow.py @@ -0,0 +1,153 @@ +"""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 + +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.""" + 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) + + 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) + + +@callback +def websocket_setup_mfa(hass, connection, msg): + """Return a setup flow for mfa auth module.""" + 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 + if user.system_generated: + connection.to_write.put_nowait(websocket_api.error_message( + msg['id'], 'no_system_user', 'System user cannot enable MFA')) + return + + async def async_setup_flow(msg): + """Helper to return a setup flow for mfa auth module.""" + flow_manager = hass.data.get(DATA_SETUP_FLOW_MGR) + if flow_manager is None: + connection.to_write.put_nowait(websocket_api.error_message( + msg['id'], 'not_init', + 'Setup flow manager is not initialized.')) + return + + flow_id = msg.get('flow_id') + if flow_id is not None: + result = await flow_manager.async_configure( + flow_id, msg.get('user_input')) + + else: + 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.to_write.put_nowait(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': user.id}) + + connection.to_write.put_nowait( + websocket_api.result_message( + msg['id'], _prepare_result_json(result))) + + hass.async_add_job(async_setup_flow(msg)) + + +@callback +def websocket_depose_mfa(hass, connection, msg): + """Remove user from mfa module.""" + 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 + if user.system_generated: + connection.to_write.put_nowait(websocket_api.error_message( + msg['id'], 'no_system_user', 'System user cannot enable MFA')) + return + + async def async_depose(msg): + """Helper to disable user from mfa auth module.""" + mfa_module_id = msg['mfa_module_id'] + try: + await hass.auth.async_disable_user_mfa(user, msg['mfa_module_id']) + except Exception as err: + connection.to_write.put_nowait(websocket_api.error_message( + msg['id'], 'disable_failed', + 'Cannot disable Multi-factor Authentication Module' + ' {}: {}'.format(mfa_module_id, err))) + return + + connection.to_write.put_nowait( + websocket_api.result_message( + msg['id'], 'done')) + + hass.async_add_job(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 + + elif 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 diff --git a/tests/auth/mfa_modules/test_insecure_example.py b/tests/auth/mfa_modules/test_insecure_example.py index 9d90532728afa..e6f83762cd770 100644 --- a/tests/auth/mfa_modules/test_insecure_example.py +++ b/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' diff --git a/tests/components/auth/test_mfa_setup_flow.py b/tests/components/auth/test_mfa_setup_flow.py new file mode 100644 index 0000000000000..c2b9ecef2b670 --- /dev/null +++ b/tests/components/auth/test_mfa_setup_flow.py @@ -0,0 +1,94 @@ +"""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['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['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['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['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['success'] + assert result['result'] == 'done' From 7d876590a6174acce319a79d3ad49f4a98a68489 Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Wed, 22 Aug 2018 17:09:53 -0700 Subject: [PATCH 2/8] Lint --- homeassistant/auth/mfa_modules/__init__.py | 6 ++++-- homeassistant/components/auth/__init__.py | 1 - homeassistant/components/auth/mfa_setup_flow.py | 6 +++--- 3 files changed, 7 insertions(+), 6 deletions(-) diff --git a/homeassistant/auth/mfa_modules/__init__.py b/homeassistant/auth/mfa_modules/__init__.py index d8223e8e4eb5b..cec1d29f51f99 100644 --- a/homeassistant/auth/mfa_modules/__init__.py +++ b/homeassistant/auth/mfa_modules/__init__.py @@ -106,13 +106,15 @@ def __init__(self, auth_module: MultiFactorAuthModule, self._auth_module = auth_module self._user_id = user_id - async def async_step_init(self, user_input=None): + 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 await self.async_finish(flow_result) if finish. """ - errors = {} + errors = {} # type: Dict[str, str] if user_input: result = await self._auth_module.async_setup_user( diff --git a/homeassistant/components/auth/__init__.py b/homeassistant/components/auth/__init__.py index 5f7544181dd05..93de70f0ca919 100644 --- a/homeassistant/components/auth/__init__.py +++ b/homeassistant/components/auth/__init__.py @@ -330,7 +330,6 @@ def websocket_current_user(hass: HomeAssistant, connection, msg): async def async_get_current_user(user): """Get current user.""" - enabled_modules = await hass.auth.async_get_enabled_mfa(user) connection.to_write.put_nowait( diff --git a/homeassistant/components/auth/mfa_setup_flow.py b/homeassistant/components/auth/mfa_setup_flow.py index 62a9a75a9977d..3360206a4efd0 100644 --- a/homeassistant/components/auth/mfa_setup_flow.py +++ b/homeassistant/components/auth/mfa_setup_flow.py @@ -35,7 +35,7 @@ async def async_setup(hass): WS_TYPE_DEPOSE_MFA, websocket_depose_mfa, SCHEMA_WS_DEPOSE_MFA) async def _async_create_setup_flow(handler, context, data): - """Create a setup flow. hanlder is a mfa module""" + """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)) @@ -117,7 +117,7 @@ async def async_depose(msg): mfa_module_id = msg['mfa_module_id'] try: await hass.auth.async_disable_user_mfa(user, msg['mfa_module_id']) - except Exception as err: + except ValueError as err: connection.to_write.put_nowait(websocket_api.error_message( msg['id'], 'disable_failed', 'Cannot disable Multi-factor Authentication Module' @@ -137,7 +137,7 @@ def _prepare_result_json(result): data = result.copy() return data - elif result['type'] != data_entry_flow.RESULT_TYPE_FORM: + if result['type'] != data_entry_flow.RESULT_TYPE_FORM: return result import voluptuous_serialize From 1a681173966ce7a6adfa07b2da38d7535e1881d4 Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Thu, 23 Aug 2018 10:27:19 -0700 Subject: [PATCH 3/8] Address code review comment --- homeassistant/components/auth/__init__.py | 18 ++-- .../components/auth/mfa_setup_flow.py | 96 ++++++++----------- homeassistant/components/auth/util.py | 61 ++++++++++++ 3 files changed, 107 insertions(+), 68 deletions(-) create mode 100644 homeassistant/components/auth/util.py diff --git a/homeassistant/components/auth/__init__.py b/homeassistant/components/auth/__init__.py index 93de70f0ca919..7ea69e85591c2 100644 --- a/homeassistant/components/auth/__init__.py +++ b/homeassistant/components/auth/__init__.py @@ -74,9 +74,10 @@ from . import indieauth from . import login_flow from . import mfa_setup_flow +from . import util DOMAIN = 'auth' -DEPENDENCIES = ['http'] +DEPENDENCIES = ['http', 'websocket_api'] WS_TYPE_CURRENT_USER = 'auth/current_user' SCHEMA_WS_CURRENT_USER = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({ @@ -318,21 +319,16 @@ def retrieve_result(client_id, result_type, code): return store_result, retrieve_result +@util.validate_current_user() @callback -def websocket_current_user(hass: HomeAssistant, 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 - async def async_get_current_user(user): """Get current user.""" enabled_modules = await hass.auth.async_get_enabled_mfa(user) - connection.to_write.put_nowait( + connection.send_message_outside( websocket_api.result_message(msg['id'], { 'id': user.id, 'name': user.name, @@ -347,4 +343,4 @@ async def async_get_current_user(user): } for module in hass.auth.auth_mfa_modules], })) - hass.async_create_task(async_get_current_user(user)) + hass.async_create_task(async_get_current_user(connection.user)) diff --git a/homeassistant/components/auth/mfa_setup_flow.py b/homeassistant/components/auth/mfa_setup_flow.py index 3360206a4efd0..52f85036a15dc 100644 --- a/homeassistant/components/auth/mfa_setup_flow.py +++ b/homeassistant/components/auth/mfa_setup_flow.py @@ -5,7 +5,9 @@ from homeassistant import data_entry_flow from homeassistant.components import websocket_api -from homeassistant.core import callback +from homeassistant.core import callback, HomeAssistant + +from . import util WS_TYPE_SETUP_MFA = 'auth/setup_mfa' SCHEMA_WS_SETUP_MFA = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({ @@ -28,12 +30,6 @@ async def async_setup(hass): """Init mfa setup flow manager.""" - 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) - 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) @@ -50,85 +46,71 @@ async def _async_finish_setup_flow(flow, 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 -def websocket_setup_mfa(hass, connection, msg): +@util.validate_current_user(allow_system_user=False) +def websocket_setup_mfa( + hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg): """Return a setup flow for mfa auth module.""" - 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 - if user.system_generated: - connection.to_write.put_nowait(websocket_api.error_message( - msg['id'], 'no_system_user', 'System user cannot enable MFA')) - return - async def async_setup_flow(msg): """Helper to return a setup flow for mfa auth module.""" - flow_manager = hass.data.get(DATA_SETUP_FLOW_MGR) - if flow_manager is None: - connection.to_write.put_nowait(websocket_api.error_message( - msg['id'], 'not_init', - 'Setup flow manager is not initialized.')) - return + 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))) - else: - 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.to_write.put_nowait(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': user.id}) - - connection.to_write.put_nowait( + 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_add_job(async_setup_flow(msg)) + hass.async_create_task(async_setup_flow(msg)) @callback -def websocket_depose_mfa(hass, connection, msg): +@util.validate_current_user(allow_system_user=False) +def websocket_depose_mfa( + hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg): """Remove user from mfa module.""" - 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 - if user.system_generated: - connection.to_write.put_nowait(websocket_api.error_message( - msg['id'], 'no_system_user', 'System user cannot enable MFA')) - return - async def async_depose(msg): """Helper to disable user from mfa auth module.""" mfa_module_id = msg['mfa_module_id'] try: - await hass.auth.async_disable_user_mfa(user, msg['mfa_module_id']) + await hass.auth.async_disable_user_mfa( + connection.user, msg['mfa_module_id']) except ValueError as err: - connection.to_write.put_nowait(websocket_api.error_message( + connection.send_message_outside(websocket_api.error_message( msg['id'], 'disable_failed', - 'Cannot disable Multi-factor Authentication Module' - ' {}: {}'.format(mfa_module_id, err))) + 'Cannot disable MFA Module {}: {}'.format( + mfa_module_id, err))) return - connection.to_write.put_nowait( + connection.send_message_outside( websocket_api.result_message( msg['id'], 'done')) - hass.async_add_job(async_depose(msg)) + hass.async_create_task(async_depose(msg)) def _prepare_result_json(result): diff --git a/homeassistant/components/auth/util.py b/homeassistant/components/auth/util.py new file mode 100644 index 0000000000000..1c4070356c11d --- /dev/null +++ b/homeassistant/components/auth/util.py @@ -0,0 +1,61 @@ +"""Auth component utils.""" +from functools import wraps + +from homeassistant.components import websocket_api +from homeassistant.core import HomeAssistant + + +def validate_current_user( + only_owner=False, only_system_user=False, allow_system_user=True, + only_active_user=True, only_inactive_user=False): + """Decorator that will validate login user exist in current WS connection. + + Will write out error message if not authenticated. + """ + def validator(func): + """Decorator be called.""" + @wraps(func) + def check_current_user(hass: HomeAssistant, + connection: websocket_api.ActiveConnection, + msg): + """Check current user.""" + def output_error(message_id, message): + """Output error message.""" + connection.send_message_outside(websocket_api.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 From 8f7c5b250c42d541343561d31a865959c0b9823f Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Thu, 23 Aug 2018 13:20:49 -0700 Subject: [PATCH 4/8] Fix unit test --- tests/components/auth/test_mfa_setup_flow.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/components/auth/test_mfa_setup_flow.py b/tests/components/auth/test_mfa_setup_flow.py index c2b9ecef2b670..bcc3eac1bfd96 100644 --- a/tests/components/auth/test_mfa_setup_flow.py +++ b/tests/components/auth/test_mfa_setup_flow.py @@ -81,7 +81,7 @@ async def test_ws_setup_depose_mfa(hass, hass_ws_client): result = await client.receive_json() assert result['success'] is False - assert result['error']['code'] == 'disable_failed' + assert result['error']['code'] == 'no_module' await client.send_json({ 'id': 14, From d4e01f94424f7cc9391b8a2d61f56d3f04b05f4e Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Thu, 23 Aug 2018 22:01:14 +0000 Subject: [PATCH 5/8] Add assertion for WS response ordering --- tests/components/auth/test_mfa_setup_flow.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/components/auth/test_mfa_setup_flow.py b/tests/components/auth/test_mfa_setup_flow.py index bcc3eac1bfd96..93b5cdf7bb919 100644 --- a/tests/components/auth/test_mfa_setup_flow.py +++ b/tests/components/auth/test_mfa_setup_flow.py @@ -40,6 +40,7 @@ async def test_ws_setup_depose_mfa(hass, hass_ws_client): }) result = await client.receive_json() + assert result['id'] == 10 assert result['success'] is False assert result['error']['code'] == 'no_module' @@ -50,6 +51,7 @@ async def test_ws_setup_depose_mfa(hass, hass_ws_client): }) result = await client.receive_json() + assert result['id'] == 11 assert result['success'] flow = result['result'] @@ -66,6 +68,7 @@ async def test_ws_setup_depose_mfa(hass, hass_ws_client): }) result = await client.receive_json() + assert result['id'] == 12 assert result['success'] flow = result['result'] @@ -80,8 +83,9 @@ async def test_ws_setup_depose_mfa(hass, hass_ws_client): }) result = await client.receive_json() + assert result['id'] == 13 assert result['success'] is False - assert result['error']['code'] == 'no_module' + assert result['error']['code'] == 'disable_failed' await client.send_json({ 'id': 14, @@ -90,5 +94,6 @@ async def test_ws_setup_depose_mfa(hass, hass_ws_client): }) result = await client.receive_json() + assert result['id'] == 14 assert result['success'] assert result['result'] == 'done' From af963a37d40eb8a078b08d2ed13c212cd544a96c Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Thu, 23 Aug 2018 22:46:07 +0000 Subject: [PATCH 6/8] Missed a return --- homeassistant/components/auth/mfa_setup_flow.py | 1 + 1 file changed, 1 insertion(+) diff --git a/homeassistant/components/auth/mfa_setup_flow.py b/homeassistant/components/auth/mfa_setup_flow.py index 52f85036a15dc..0d287d1dd30e8 100644 --- a/homeassistant/components/auth/mfa_setup_flow.py +++ b/homeassistant/components/auth/mfa_setup_flow.py @@ -69,6 +69,7 @@ async def async_setup_flow(msg): 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) From 409a76dfa0bf0fa52f7f5063201f970c57dbedfd Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Fri, 24 Aug 2018 02:39:24 -0700 Subject: [PATCH 7/8] Remove setup_schema from MFA base class --- homeassistant/auth/__init__.py | 7 ------- homeassistant/auth/mfa_modules/__init__.py | 14 ++++---------- homeassistant/auth/mfa_modules/insecure_example.py | 11 +++++++++-- 3 files changed, 13 insertions(+), 19 deletions(-) diff --git a/homeassistant/auth/__init__.py b/homeassistant/auth/__init__.py index c0beba1a22760..5f50bfd5fb3fc 100644 --- a/homeassistant/auth/__init__.py +++ b/homeassistant/auth/__init__.py @@ -235,13 +235,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, diff --git a/homeassistant/auth/mfa_modules/__init__.py b/homeassistant/auth/mfa_modules/__init__.py index cec1d29f51f99..a250d83eed56e 100644 --- a/homeassistant/auth/mfa_modules/__init__.py +++ b/homeassistant/auth/mfa_modules/__init__.py @@ -64,20 +64,12 @@ 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. - - Optional - """ - return 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, user_id) + raise NotImplementedError async def async_setup_user(self, user_id: str, setup_data: Any) -> Any: """Set up user for mfa auth module.""" @@ -101,9 +93,11 @@ 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( @@ -126,7 +120,7 @@ async def async_step_init( return self.async_show_form( step_id='init', - data_schema=self._auth_module.setup_schema, + data_schema=self._setup_schema, errors=errors ) diff --git a/homeassistant/auth/mfa_modules/insecure_example.py b/homeassistant/auth/mfa_modules/insecure_example.py index e870fa2ee2785..8bd68c565ac6e 100644 --- a/homeassistant/auth/mfa_modules/insecure_example.py +++ b/homeassistant/auth/mfa_modules/insecure_example.py @@ -7,7 +7,7 @@ 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,10 +36,17 @@ 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_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 From 58e6b4a8be090725c3b1ee0b2a46271cbf7a99d5 Mon Sep 17 00:00:00 2001 From: Jason Hu Date: Fri, 24 Aug 2018 09:52:24 -0700 Subject: [PATCH 8/8] Move auth.util.validate_current_user -> webscoket_api.ws_require_user --- homeassistant/auth/__init__.py | 2 - homeassistant/auth/mfa_modules/__init__.py | 2 +- .../auth/mfa_modules/insecure_example.py | 2 +- homeassistant/components/auth/__init__.py | 5 +- .../components/auth/mfa_setup_flow.py | 10 ++- homeassistant/components/auth/util.py | 61 ------------------- homeassistant/components/websocket_api.py | 58 +++++++++++++++++- 7 files changed, 65 insertions(+), 75 deletions(-) delete mode 100644 homeassistant/components/auth/util.py diff --git a/homeassistant/auth/__init__.py b/homeassistant/auth/__init__.py index 5f50bfd5fb3fc..0cd638b4e9c8e 100644 --- a/homeassistant/auth/__init__.py +++ b/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 diff --git a/homeassistant/auth/mfa_modules/__init__.py b/homeassistant/auth/mfa_modules/__init__.py index a250d83eed56e..cb0758e3ef8c0 100644 --- a/homeassistant/auth/mfa_modules/__init__.py +++ b/homeassistant/auth/mfa_modules/__init__.py @@ -106,7 +106,7 @@ async def async_step_init( """Handle the first step of setup flow. Return self.async_show_form(step_id='init') if user_input == None. - Return await self.async_finish(flow_result) if finish. + Return self.async_create_entry(data={'result': result}) if finish. """ errors = {} # type: Dict[str, str] diff --git a/homeassistant/auth/mfa_modules/insecure_example.py b/homeassistant/auth/mfa_modules/insecure_example.py index 8bd68c565ac6e..9c72111ef9697 100644 --- a/homeassistant/auth/mfa_modules/insecure_example.py +++ b/homeassistant/auth/mfa_modules/insecure_example.py @@ -1,6 +1,6 @@ """Example auth module.""" import logging -from typing import Any, Dict, Optional +from typing import Any, Dict import voluptuous as vol diff --git a/homeassistant/components/auth/__init__.py b/homeassistant/components/auth/__init__.py index 7ea69e85591c2..a87e646761c47 100644 --- a/homeassistant/components/auth/__init__.py +++ b/homeassistant/components/auth/__init__.py @@ -74,10 +74,9 @@ from . import indieauth from . import login_flow from . import mfa_setup_flow -from . import util DOMAIN = 'auth' -DEPENDENCIES = ['http', 'websocket_api'] +DEPENDENCIES = ['http'] WS_TYPE_CURRENT_USER = 'auth/current_user' SCHEMA_WS_CURRENT_USER = websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({ @@ -319,7 +318,7 @@ def retrieve_result(client_id, result_type, code): return store_result, retrieve_result -@util.validate_current_user() +@websocket_api.ws_require_user() @callback def websocket_current_user( hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg): diff --git a/homeassistant/components/auth/mfa_setup_flow.py b/homeassistant/components/auth/mfa_setup_flow.py index 0d287d1dd30e8..82eb913d89082 100644 --- a/homeassistant/components/auth/mfa_setup_flow.py +++ b/homeassistant/components/auth/mfa_setup_flow.py @@ -7,8 +7,6 @@ from homeassistant.components import websocket_api from homeassistant.core import callback, HomeAssistant -from . import util - 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, @@ -54,12 +52,12 @@ async def _async_finish_setup_flow(flow, flow_result): @callback -@util.validate_current_user(allow_system_user=False) +@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): - """Helper to return a setup flow for mfa auth module.""" + """Return a setup flow for mfa auth module.""" flow_manager = hass.data[DATA_SETUP_FLOW_MGR] flow_id = msg.get('flow_id') @@ -90,12 +88,12 @@ async def async_setup_flow(msg): @callback -@util.validate_current_user(allow_system_user=False) +@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): - """Helper to disable user from mfa auth module.""" + """Remove user from mfa auth module.""" mfa_module_id = msg['mfa_module_id'] try: await hass.auth.async_disable_user_mfa( diff --git a/homeassistant/components/auth/util.py b/homeassistant/components/auth/util.py deleted file mode 100644 index 1c4070356c11d..0000000000000 --- a/homeassistant/components/auth/util.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Auth component utils.""" -from functools import wraps - -from homeassistant.components import websocket_api -from homeassistant.core import HomeAssistant - - -def validate_current_user( - only_owner=False, only_system_user=False, allow_system_user=True, - only_active_user=True, only_inactive_user=False): - """Decorator that will validate login user exist in current WS connection. - - Will write out error message if not authenticated. - """ - def validator(func): - """Decorator be called.""" - @wraps(func) - def check_current_user(hass: HomeAssistant, - connection: websocket_api.ActiveConnection, - msg): - """Check current user.""" - def output_error(message_id, message): - """Output error message.""" - connection.send_message_outside(websocket_api.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 diff --git a/homeassistant/components/websocket_api.py b/homeassistant/components/websocket_api.py index 1ba0e20d55345..0c9ab366534f2 100644 --- a/homeassistant/components/websocket_api.py +++ b/homeassistant/components/websocket_api.py @@ -18,7 +18,7 @@ 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