Skip to content

Commit

Permalink
Add websocket commands for refresh tokens (#16559)
Browse files Browse the repository at this point in the history
* Add websocket commands for refresh tokens

* Comment
  • Loading branch information
balloob authored and awarecan committed Sep 11, 2018
1 parent 4e3faf6 commit 0db13a9
Show file tree
Hide file tree
Showing 3 changed files with 140 additions and 44 deletions.
60 changes: 60 additions & 0 deletions homeassistant/components/auth/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -156,6 +156,19 @@
vol.Optional('client_icon'): str,
})

WS_TYPE_REFRESH_TOKENS = 'auth/refresh_tokens'
SCHEMA_WS_REFRESH_TOKENS = \
websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): WS_TYPE_REFRESH_TOKENS,
})

WS_TYPE_DELETE_REFRESH_TOKEN = 'auth/delete_refresh_token'
SCHEMA_WS_DELETE_REFRESH_TOKEN = \
websocket_api.BASE_COMMAND_MESSAGE_SCHEMA.extend({
vol.Required('type'): WS_TYPE_DELETE_REFRESH_TOKEN,
vol.Required('refresh_token_id'): str,
})

RESULT_TYPE_CREDENTIALS = 'credentials'
RESULT_TYPE_USER = 'user'

Expand All @@ -178,6 +191,16 @@ async def async_setup(hass, config):
websocket_create_long_lived_access_token,
SCHEMA_WS_LONG_LIVED_ACCESS_TOKEN
)
hass.components.websocket_api.async_register_command(
WS_TYPE_REFRESH_TOKENS,
websocket_refresh_tokens,
SCHEMA_WS_REFRESH_TOKENS
)
hass.components.websocket_api.async_register_command(
WS_TYPE_DELETE_REFRESH_TOKEN,
websocket_delete_refresh_token,
SCHEMA_WS_DELETE_REFRESH_TOKEN
)

await login_flow.async_setup(hass, store_result)
await mfa_setup_flow.async_setup(hass)
Expand Down Expand Up @@ -445,3 +468,40 @@ async def async_create_long_lived_access_token(user):

hass.async_create_task(
async_create_long_lived_access_token(connection.user))


@websocket_api.ws_require_user()
@callback
def websocket_refresh_tokens(
hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg):
"""Return metadata of users refresh tokens."""
connection.to_write.put_nowait(websocket_api.result_message(msg['id'], [{
'id': refresh.id,
'client_id': refresh.client_id,
'client_name': refresh.client_name,
'client_icon': refresh.client_icon,
'type': refresh.token_type,
'created_at': refresh.created_at,
} for refresh in connection.user.refresh_tokens.values()]))


@websocket_api.ws_require_user()
@callback
def websocket_delete_refresh_token(
hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg):
"""Handle a delete refresh token request."""
async def async_delete_refresh_token(user, refresh_token_id):
"""Delete a refresh token."""
refresh_token = connection.user.refresh_tokens.get(refresh_token_id)

if refresh_token is None:
return websocket_api.error_message(
msg['id'], 'invalid_token_id', 'Received invalid token')

await hass.auth.async_remove_refresh_token(refresh_token)

connection.send_message_outside(
websocket_api.result_message(msg['id'], {}))

hass.async_create_task(
async_delete_refresh_token(connection.user, msg['refresh_token_id']))
82 changes: 51 additions & 31 deletions tests/components/auth/test_init.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,18 +2,15 @@
from datetime import timedelta
from unittest.mock import patch

from homeassistant import const
from homeassistant.auth import auth_manager_from_config
from homeassistant.auth.models import Credentials
from homeassistant.components.auth import RESULT_TYPE_USER
from homeassistant.setup import async_setup_component
from homeassistant.util.dt import utcnow
from homeassistant.components import auth

from . import async_setup_auth
from tests.common import CLIENT_ID, CLIENT_REDIRECT_URI, MockUser

from tests.common import CLIENT_ID, CLIENT_REDIRECT_URI, MockUser, \
ensure_auth_manager_loaded
from . import async_setup_auth


async def test_login_new_user_and_trying_refresh_token(hass, aiohttp_client):
Expand Down Expand Up @@ -272,28 +269,12 @@ async def test_revoking_refresh_token(hass, aiohttp_client):
assert resp.status == 400


async def test_ws_long_lived_access_token(hass, hass_ws_client):
async def test_ws_long_lived_access_token(hass, hass_ws_client,
hass_access_token):
"""Test generate long-lived access token."""
hass.auth = await auth_manager_from_config(
hass, provider_configs=[{
'type': 'insecure_example',
'users': [{
'username': 'test-user',
'password': 'test-pass',
'name': 'Test Name',
}]
}], module_configs=[])
ensure_auth_manager_loaded(hass.auth)
assert await async_setup_component(hass, 'auth', {'http': {}})
assert await async_setup_component(hass, 'api', {'http': {}})

user = MockUser(id='mock-user').add_to_hass(hass)
cred = await hass.auth.auth_providers[0].async_get_or_create_credentials(
{'username': 'test-user'})
await hass.auth.async_link_user(user, cred)

ws_client = await hass_ws_client(hass, hass.auth.async_create_access_token(
await hass.auth.async_create_refresh_token(user, CLIENT_ID)))
ws_client = await hass_ws_client(hass, hass_access_token)

# verify create long-lived access token
await ws_client.send_json({
Expand All @@ -315,12 +296,51 @@ async def test_ws_long_lived_access_token(hass, hass_ws_client):
assert refresh_token.client_name == 'GPS Logger'
assert refresh_token.client_icon is None

# verify long-lived access token can be used as bearer token
api_client = ws_client.client
resp = await api_client.get(const.URL_API)
assert resp.status == 401

resp = await api_client.get(const.URL_API, headers={
'Authorization': 'Bearer {}'.format(long_lived_access_token)
async def test_ws_refresh_tokens(hass, hass_ws_client, hass_access_token):
"""Test fetching refresh token metadata."""
assert await async_setup_component(hass, 'auth', {'http': {}})

ws_client = await hass_ws_client(hass, hass_access_token)

await ws_client.send_json({
'id': 5,
'type': auth.WS_TYPE_REFRESH_TOKENS,
})
assert resp.status == 200

result = await ws_client.receive_json()
assert result['success'], result
assert len(result['result']) == 1
token = result['result'][0]
refresh_token = await hass.auth.async_validate_access_token(
hass_access_token)
assert token['id'] == refresh_token.id
assert token['type'] == refresh_token.token_type
assert token['client_id'] == refresh_token.client_id
assert token['client_name'] == refresh_token.client_name
assert token['client_icon'] == refresh_token.client_icon
assert token['created_at'] == refresh_token.created_at.isoformat()


async def test_ws_delete_refresh_token(hass, hass_ws_client,
hass_access_token):
"""Test deleting a refresh token."""
assert await async_setup_component(hass, 'auth', {'http': {}})

refresh_token = await hass.auth.async_validate_access_token(
hass_access_token)

ws_client = await hass_ws_client(hass, hass_access_token)

# verify create long-lived access token
await ws_client.send_json({
'id': 5,
'type': auth.WS_TYPE_DELETE_REFRESH_TOKEN,
'refresh_token_id': refresh_token.id
})

result = await ws_client.receive_json()
assert result['success'], result
refresh_token = await hass.auth.async_validate_access_token(
hass_access_token)
assert refresh_token is None
42 changes: 29 additions & 13 deletions tests/components/conftest.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
"""Fixtures for component testing."""
from unittest.mock import patch

import pytest

from homeassistant.setup import async_setup_component
Expand All @@ -16,23 +18,37 @@ async def create_client(hass, access_token=None):
assert await async_setup_component(hass, 'websocket_api')

client = await aiohttp_client(hass.http.app)
websocket = await client.ws_connect(wapi.URL)
auth_resp = await websocket.receive_json()

if auth_resp['type'] == wapi.TYPE_AUTH_OK:
assert access_token is None, \
'Access token given but no auth required'
return websocket
patching = None

if access_token is not None:
patching = patch('homeassistant.auth.AuthManager.active',
return_value=True)
patching.start()

try:
websocket = await client.ws_connect(wapi.URL)
auth_resp = await websocket.receive_json()

if auth_resp['type'] == wapi.TYPE_AUTH_OK:
assert access_token is None, \
'Access token given but no auth required'
return websocket

assert access_token is not None, \
'Access token required for fixture'

assert access_token is not None, 'Access token required for fixture'
await websocket.send_json({
'type': websocket_api.TYPE_AUTH,
'access_token': access_token
})

await websocket.send_json({
'type': websocket_api.TYPE_AUTH,
'access_token': access_token
})
auth_ok = await websocket.receive_json()
assert auth_ok['type'] == wapi.TYPE_AUTH_OK

auth_ok = await websocket.receive_json()
assert auth_ok['type'] == wapi.TYPE_AUTH_OK
finally:
if patching is not None:
patching.stop()

# wrap in client
websocket.client = client
Expand Down

0 comments on commit 0db13a9

Please sign in to comment.