Extend hass.io services / updater (#11549)

* Extend hass.io services

* Add warning for carfuly options with hass.io

* update tests

* finish tests

* remove update calls

* address comments

* address comments p2

* fix tests

* fix tests

* Use token also for proxy

* Add test for server_host

* Fix test

* Fix tests

* Add test for version

* Address comments

homeassistant/components/hassio.py

@@ -5,6 +5,7 @@ For more details about this platform, please refer to the documentation at
 https://home-assistant.io/components/hassio/
 """
 import asyncio
+from datetime import timedelta
 import logging
 import os
 import re
@@ -21,23 +22,38 @@ from homeassistant.const import (
     CONTENT_TYPE_TEXT_PLAIN, SERVER_PORT, CONF_TIME_ZONE)
 from homeassistant.components.http import (
     HomeAssistantView, KEY_AUTHENTICATED, CONF_API_PASSWORD, CONF_SERVER_PORT,
-    CONF_SSL_CERTIFICATE)
+    CONF_SERVER_HOST, CONF_SSL_CERTIFICATE)
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.loader import bind_hass
+from homeassistant.util.dt import utcnow
 
 _LOGGER = logging.getLogger(__name__)
 
 DOMAIN = 'hassio'
 DEPENDENCIES = ['http']
 
+X_HASSIO = 'X-HASSIO-KEY'
+
+DATA_HOMEASSISTANT_VERSION = 'hassio_hass_version'
+HASSIO_UPDATE_INTERVAL = timedelta(hours=1)
+
 SERVICE_ADDON_START = 'addon_start'
 SERVICE_ADDON_STOP = 'addon_stop'
 SERVICE_ADDON_RESTART = 'addon_restart'
 SERVICE_ADDON_STDIN = 'addon_stdin'
 SERVICE_HOST_SHUTDOWN = 'host_shutdown'
 SERVICE_HOST_REBOOT = 'host_reboot'
+SERVICE_SNAPSHOT_FULL = 'snapshot_full'
+SERVICE_SNAPSHOT_PARTIAL = 'snapshot_partial'
+SERVICE_RESTORE_FULL = 'restore_full'
+SERVICE_RESTORE_PARTIAL = 'restore_partial'
 
 ATTR_ADDON = 'addon'
 ATTR_INPUT = 'input'
+ATTR_SNAPSHOT = 'snapshot'
+ATTR_ADDONS = 'addons'
+ATTR_FOLDERS = 'folders'
+ATTR_HOMEASSISTANT = 'homeassistant'
+ATTR_NAME = 'name'
 
 NO_TIMEOUT = {
     re.compile(r'^homeassistant/update$'),
@@ -45,13 +61,17 @@ NO_TIMEOUT = {
     re.compile(r'^supervisor/update$'),
     re.compile(r'^addons/[^/]*/update$'),
     re.compile(r'^addons/[^/]*/install$'),
-    re.compile(r'^addons/[^/]*/rebuild$')
+    re.compile(r'^addons/[^/]*/rebuild$'),
+    re.compile(r'^snapshots/.*/full$'),
+    re.compile(r'^snapshots/.*/partial$'),
 }
 
 NO_AUTH = {
     re.compile(r'^panel_(es5|latest)$'), re.compile(r'^addons/[^/]*/logo$')
 }
 
+SCHEMA_NO_DATA = vol.Schema({})
+
 SCHEMA_ADDON = vol.Schema({
     vol.Required(ATTR_ADDON): cv.slug,
 })
@@ -60,16 +80,52 @@ SCHEMA_ADDON_STDIN = SCHEMA_ADDON.extend({
     vol.Required(ATTR_INPUT): vol.Any(dict, cv.string)
 })
 
+SCHEMA_SNAPSHOT_FULL = vol.Schema({
+    vol.Optional(ATTR_NAME): cv.string,
+})
+
+SCHEMA_SNAPSHOT_PARTIAL = SCHEMA_SNAPSHOT_FULL.extend({
+    vol.Optional(ATTR_FOLDERS): vol.All(cv.ensure_list, [cv.string]),
+    vol.Optional(ATTR_ADDONS): vol.All(cv.ensure_list, [cv.string]),
+})
+
+SCHEMA_RESTORE_FULL = vol.Schema({
+    vol.Required(ATTR_SNAPSHOT): cv.slug,
+})
+
+SCHEMA_RESTORE_PARTIAL = SCHEMA_RESTORE_FULL.extend({
+    vol.Optional(ATTR_HOMEASSISTANT): cv.boolean,
+    vol.Optional(ATTR_FOLDERS): vol.All(cv.ensure_list, [cv.string]),
+    vol.Optional(ATTR_ADDONS): vol.All(cv.ensure_list, [cv.string]),
+})
+
 MAP_SERVICE_API = {
-    SERVICE_ADDON_START: ('/addons/{addon}/start', SCHEMA_ADDON),
+    SERVICE_ADDON_START: ('/addons/{addon}/start', SCHEMA_ADDON, 60, False),
-    SERVICE_ADDON_STOP: ('/addons/{addon}/stop', SCHEMA_ADDON),
+    SERVICE_ADDON_STOP: ('/addons/{addon}/stop', SCHEMA_ADDON, 60, False),
-    SERVICE_ADDON_RESTART: ('/addons/{addon}/restart', SCHEMA_ADDON),
+    SERVICE_ADDON_RESTART:
-    SERVICE_ADDON_STDIN: ('/addons/{addon}/stdin', SCHEMA_ADDON_STDIN),
+        ('/addons/{addon}/restart', SCHEMA_ADDON, 60, False),
-    SERVICE_HOST_SHUTDOWN: ('/host/shutdown', None),
+    SERVICE_ADDON_STDIN:
-    SERVICE_HOST_REBOOT: ('/host/reboot', None),
+        ('/addons/{addon}/stdin', SCHEMA_ADDON_STDIN, 60, False),
+    SERVICE_HOST_SHUTDOWN: ('/host/shutdown', SCHEMA_NO_DATA, 60, False),
+    SERVICE_HOST_REBOOT: ('/host/reboot', SCHEMA_NO_DATA, 60, False),
+    SERVICE_SNAPSHOT_FULL:
+        ('/snapshots/new/full', SCHEMA_SNAPSHOT_FULL, 300, True),
+    SERVICE_SNAPSHOT_PARTIAL:
+        ('/snapshots/new/partial', SCHEMA_SNAPSHOT_PARTIAL, 300, True),
+    SERVICE_RESTORE_FULL:
+        ('/snapshots/{snapshot}/restore/full', SCHEMA_RESTORE_FULL, 300, True),
+    SERVICE_RESTORE_PARTIAL:
+        ('/snapshots/{snapshot}/restore/partial', SCHEMA_RESTORE_PARTIAL, 300,
+         True),
 }
 
 
+@bind_hass
+def get_homeassistant_version(hass):
+    """Return last available HomeAssistant version."""
+    return hass.data.get(DATA_HOMEASSISTANT_VERSION)
+
+
 @asyncio.coroutine
 def async_setup(hass, config):
     """Set up the HASSio component."""
@@ -79,7 +135,7 @@ def async_setup(hass, config):
         _LOGGER.error("No HassIO supervisor detect!")
         return False
 
-    websession = async_get_clientsession(hass)
+    websession = hass.helpers.aiohttp_client.async_get_clientsession()
     hassio = HassIO(hass.loop, websession, host)
 
     if not (yield from hassio.is_connected()):
@@ -102,16 +158,41 @@ def async_setup(hass, config):
     def async_service_handler(service):
         """Handle service calls for HassIO."""
         api_command = MAP_SERVICE_API[service.service][0]
-        addon = service.data.get(ATTR_ADDON)
+        data = service.data.copy()
-        data = service.data[ATTR_INPUT] if ATTR_INPUT in service.data else None
+        addon = data.pop(ATTR_ADDON, None)
+        snapshot = data.pop(ATTR_SNAPSHOT, None)
+        payload = None
 
+        # Pass data to hass.io API
+        if service.service == SERVICE_ADDON_STDIN:
+            payload = data[ATTR_INPUT]
+        elif MAP_SERVICE_API[service.service][3]:
+            payload = data
+
+        # Call API
         yield from hassio.send_command(
-            api_command.format(addon=addon), payload=data, timeout=60)
+            api_command.format(addon=addon, snapshot=snapshot),
+            payload=payload, timeout=MAP_SERVICE_API[service.service][2]
+        )
 
     for service, settings in MAP_SERVICE_API.items():
         hass.services.async_register(
             DOMAIN, service, async_service_handler, schema=settings[1])
 
+    @asyncio.coroutine
+    def update_homeassistant_version(now):
+        """Update last available HomeAssistant version."""
+        data = yield from hassio.get_homeassistant_info()
+        if data:
+            hass.data[DATA_HOMEASSISTANT_VERSION] = \
+                data['data']['last_version']
+
+        hass.helpers.event.async_track_point_in_utc_time(
+            update_homeassistant_version, utcnow() + HASSIO_UPDATE_INTERVAL)
+
+    # Fetch last version
+    yield from update_homeassistant_version(None)
+
     return True
 
 
@@ -131,6 +212,13 @@ class HassIO(object):
         """
         return self.send_command("/supervisor/ping", method="get")
 
+    def get_homeassistant_info(self):
+        """Return data for HomeAssistant.
+
+        This method return a coroutine.
+        """
+        return self.send_command("/homeassistant/info", method="get")
+
     def update_hass_api(self, http_config):
         """Update Home-Assistant API data on HassIO.
 
@@ -141,8 +229,13 @@ class HassIO(object):
             'ssl': CONF_SSL_CERTIFICATE in http_config,
             'port': port,
             'password': http_config.get(CONF_API_PASSWORD),
+            'watchdog': True,
         }
 
+        if CONF_SERVER_HOST in http_config:
+            options['watchdog'] = False
+            _LOGGER.warning("Don't use 'server_host' options with Hass.io!")
+
         return self.send_command("/homeassistant/options", payload=options)
 
     def update_hass_timezone(self, core_config):
@@ -164,15 +257,17 @@ class HassIO(object):
             with async_timeout.timeout(timeout, loop=self.loop):
                 request = yield from self.websession.request(
                     method, "http://{}{}".format(self._ip, command),
-                    json=payload)
+                    json=payload, headers={
+                        X_HASSIO: os.environ.get('HASSIO_TOKEN')
+                    })
 
                 if request.status != 200:
                     _LOGGER.error(
                         "%s return code %d.", command, request.status)
-                    return False
+                    return None
 
                 answer = yield from request.json()
-                return answer and answer['result'] == 'ok'
+                return answer
 
         except asyncio.TimeoutError:
             _LOGGER.error("Timeout on %s request", command)
@@ -180,7 +275,7 @@ class HassIO(object):
         except aiohttp.ClientError as err:
             _LOGGER.error("Client error on %s request %s", command, err)
 
-        return False
+        return None
 
     @asyncio.coroutine
     def command_proxy(self, path, request):
@@ -192,11 +287,11 @@ class HassIO(object):
 
         try:
             data = None
-            headers = None
+            headers = {X_HASSIO: os.environ.get('HASSIO_TOKEN')}
             with async_timeout.timeout(10, loop=self.loop):
                 data = yield from request.read()
                 if data:
-                    headers = {CONTENT_TYPE: request.content_type}
+                    headers[CONTENT_TYPE] = request.content_type
                 else:
                     data = None
 

tests/components/test_hassio.py

@@ -18,7 +18,12 @@ def hassio_env():
     """Fixture to inject hassio env."""
     with patch.dict(os.environ, {'HASSIO': "127.0.0.1"}), \
             patch('homeassistant.components.hassio.HassIO.is_connected',
-                  Mock(return_value=mock_coro(True))):
+                  Mock(return_value=mock_coro(
+                    {"result": "ok", "data": {}}))), \
+            patch.dict(os.environ, {'HASSIO_TOKEN': "123456"}), \
+            patch('homeassistant.components.hassio.HassIO.'
+                  'get_homeassistant_info',
+                  Mock(return_value=mock_coro(None))):
         yield
 
 
@@ -26,7 +31,10 @@ def hassio_env():
 def hassio_client(hassio_env, hass, test_client):
     """Create mock hassio http client."""
     with patch('homeassistant.components.hassio.HassIO.update_hass_api',
-               Mock(return_value=mock_coro(True))):
+               Mock(return_value=mock_coro({"result": "ok"}))), \
+            patch('homeassistant.components.hassio.HassIO.'
+                  'get_homeassistant_info',
+                  Mock(return_value=mock_coro(None))):
         hass.loop.run_until_complete(async_setup_component(hass, 'hassio', {
             'http': {
                 'api_password': API_PASSWORD
@@ -48,7 +56,7 @@ def test_fail_setup_cannot_connect(hass):
     """Fail setup if cannot connect."""
     with patch.dict(os.environ, {'HASSIO': "127.0.0.1"}), \
             patch('homeassistant.components.hassio.HassIO.is_connected',
-                  Mock(return_value=mock_coro(False))):
+                  Mock(return_value=mock_coro(None))):
         result = yield from async_setup_component(hass, 'hassio', {})
         assert not result
 
@@ -58,12 +66,16 @@ def test_setup_api_ping(hass, aioclient_mock):
     """Test setup with API ping."""
     aioclient_mock.get(
         "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={
+            'result': 'ok', 'data': {'last_version': '10.0'}})
 
     with patch.dict(os.environ, {'HASSIO': "127.0.0.1"}):
         result = yield from async_setup_component(hass, 'hassio', {})
         assert result
 
-    assert aioclient_mock.call_count == 1
+    assert aioclient_mock.call_count == 2
+    assert hass.data['hassio_hass_version'] == "10.0"
 
 
 @asyncio.coroutine
@@ -71,6 +83,9 @@ def test_setup_api_push_api_data(hass, aioclient_mock):
     """Test setup with API push."""
     aioclient_mock.get(
         "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={
+            'result': 'ok', 'data': {'last_version': '10.0'}})
     aioclient_mock.post(
         "http://127.0.0.1/homeassistant/options", json={'result': 'ok'})
 
@@ -84,10 +99,40 @@ def test_setup_api_push_api_data(hass, aioclient_mock):
         })
         assert result
 
-    assert aioclient_mock.call_count == 2
+    assert aioclient_mock.call_count == 3
-    assert not aioclient_mock.mock_calls[-1][2]['ssl']
+    assert not aioclient_mock.mock_calls[1][2]['ssl']
-    assert aioclient_mock.mock_calls[-1][2]['password'] == "123456"
+    assert aioclient_mock.mock_calls[1][2]['password'] == "123456"
-    assert aioclient_mock.mock_calls[-1][2]['port'] == 9999
+    assert aioclient_mock.mock_calls[1][2]['port'] == 9999
+    assert aioclient_mock.mock_calls[1][2]['watchdog']
+
+
+@asyncio.coroutine
+def test_setup_api_push_api_data_server_host(hass, aioclient_mock):
+    """Test setup with API push with active server host."""
+    aioclient_mock.get(
+        "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={
+            'result': 'ok', 'data': {'last_version': '10.0'}})
+    aioclient_mock.post(
+        "http://127.0.0.1/homeassistant/options", json={'result': 'ok'})
+
+    with patch.dict(os.environ, {'HASSIO': "127.0.0.1"}):
+        result = yield from async_setup_component(hass, 'hassio', {
+            'http': {
+                'api_password': "123456",
+                'server_port': 9999,
+                'server_host': "127.0.0.1"
+            },
+            'hassio': {}
+        })
+        assert result
+
+    assert aioclient_mock.call_count == 3
+    assert not aioclient_mock.mock_calls[1][2]['ssl']
+    assert aioclient_mock.mock_calls[1][2]['password'] == "123456"
+    assert aioclient_mock.mock_calls[1][2]['port'] == 9999
+    assert not aioclient_mock.mock_calls[1][2]['watchdog']
 
 
 @asyncio.coroutine
@@ -95,6 +140,9 @@ def test_setup_api_push_api_data_default(hass, aioclient_mock):
     """Test setup with API push default data."""
     aioclient_mock.get(
         "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={
+            'result': 'ok', 'data': {'last_version': '10.0'}})
     aioclient_mock.post(
         "http://127.0.0.1/homeassistant/options", json={'result': 'ok'})
 
@@ -105,10 +153,10 @@ def test_setup_api_push_api_data_default(hass, aioclient_mock):
         })
         assert result
 
-    assert aioclient_mock.call_count == 2
+    assert aioclient_mock.call_count == 3
-    assert not aioclient_mock.mock_calls[-1][2]['ssl']
+    assert not aioclient_mock.mock_calls[1][2]['ssl']
-    assert aioclient_mock.mock_calls[-1][2]['password'] is None
+    assert aioclient_mock.mock_calls[1][2]['password'] is None
-    assert aioclient_mock.mock_calls[-1][2]['port'] == 8123
+    assert aioclient_mock.mock_calls[1][2]['port'] == 8123
 
 
 @asyncio.coroutine
@@ -116,6 +164,9 @@ def test_setup_core_push_timezone(hass, aioclient_mock):
     """Test setup with API push default data."""
     aioclient_mock.get(
         "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={
+            'result': 'ok', 'data': {'last_version': '10.0'}})
     aioclient_mock.post(
         "http://127.0.0.1/supervisor/options", json={'result': 'ok'})
 
@@ -128,8 +179,8 @@ def test_setup_core_push_timezone(hass, aioclient_mock):
         })
         assert result
 
-    assert aioclient_mock.call_count == 2
+    assert aioclient_mock.call_count == 3
-    assert aioclient_mock.mock_calls[-1][2]['timezone'] == "testzone"
+    assert aioclient_mock.mock_calls[1][2]['timezone'] == "testzone"
 
 
 @asyncio.coroutine
@@ -137,14 +188,21 @@ def test_setup_hassio_no_additional_data(hass, aioclient_mock):
     """Test setup with API push default data."""
     aioclient_mock.get(
         "http://127.0.0.1/supervisor/ping", json={'result': 'ok'})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={
+            'result': 'ok', 'data': {'last_version': '10.0'}})
+    aioclient_mock.get(
+        "http://127.0.0.1/homeassistant/info", json={'result': 'ok'})
 
-    with patch.dict(os.environ, {'HASSIO': "127.0.0.1"}):
+    with patch.dict(os.environ, {'HASSIO': "127.0.0.1"}), \
+            patch.dict(os.environ, {'HASSIO_TOKEN': "123456"}):
         result = yield from async_setup_component(hass, 'hassio', {
             'hassio': {},
         })
         assert result
 
-    assert aioclient_mock.call_count == 1
+    assert aioclient_mock.call_count == 2
+    assert aioclient_mock.mock_calls[-1][3]['X-HASSIO-KEY'] == "123456"
 
 
 @asyncio.coroutine
@@ -157,6 +215,10 @@ def test_service_register(hassio_env, hass):
     assert hass.services.has_service('hassio', 'addon_stdin')
     assert hass.services.has_service('hassio', 'host_shutdown')
     assert hass.services.has_service('hassio', 'host_reboot')
+    assert hass.services.has_service('hassio', 'snapshot_full')
+    assert hass.services.has_service('hassio', 'snapshot_partial')
+    assert hass.services.has_service('hassio', 'restore_full')
+    assert hass.services.has_service('hassio', 'restore_partial')
 
 
 @asyncio.coroutine
@@ -176,6 +238,15 @@ def test_service_calls(hassio_env, hass, aioclient_mock):
         "http://127.0.0.1/host/shutdown", json={'result': 'ok'})
     aioclient_mock.post(
         "http://127.0.0.1/host/reboot", json={'result': 'ok'})
+    aioclient_mock.post(
+        "http://127.0.0.1/snapshots/new/full", json={'result': 'ok'})
+    aioclient_mock.post(
+        "http://127.0.0.1/snapshots/new/partial", json={'result': 'ok'})
+    aioclient_mock.post(
+        "http://127.0.0.1/snapshots/test/restore/full", json={'result': 'ok'})
+    aioclient_mock.post(
+        "http://127.0.0.1/snapshots/test/restore/partial",
+        json={'result': 'ok'})
 
     yield from hass.services.async_call(
         'hassio', 'addon_start', {'addon': 'test'})
@@ -196,6 +267,32 @@ def test_service_calls(hassio_env, hass, aioclient_mock):
 
     assert aioclient_mock.call_count == 6
 
+    yield from hass.services.async_call('hassio', 'snapshot_full', {})
+    yield from hass.services.async_call('hassio', 'snapshot_partial', {
+        'addons': ['test'],
+        'folders': ['ssl'],
+    })
+    yield from hass.async_block_till_done()
+
+    assert aioclient_mock.call_count == 8
+    assert aioclient_mock.mock_calls[-1][2] == {
+        'addons': ['test'], 'folders': ['ssl']}
+
+    yield from hass.services.async_call('hassio', 'restore_full', {
+        'snapshot': 'test',
+    })
+    yield from hass.services.async_call('hassio', 'restore_partial', {
+        'snapshot': 'test',
+        'homeassistant': False,
+        'addons': ['test'],
+        'folders': ['ssl'],
+    })
+    yield from hass.async_block_till_done()
+
+    assert aioclient_mock.call_count == 10
+    assert aioclient_mock.mock_calls[-1][2] == {
+        'addons': ['test'], 'folders': ['ssl'], 'homeassistant': False}
+
 
 @asyncio.coroutine
 def test_forward_request(hassio_client):

tests/test_util/aiohttp.py

@@ -83,7 +83,7 @@ class AiohttpClientMocker:
         data = data or json
         for response in self._mocks:
             if response.match_request(method, url, params):
-                self.mock_calls.append((method, url, data))
+                self.mock_calls.append((method, url, data, headers))
 
                 if response.exc:
                     raise response.exc