-
-
Notifications
You must be signed in to change notification settings - Fork 38.3k
Add Notify MFA module #16314
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add Notify MFA module #16314
Changes from all commits
08af92d
516671e
6c5536b
9984fa5
2d8e6da
0bf7bc6
e2da768
fd65a1b
471b73a
942573d
8a87465
24f10cd
fd20a9b
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,325 @@ | ||
| """HMAC-based One-time Password auth module. | ||
|
|
||
| Sending HOTP through notify service | ||
| """ | ||
| import logging | ||
| from collections import OrderedDict | ||
| from typing import Any, Dict, Optional, Tuple, List # noqa: F401 | ||
|
|
||
| import attr | ||
| import voluptuous as vol | ||
|
|
||
| from homeassistant.const import CONF_EXCLUDE, CONF_INCLUDE | ||
| from homeassistant.core import HomeAssistant, callback | ||
| from homeassistant.helpers import config_validation as cv | ||
|
|
||
| from . import MultiFactorAuthModule, MULTI_FACTOR_AUTH_MODULES, \ | ||
| MULTI_FACTOR_AUTH_MODULE_SCHEMA, SetupFlow | ||
|
|
||
| REQUIREMENTS = ['pyotp==2.2.6'] | ||
|
|
||
| CONF_MESSAGE = 'message' | ||
|
|
||
| CONFIG_SCHEMA = MULTI_FACTOR_AUTH_MODULE_SCHEMA.extend({ | ||
| vol.Optional(CONF_INCLUDE): vol.All(cv.ensure_list, [cv.string]), | ||
| vol.Optional(CONF_EXCLUDE): vol.All(cv.ensure_list, [cv.string]), | ||
| vol.Optional(CONF_MESSAGE, | ||
| default='{} is your Home Assistant login code'): str | ||
| }, extra=vol.PREVENT_EXTRA) | ||
|
|
||
| STORAGE_VERSION = 1 | ||
| STORAGE_KEY = 'auth_module.notify' | ||
| STORAGE_USERS = 'users' | ||
| STORAGE_USER_ID = 'user_id' | ||
|
|
||
| INPUT_FIELD_CODE = 'code' | ||
|
|
||
| _LOGGER = logging.getLogger(__name__) | ||
|
|
||
|
|
||
| def _generate_secret() -> str: | ||
| """Generate a secret.""" | ||
| import pyotp | ||
| return str(pyotp.random_base32()) | ||
|
|
||
|
|
||
| def _generate_random() -> int: | ||
| """Generate a 8 digit number.""" | ||
| import pyotp | ||
| return int(pyotp.random_base32(length=8, chars=list('1234567890'))) | ||
|
|
||
|
|
||
| def _generate_otp(secret: str, count: int) -> str: | ||
| """Generate one time password.""" | ||
| import pyotp | ||
| return str(pyotp.HOTP(secret).at(count)) | ||
|
|
||
|
|
||
| def _verify_otp(secret: str, otp: str, count: int) -> bool: | ||
| """Verify one time password.""" | ||
| import pyotp | ||
| return bool(pyotp.HOTP(secret).verify(otp, count)) | ||
|
|
||
|
|
||
| @attr.s(slots=True) | ||
| class NotifySetting: | ||
| """Store notify setting for one user.""" | ||
|
|
||
| secret = attr.ib(type=str, factory=_generate_secret) # not persistent | ||
| counter = attr.ib(type=int, factory=_generate_random) # not persistent | ||
| notify_service = attr.ib(type=Optional[str], default=None) | ||
| target = attr.ib(type=Optional[str], default=None) | ||
|
|
||
|
|
||
| _UsersDict = Dict[str, NotifySetting] | ||
|
|
||
|
|
||
| @MULTI_FACTOR_AUTH_MODULES.register('notify') | ||
| class NotifyAuthModule(MultiFactorAuthModule): | ||
| """Auth module send hmac-based one time password by notify service.""" | ||
|
|
||
| DEFAULT_TITLE = 'Notify One-Time Password' | ||
|
|
||
| def __init__(self, hass: HomeAssistant, config: Dict[str, Any]) -> None: | ||
| """Initialize the user data store.""" | ||
| super().__init__(hass, config) | ||
| self._user_settings = None # type: Optional[_UsersDict] | ||
| self._user_store = hass.helpers.storage.Store( | ||
| STORAGE_VERSION, STORAGE_KEY) | ||
| self._include = config.get(CONF_INCLUDE, []) | ||
| self._exclude = config.get(CONF_EXCLUDE, []) | ||
| self._message_template = config[CONF_MESSAGE] | ||
|
|
||
| @property | ||
| def input_schema(self) -> vol.Schema: | ||
| """Validate login flow input data.""" | ||
| return vol.Schema({INPUT_FIELD_CODE: str}) | ||
|
|
||
| async def _async_load(self) -> None: | ||
| """Load stored data.""" | ||
| data = await self._user_store.async_load() | ||
|
|
||
| if data is None: | ||
| data = {STORAGE_USERS: {}} | ||
|
|
||
| self._user_settings = { | ||
| user_id: NotifySetting(**setting) | ||
| for user_id, setting in data.get(STORAGE_USERS, {}).items() | ||
| } | ||
|
|
||
| async def _async_save(self) -> None: | ||
| """Save data.""" | ||
| if self._user_settings is None: | ||
| return | ||
|
|
||
| await self._user_store.async_save({STORAGE_USERS: { | ||
| user_id: attr.asdict( | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It's shorter code to just write the dict out |
||
| notify_setting, filter=attr.filters.exclude( | ||
| attr.fields(NotifySetting).secret, | ||
| attr.fields(NotifySetting).counter, | ||
| )) | ||
| for user_id, notify_setting | ||
| in self._user_settings.items() | ||
| }}) | ||
|
|
||
| @callback | ||
| def aync_get_available_notify_services(self) -> List[str]: | ||
| """Return list of notify services.""" | ||
| unordered_services = set() | ||
|
|
||
| for service in self.hass.services.async_services().get('notify', {}): | ||
| if service not in self._exclude: | ||
| unordered_services.add(service) | ||
|
|
||
| if self._include: | ||
| unordered_services &= set(self._include) | ||
|
|
||
| return sorted(unordered_services) | ||
|
|
||
| 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 NotifySetupFlow( | ||
| self, self.input_schema, user_id, | ||
| self.aync_get_available_notify_services()) | ||
|
|
||
| async def async_setup_user(self, user_id: str, setup_data: Any) -> Any: | ||
| """Set up auth module for user.""" | ||
| if self._user_settings is None: | ||
| await self._async_load() | ||
| assert self._user_settings is not None | ||
|
|
||
| self._user_settings[user_id] = NotifySetting( | ||
| notify_service=setup_data.get('notify_service'), | ||
| target=setup_data.get('target'), | ||
| ) | ||
|
|
||
| await self._async_save() | ||
|
|
||
| async def async_depose_user(self, user_id: str) -> None: | ||
| """Depose auth module for user.""" | ||
| if self._user_settings is None: | ||
| await self._async_load() | ||
| assert self._user_settings is not None | ||
|
|
||
| if self._user_settings.pop(user_id, None): | ||
| await self._async_save() | ||
|
|
||
| async def async_is_user_setup(self, user_id: str) -> bool: | ||
| """Return whether user is setup.""" | ||
| if self._user_settings is None: | ||
| await self._async_load() | ||
| assert self._user_settings is not None | ||
|
|
||
| return user_id in self._user_settings | ||
|
|
||
| async def async_validate( | ||
| self, user_id: str, user_input: Dict[str, Any]) -> bool: | ||
| """Return True if validation passed.""" | ||
| if self._user_settings is None: | ||
| await self._async_load() | ||
| assert self._user_settings is not None | ||
|
|
||
| notify_setting = self._user_settings.get(user_id, None) | ||
| if notify_setting is None: | ||
| return False | ||
|
|
||
| # user_input has been validate in caller | ||
| return await self.hass.async_add_executor_job( | ||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I wonder if we did the right thing with MFA architecture. Would it be possible to get an instance of a flow which can store instance variables instead of storing temp values in our storage object. The latter feels hacky.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I even wrote a version of that somewhere in the history. The reason I give it up because consider following scenario:
I had seen this on other systems, most of them won't allow you to use the first code anyway, for that case, the code is associated with user, not login flow instance.
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I think it's fine to have a code be bound to a login flow. That's what I would have expected. |
||
| _verify_otp, notify_setting.secret, | ||
| user_input.get(INPUT_FIELD_CODE, ''), | ||
| notify_setting.counter) | ||
|
|
||
| async def async_initialize_login_mfa_step(self, user_id: str) -> None: | ||
| """Generate code and notify user.""" | ||
| if self._user_settings is None: | ||
| await self._async_load() | ||
| assert self._user_settings is not None | ||
|
|
||
| notify_setting = self._user_settings.get(user_id, None) | ||
| if notify_setting is None: | ||
| raise ValueError('Cannot find user_id') | ||
|
|
||
| def generate_secret_and_one_time_password() -> str: | ||
| """Generate and send one time password.""" | ||
| assert notify_setting | ||
| # secret and counter are not persistent | ||
| notify_setting.secret = _generate_secret() | ||
| notify_setting.counter = _generate_random() | ||
| return _generate_otp( | ||
| notify_setting.secret, notify_setting.counter) | ||
|
|
||
| code = await self.hass.async_add_executor_job( | ||
| generate_secret_and_one_time_password) | ||
|
|
||
| await self.async_notify_user(user_id, code) | ||
|
|
||
| async def async_notify_user(self, user_id: str, code: str) -> None: | ||
| """Send code by user's notify service.""" | ||
| if self._user_settings is None: | ||
| await self._async_load() | ||
| assert self._user_settings is not None | ||
|
|
||
| notify_setting = self._user_settings.get(user_id, None) | ||
| if notify_setting is None: | ||
| _LOGGER.error('Cannot find user %s', user_id) | ||
| return | ||
|
|
||
| await self.async_notify( # type: ignore | ||
| code, notify_setting.notify_service, notify_setting.target) | ||
|
|
||
| async def async_notify(self, code: str, notify_service: str, | ||
| target: Optional[str] = None) -> None: | ||
| """Send code by notify service.""" | ||
| data = {'message': self._message_template.format(code)} | ||
| if target: | ||
| data['target'] = [target] | ||
|
|
||
| await self.hass.services.async_call('notify', notify_service, data) | ||
|
|
||
|
|
||
| class NotifySetupFlow(SetupFlow): | ||
| """Handler for the setup flow.""" | ||
|
|
||
| def __init__(self, auth_module: NotifyAuthModule, | ||
| setup_schema: vol.Schema, | ||
| user_id: str, | ||
| available_notify_services: List[str]) -> None: | ||
| """Initialize the setup flow.""" | ||
| super().__init__(auth_module, setup_schema, user_id) | ||
| # to fix typing complaint | ||
| self._auth_module = auth_module # type: NotifyAuthModule | ||
| self._available_notify_services = available_notify_services | ||
| self._secret = None # type: Optional[str] | ||
| self._count = None # type: Optional[int] | ||
| self._notify_service = None # type: Optional[str] | ||
| self._target = None # type: Optional[str] | ||
|
|
||
| async def async_step_init( | ||
| self, user_input: Optional[Dict[str, str]] = None) \ | ||
| -> Dict[str, Any]: | ||
| """Let user select available notify services.""" | ||
| errors = {} # type: Dict[str, str] | ||
|
|
||
| hass = self._auth_module.hass | ||
| if user_input: | ||
| self._notify_service = user_input['notify_service'] | ||
| self._target = user_input.get('target') | ||
| self._secret = await hass.async_add_executor_job(_generate_secret) | ||
| self._count = await hass.async_add_executor_job(_generate_random) | ||
|
|
||
| return await self.async_step_setup() | ||
|
|
||
| if not self._available_notify_services: | ||
| return self.async_abort(reason='no_available_service') | ||
|
|
||
| schema = OrderedDict() # type: Dict[str, Any] | ||
| schema['notify_service'] = vol.In(self._available_notify_services) | ||
| schema['target'] = vol.Optional(str) | ||
|
|
||
| return self.async_show_form( | ||
| step_id='init', | ||
| data_schema=vol.Schema(schema), | ||
| errors=errors | ||
| ) | ||
|
|
||
| async def async_step_setup( | ||
| self, user_input: Optional[Dict[str, str]] = None) \ | ||
| -> Dict[str, Any]: | ||
| """Verify user can recevie one-time password.""" | ||
| errors = {} # type: Dict[str, str] | ||
|
|
||
| hass = self._auth_module.hass | ||
| if user_input: | ||
| verified = await hass.async_add_executor_job( | ||
| _verify_otp, self._secret, user_input['code'], self._count) | ||
| if verified: | ||
| await self._auth_module.async_setup_user( | ||
| self._user_id, { | ||
| 'notify_service': self._notify_service, | ||
| 'target': self._target, | ||
| }) | ||
| return self.async_create_entry( | ||
| title=self._auth_module.name, | ||
| data={} | ||
| ) | ||
|
|
||
| errors['base'] = 'invalid_code' | ||
|
|
||
| # generate code every time, no retry logic | ||
| assert self._secret and self._count | ||
| code = await hass.async_add_executor_job( | ||
| _generate_otp, self._secret, self._count) | ||
|
|
||
| assert self._notify_service | ||
| await self._auth_module.async_notify( | ||
| code, self._notify_service, self._target) | ||
|
|
||
| return self.async_show_form( | ||
| step_id='setup', | ||
| data_schema=self._setup_schema, | ||
| description_placeholders={'notify_service': self._notify_service}, | ||
| errors=errors, | ||
| ) | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If it's not persistent, we shouldn't put it on the storage class.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Because instantiating it will have the random side effect of generating the secret / random. Meaning during load, all these things get generated without maybe not being needed.