Add multi-factor authentication modules (#15489)

* Get user after login flow finished

* Add multi factor authentication support

* Typings

homeassistant/auth/__init__.py

@@ -6,21 +6,26 @@
 
 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
 
 from . import auth_store, models
-from .providers import auth_provider_from_config, AuthProvider
+from .mfa_modules import auth_mfa_module_from_config, MultiFactorAuthModule
+from .providers import auth_provider_from_config, AuthProvider, LoginFlow
 
 _LOGGER = logging.getLogger(__name__)
+_MfaModuleDict = Dict[str, MultiFactorAuthModule]
 _ProviderKey = Tuple[str, Optional[str]]
 _ProviderDict = Dict[_ProviderKey, AuthProvider]
 
 
 async def auth_manager_from_config(
         hass: HomeAssistant,
-        provider_configs: List[Dict[str, Any]]) -> 'AuthManager':
+        provider_configs: List[Dict[str, Any]],
+        module_configs: List[Dict[str, Any]]) -> 'AuthManager':
     """Initialize an auth manager from config."""
     store = auth_store.AuthStore(hass)
     if provider_configs:
@@ -44,18 +49,42 @@ async def auth_manager_from_config(
             continue
 
         provider_hash[key] = provider
-    manager = AuthManager(hass, store, provider_hash)
+
+    if module_configs:
+        modules = await asyncio.gather(
+            *[auth_mfa_module_from_config(hass, config)
+              for config in module_configs])
+    else:
+        modules = ()
+    # So returned auth modules are in same order as config
+    module_hash = OrderedDict()  # type: _MfaModuleDict
+    for module in modules:
+        if module is None:
+            continue
+
+        if module.id in module_hash:
+            _LOGGER.error(
+                'Found duplicate multi-factor module: %s. Please add unique '
+                'IDs if you want to have the same module twice.', module.id)
+            continue
+
+        module_hash[module.id] = module
+
+    manager = AuthManager(hass, store, provider_hash, module_hash)
     return manager
 
 
 class AuthManager:
     """Manage the authentication for Home Assistant."""
 
     def __init__(self, hass: HomeAssistant, store: auth_store.AuthStore,
-                 providers: _ProviderDict) -> None:
+                 providers: _ProviderDict, mfa_modules: _MfaModuleDict) \
+            -> None:
         """Initialize the auth manager."""
+        self.hass = hass
         self._store = store
         self._providers = providers
+        self._mfa_modules = mfa_modules
         self.login_flow = data_entry_flow.FlowManager(
             hass, self._async_create_login_flow,
             self._async_finish_login_flow)
@@ -82,6 +111,16 @@ def auth_providers(self) -> List[AuthProvider]:
         """Return a list of available auth providers."""
         return list(self._providers.values())
 
+    @property
+    def auth_mfa_modules(self) -> List[MultiFactorAuthModule]:
+        """Return a list of available auth modules."""
+        return list(self._mfa_modules.values())
+
+    def get_auth_mfa_module(self, module_id: str) \
+            -> Optional[MultiFactorAuthModule]:
+        """Return an multi-factor auth module, None if not found."""
+        return self._mfa_modules.get(module_id)
+
     async def async_get_users(self) -> List[models.User]:
         """Retrieve all users."""
         return await self._store.async_get_users()
@@ -90,6 +129,16 @@ async def async_get_user(self, user_id: str) -> Optional[models.User]:
         """Retrieve a user."""
         return await self._store.async_get_user(user_id)
 
+    async def async_get_user_by_credentials(
+            self, credentials: models.Credentials) -> Optional[models.User]:
+        """Get a user by credential, return None if not found."""
+        for user in await self.async_get_users():
+            for creds in user.credentials:
+                if creds.id == credentials.id:
+                    return user
+
+        return None
+
     async def async_create_system_user(self, name: str) -> models.User:
         """Create a system user."""
         return await self._store.async_create_user(
@@ -114,12 +163,11 @@ async def async_get_or_create_user(self, credentials: models.Credentials) \
             -> models.User:
         """Get or create a user."""
         if not credentials.is_new:
-            for user in await self._store.async_get_users():
-                for creds in user.credentials:
-                    if creds.id == credentials.id:
-                        return user
-
-            raise ValueError('Unable to find the user.')
+            user = await self.async_get_user_by_credentials(credentials)
+            if user is None:
+                raise ValueError('Unable to find the user.')
+            else:
+                return user
 
         auth_provider = self._async_get_auth_provider(credentials)
 
@@ -175,6 +223,49 @@ async def async_remove_credentials(
 
         await self._store.async_remove_credentials(credentials)
 
+    async def async_enable_user_mfa(self, user: models.User,
+                                    mfa_module_id: str, data: Any) -> None:
+        """Enable a multi-factor auth module for user."""
+        if user.system_generated:
+            raise ValueError('System generated users cannot enable '
+                             'multi-factor auth module.')
+
+        module = self.get_auth_mfa_module(mfa_module_id)
+        if module is None:
+            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,
+                                     mfa_module_id: str) -> None:
+        """Disable a multi-factor auth module for user."""
+        if user.system_generated:
+            raise ValueError('System generated users cannot disable '
+                             'multi-factor auth module.')
+
+        module = self.get_auth_mfa_module(mfa_module_id)
+        if module is None:
+            raise ValueError('Unable find multi-factor auth module: {}'
+                             .format(mfa_module_id))
+
+        await module.async_depose_user(user.id)
+
+    async def async_get_enabled_mfa(self, user: models.User) -> List[str]:
+        """List enabled mfa modules for user."""
+        module_ids = []
+        for module_id, module in self._mfa_modules.items():
+            if await module.async_is_user_setup(user.id):
+                module_ids.append(module_id)
+        return module_ids
+
     async def async_create_refresh_token(self, user: models.User,
                                          client_id: Optional[str] = None) \
             -> models.RefreshToken:
@@ -262,12 +353,17 @@ async def _async_create_login_flow(
         return await auth_provider.async_login_flow(context)
 
     async def _async_finish_login_flow(
-            self, flow: data_entry_flow.FlowHandler, result: Dict[str, Any]) \
+            self, flow: LoginFlow, result: Dict[str, Any]) \
             -> Dict[str, Any]:
         """Return a user as result of login flow."""
         if result['type'] != data_entry_flow.RESULT_TYPE_CREATE_ENTRY:
             return result
 
+        # we got final result
+        if isinstance(result['data'], models.User):
+            result['result'] = result['data']
+            return result
+
         auth_provider = self._providers[result['handler']]
         credentials = await auth_provider.async_get_or_create_credentials(
             result['data'])
@@ -276,8 +372,19 @@ async def _async_finish_login_flow(
             result['result'] = credentials
             return result
 
-        user = await self.async_get_or_create_user(credentials)
-        result['result'] = user
+        # multi-factor module cannot enabled for new credential
+        # which has not linked to a user yet
+        if auth_provider.support_mfa and not credentials.is_new:
+            user = await self.async_get_user_by_credentials(credentials)
+            if user is not None:
+                modules = await self.async_get_enabled_mfa(user)
+
+                if modules:
+                    flow.user = user
+                    flow.available_mfa_modules = modules
+                    return await flow.async_step_select_mfa_module()
+
+        result['result'] = await self.async_get_or_create_user(credentials)
         return result
 
     @callback

homeassistant/auth/mfa_modules/__init__.py

@@ -0,0 +1,141 @@
+"""Plugable auth modules for Home Assistant."""
+from datetime import timedelta
+import importlib
+import logging
+import types
+from typing import Any, Dict, Optional
+
+import voluptuous as vol
+from voluptuous.humanize import humanize_error
+
+from homeassistant import requirements
+from homeassistant.const import CONF_ID, CONF_NAME, CONF_TYPE
+from homeassistant.core import HomeAssistant
+from homeassistant.util.decorator import Registry
+
+MULTI_FACTOR_AUTH_MODULES = Registry()
+
+MULTI_FACTOR_AUTH_MODULE_SCHEMA = vol.Schema({
+    vol.Required(CONF_TYPE): str,
+    vol.Optional(CONF_NAME): str,
+    # Specify ID if you have two mfa auth module for same type.
+    vol.Optional(CONF_ID): str,
+}, extra=vol.ALLOW_EXTRA)
+
+SESSION_EXPIRATION = timedelta(minutes=5)
+
+DATA_REQS = 'mfa_auth_module_reqs_processed'
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class MultiFactorAuthModule:
+    """Multi-factor Auth Module of validation function."""
+
+    DEFAULT_TITLE = 'Unnamed auth module'
+
+    def __init__(self, hass: HomeAssistant, config: Dict[str, Any]) -> None:
+        """Initialize an auth module."""
+        self.hass = hass
+        self.config = config
+
+    @property
+    def id(self) -> str:  # pylint: disable=invalid-name
+        """Return id of the auth module.
+
+        Default is same as type
+        """
+        return self.config.get(CONF_ID, self.type)
+
+    @property
+    def type(self) -> str:
+        """Return type of the module."""
+        return self.config[CONF_TYPE]  # type: ignore
+
+    @property
+    def name(self) -> str:
+        """Return the name of the auth module."""
+        return self.config.get(CONF_NAME, self.DEFAULT_TITLE)
+
+    # Implement by extending class
+
+    @property
+    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_user(self, user_id: str, setup_data: Any) -> None:
+        """Set up user for mfa auth module."""
+        raise NotImplementedError
+
+    async def async_depose_user(self, user_id: str) -> None:
+        """Remove user from mfa module."""
+        raise NotImplementedError
+
+    async def async_is_user_setup(self, user_id: str) -> bool:
+        """Return whether user is setup."""
+        raise NotImplementedError
+
+    async def async_validation(
+            self, user_id: str, user_input: Dict[str, Any]) -> bool:
+        """Return True if validation passed."""
+        raise NotImplementedError
+
+
+async def auth_mfa_module_from_config(
+        hass: HomeAssistant, config: Dict[str, Any]) \
+        -> Optional[MultiFactorAuthModule]:
+    """Initialize an auth module from a config."""
+    module_name = config[CONF_TYPE]
+    module = await _load_mfa_module(hass, module_name)
+
+    if module is None:
+        return None
+
+    try:
+        config = module.CONFIG_SCHEMA(config)  # type: ignore
+    except vol.Invalid as err:
+        _LOGGER.error('Invalid configuration for multi-factor module %s: %s',
+                      module_name, humanize_error(config, err))
+        return None
+
+    return MULTI_FACTOR_AUTH_MODULES[module_name](hass, config)  # type: ignore
+
+
+async def _load_mfa_module(hass: HomeAssistant, module_name: str) \
+        -> Optional[types.ModuleType]:
+    """Load an mfa auth module."""
+    module_path = 'homeassistant.auth.mfa_modules.{}'.format(module_name)
+
+    try:
+        module = importlib.import_module(module_path)
+    except ImportError:
+        _LOGGER.warning('Unable to find %s', module_path)
+        return None
+
+    if hass.config.skip_pip or not hasattr(module, 'REQUIREMENTS'):
+        return module
+
+    processed = hass.data.get(DATA_REQS)
+    if processed and module_name in processed:
+        return module
+
+    processed = hass.data[DATA_REQS] = set()
+
+    # https://github.com/python/mypy/issues/1424
+    req_success = await requirements.async_process_requirements(
+        hass, module_path, module.REQUIREMENTS)    # type: ignore
+
+    if not req_success:
+        return None
+
+    processed.add(module_name)
+    return module