"""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, private=True) 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( 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( _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, )