Skip to content

Commit

Permalink
Add Time-based Onetime Password Multi-factor Auth
Browse files Browse the repository at this point in the history
Add TOTP setup flow, generate QR code
  • Loading branch information
awarecan committed Aug 24, 2018
1 parent e8775ba commit 37842c2
Show file tree
Hide file tree
Showing 10 changed files with 435 additions and 7 deletions.
8 changes: 4 additions & 4 deletions homeassistant/auth/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -249,13 +249,13 @@ async def async_disable_user_mfa(self, user: models.User,

await module.async_depose_user(user.id)

async def async_get_enabled_mfa(self, user: models.User) -> List[str]:
async def async_get_enabled_mfa(self, user: models.User) -> Dict[str, str]:
"""List enabled mfa modules for user."""
module_ids = []
modules = OrderedDict()
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
modules[module_id] = module.name
return modules

async def async_create_refresh_token(self, user: models.User,
client_id: Optional[str] = None) \
Expand Down
209 changes: 209 additions & 0 deletions homeassistant/auth/mfa_modules/totp.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,209 @@
"""Time-based One Time Password auth module."""
import logging
from typing import Any, Dict, Optional, Tuple # noqa: F401

import voluptuous as vol

from homeassistant.auth.models import User
from homeassistant.core import HomeAssistant
from homeassistant.loader import bind_hass

from . import MultiFactorAuthModule, MULTI_FACTOR_AUTH_MODULES, \
MULTI_FACTOR_AUTH_MODULE_SCHEMA, SetupFlow

REQUIREMENTS = ['pyotp==2.2.6', 'PyQRCode==1.2.1', 'pypng==0.0.18']

CONFIG_SCHEMA = MULTI_FACTOR_AUTH_MODULE_SCHEMA.extend({
}, extra=vol.PREVENT_EXTRA)

STORAGE_VERSION = 1
STORAGE_KEY = 'auth_module.totp'
STORAGE_USERS = 'users'
STORAGE_USER_ID = 'user_id'
STORAGE_OTA_SECRET = 'ota_secret'

INPUT_FIELD_CODE = 'code'

DUMMY_SECRET = 'FPPTH34D4E3MI2HG'

_LOGGER = logging.getLogger(__name__)


def _generate_qr_code(data: str) -> str:
"""Generate a base64 PNG string represent QR Code image of data."""
import pyqrcode

qr_code = pyqrcode.create(data)
return str(qr_code.png_as_base64_str(scale=4))


@bind_hass
async def _async_generate_secret(hass: HomeAssistant, username: str) \
-> Tuple[str, str, str]:
"""Generate a secret, url, and QR code."""
def generate_secret_helper(username: str) -> Tuple[str, str, str]:
"""Generate a secret, url, and QR code."""
import pyotp

ota_secret = pyotp.random_base32()
url = pyotp.totp.TOTP(ota_secret).provisioning_uri(
username, issuer_name="Home Assistant")
image = _generate_qr_code(url)
return ota_secret, url, image

return await hass.async_add_executor_job(
generate_secret_helper, username)


@MULTI_FACTOR_AUTH_MODULES.register('totp')
class TotpAuthModule(MultiFactorAuthModule):
"""Auth module validate time-based one time password."""

DEFAULT_TITLE = 'Time-based One Time Password'

def __init__(self, hass: HomeAssistant, config: Dict[str, Any]) -> None:
"""Initialize the user data store."""
super().__init__(hass, config)
self._users = None # type: Optional[Dict[str, str]]
self._user_store = hass.helpers.storage.Store(
STORAGE_VERSION, STORAGE_KEY)

@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._users = data.get(STORAGE_USERS, {})

async def _async_save(self) -> None:
"""Save data."""
await self._user_store.async_save({STORAGE_USERS: self._users})

def _add_ota_secret(self, user_id: str,
secret: Optional[str] = None) -> str:
"""Create a ota_secret for user."""
import pyotp

ota_secret = secret or pyotp.random_base32() # type: str

self._users[user_id] = ota_secret # type: ignore
return ota_secret

async def async_setup_flow(self, user_id: str) -> SetupFlow:
"""Return a data entry flow handler for setup module.
Mfa module should extend SetupFlow
"""
user = await self.hass.auth.async_get_user(user_id) # type: ignore
return TotpSetupFlow(self, user)

async def async_setup_user(self, user_id: str, setup_data: Any) -> str:
"""Set up auth module for user."""
if self._users is None:
await self._async_load()

result = await self.hass.async_add_executor_job(
self._add_ota_secret, user_id, setup_data.get('secret'))

await self._async_save()
return result

async def async_depose_user(self, user_id: str) -> None:
"""Depose auth module for user."""
if self._users is None:
await self._async_load()

if self._users.pop(user_id, None): # type: ignore
await self._async_save()

async def async_is_user_setup(self, user_id: str) -> bool:
"""Return whether user is setup."""
if self._users is None:
await self._async_load()

return user_id in self._users # type: ignore

async def async_validation(
self, user_id: str, user_input: Dict[str, Any]) -> bool:
"""Return True if validation passed."""
if self._users is None:
await self._async_load()

# user_input has been validate in caller
return await self.hass.async_add_executor_job(
self._validate_2fa, user_id, user_input[INPUT_FIELD_CODE])

def _validate_2fa(self, user_id: str, code: str) -> bool:
"""Validate two factor authentication code."""
import pyotp

ota_secret = self._users.get(user_id) # type: ignore
if ota_secret is None:
# even we cannot find user, we still do verify
# to make timing the same as if user was found.
pyotp.TOTP(DUMMY_SECRET).verify(code)
return False

return bool(pyotp.TOTP(ota_secret).verify(code))


class TotpSetupFlow(SetupFlow):
"""Handler for the setup flow."""

def __init__(self, auth_module: TotpAuthModule,
user: User) -> None:
"""Initialize the setup flow."""
super().__init__(auth_module, user.id)
self._auth_module = auth_module # type: TotpAuthModule
self._user = user
self._ota_secret = None # type: Optional[str]
self._url = None # type Optional[str]
self._image = None # type Optional[str]

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 self.async_create_entry(data={'result': result}) if finish.
"""
import pyotp

errors = {} # type: Dict[str, str]

if user_input:
verified = await self.hass.async_add_executor_job( # type: ignore
pyotp.TOTP(self._ota_secret).verify, user_input['code'])
if verified:
result = await self._auth_module.async_setup_user(
self._user_id, {'secret': self._ota_secret})
return self.async_create_entry(
title=self._auth_module.name,
data={'result': result}
)

errors['base'] = 'invalid_code'

else:
self._ota_secret, self._url, self._image = ( # type: ignore
await _async_generate_secret(
self._auth_module.hass, str(self._user.name)))

return self.async_show_form(
step_id='init',
data_schema=self._auth_module.input_schema,
description_placeholders={
'code': self._ota_secret,
'url': self._url,
'qr_code': self._image
},
errors=errors
)
4 changes: 2 additions & 2 deletions homeassistant/auth/providers/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -168,7 +168,7 @@ def __init__(self, auth_provider: AuthProvider) -> None:
self._auth_provider = auth_provider
self._auth_module_id = None # type: Optional[str]
self._auth_manager = auth_provider.hass.auth # type: ignore
self.available_mfa_modules = [] # type: List
self.available_mfa_modules = {} # type: Dict[str, str]
self.created_at = dt_util.utcnow()
self.user = None # type: Optional[User]

Expand Down Expand Up @@ -196,7 +196,7 @@ async def async_step_select_mfa_module(
errors['base'] = 'invalid_auth_module'

if len(self.available_mfa_modules) == 1:
self._auth_module_id = self.available_mfa_modules[0]
self._auth_module_id = list(self.available_mfa_modules.keys())[0]
return await self.async_step_mfa()

return self.async_show_form(
Expand Down
2 changes: 1 addition & 1 deletion homeassistant/components/auth/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@
from . import mfa_setup_flow

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({
Expand Down
16 changes: 16 additions & 0 deletions homeassistant/components/auth/strings.json
Original file line number Diff line number Diff line change
@@ -0,0 +1,16 @@
{
"mfa_setup":{
"totp": {
"title": "TOTP",
"step": {
"init": {
"title": "Scan this QR code with your app",
"description": "Scan the QR code with your authentication app, such as **Google Authenticator** or **Authy**. \nOr enter {code} instead. \n\n[QR Code]{url}",
},
},
"error": {
"invalid_code": "Invalid code, please try again. If you get this error consistently, please make sure the clock on Home Assistant system is accurate."
}
}
}
}
61 changes: 61 additions & 0 deletions homeassistant/components/auth/util.py
Original file line number Diff line number Diff line change
@@ -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):
"""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: 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
7 changes: 7 additions & 0 deletions requirements_all.txt
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,9 @@ PyMVGLive==1.1.4
# homeassistant.components.arduino
PyMata==2.14

# homeassistant.auth.mfa_modules.totp
PyQRCode==1.2.1

# homeassistant.components.sensor.rmvtransport
PyRMVtransport==0.0.7

Expand Down Expand Up @@ -985,6 +988,7 @@ pyopenuv==1.0.1
# homeassistant.components.iota
pyota==2.0.5

# homeassistant.auth.mfa_modules.totp
# homeassistant.components.sensor.otp
pyotp==2.2.6

Expand All @@ -995,6 +999,9 @@ pyowm==2.9.0
# homeassistant.components.media_player.pjlink
pypjlink2==1.2.0

# homeassistant.auth.mfa_modules.totp
pypng==0.0.18

# homeassistant.components.sensor.pollen
pypollencom==2.1.0

Expand Down
4 changes: 4 additions & 0 deletions requirements_test_all.txt
Original file line number Diff line number Diff line change
Expand Up @@ -154,6 +154,10 @@ pymonoprice==0.3
# homeassistant.components.binary_sensor.nx584
pynx584==0.4

# homeassistant.auth.mfa_modules.totp
# homeassistant.components.sensor.otp
pyotp==2.2.6

# homeassistant.components.qwikswitch
pyqwikswitch==0.8

Expand Down
1 change: 1 addition & 0 deletions script/gen_requirements_all.py
Original file line number Diff line number Diff line change
Expand Up @@ -78,6 +78,7 @@
'pylitejet',
'pymonoprice',
'pynx584',
'pyotp',
'pyqwikswitch',
'PyRMVtransport',
'python-forecastio',
Expand Down
Loading

0 comments on commit 37842c2

Please sign in to comment.