326 lines
11 KiB
Python
326 lines
11 KiB
Python
"""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,
|
|
)
|