From 79da44a6b3d5f986c90a10ab360095a4d20c7efa Mon Sep 17 00:00:00 2001
From: Lewis Juggins <ldjuggins@gmail.com>
Date: Thu, 2 Nov 2017 06:23:06 +0000
Subject: [PATCH] Support new tradfri individual DTLS identification method
 (#10282)

---
 homeassistant/components/tradfri.py | 46 +++++++++++++++++++----------
 requirements_all.txt                |  2 +-
 2 files changed, 31 insertions(+), 17 deletions(-)

diff --git a/homeassistant/components/tradfri.py b/homeassistant/components/tradfri.py
index 1ab949ed5eb..ead4924d599 100644
--- a/homeassistant/components/tradfri.py
+++ b/homeassistant/components/tradfri.py
@@ -13,17 +13,18 @@ import voluptuous as vol
 
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers import discovery
-from homeassistant.const import CONF_HOST, CONF_API_KEY
+from homeassistant.const import CONF_HOST
 from homeassistant.components.discovery import SERVICE_IKEA_TRADFRI
 
-REQUIREMENTS = ['pytradfri==3.0.2',
+REQUIREMENTS = ['pytradfri==4.0.1',
                 'DTLSSocket==0.1.4',
                 'https://github.com/chrysn/aiocoap/archive/'
                 '3286f48f0b949901c8b5c04c0719dc54ab63d431.zip'
                 '#aiocoap==0.3']
 
 DOMAIN = 'tradfri'
-CONFIG_FILE = 'tradfri.conf'
+GATEWAY_IDENTITY = 'homeassistant'
+CONFIG_FILE = '.tradfri_psk.conf'
 KEY_CONFIG = 'tradfri_configuring'
 KEY_GATEWAY = 'tradfri_gateway'
 KEY_API = 'tradfri_api'
@@ -34,7 +35,6 @@ DEFAULT_ALLOW_TRADFRI_GROUPS = True
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
         vol.Inclusive(CONF_HOST, 'gateway'): cv.string,
-        vol.Inclusive(CONF_API_KEY, 'gateway'): cv.string,
         vol.Optional(CONF_ALLOW_TRADFRI_GROUPS,
                      default=DEFAULT_ALLOW_TRADFRI_GROUPS): cv.boolean,
     })
@@ -56,9 +56,17 @@ def request_configuration(hass, config, host):
     @asyncio.coroutine
     def configuration_callback(callback_data):
         """Handle the submitted configuration."""
-        res = yield from _setup_gateway(hass, config, host,
-                                        callback_data.get('key'),
+        try:
+            from pytradfri.api.aiocoap_api import APIFactory
+        except ImportError:
+            _LOGGER.exception("Looks like something isn't installed!")
+            return
+
+        api_factory = APIFactory(host, psk_id=GATEWAY_IDENTITY)
+        psk = yield from api_factory.generate_psk(callback_data.get('key'))
+        res = yield from _setup_gateway(hass, config, host, psk,
                                         DEFAULT_ALLOW_TRADFRI_GROUPS)
+
         if not res:
             hass.async_add_job(configurator.notify_errors, instance,
                                "Unable to connect.")
@@ -67,7 +75,7 @@ def request_configuration(hass, config, host):
         def success():
             """Set up was successful."""
             conf = _read_config(hass)
-            conf[host] = {'key': callback_data.get('key')}
+            conf[host] = {'key': psk}
             _write_config(hass, conf)
             hass.async_add_job(configurator.request_done, instance)
 
@@ -87,13 +95,12 @@ def async_setup(hass, config):
     """Set up the Tradfri component."""
     conf = config.get(DOMAIN, {})
     host = conf.get(CONF_HOST)
-    key = conf.get(CONF_API_KEY)
     allow_tradfri_groups = conf.get(CONF_ALLOW_TRADFRI_GROUPS)
+    keys = yield from hass.async_add_job(_read_config, hass)
 
     @asyncio.coroutine
     def gateway_discovered(service, info):
         """Run when a gateway is discovered."""
-        keys = yield from hass.async_add_job(_read_config, hass)
         host = info['host']
 
         if host in keys:
@@ -104,11 +111,16 @@ def async_setup(hass, config):
 
     discovery.async_listen(hass, SERVICE_IKEA_TRADFRI, gateway_discovered)
 
-    if host is None:
+    if not host:
         return True
 
-    return (yield from _setup_gateway(hass, config, host, key,
-                                      allow_tradfri_groups))
+    if host and keys.get(host):
+        return (yield from _setup_gateway(hass, config, host,
+                                          keys[host]['key'],
+                                          allow_tradfri_groups))
+    else:
+        hass.async_add_job(request_configuration, hass, config, host)
+        return True
 
 
 @asyncio.coroutine
@@ -116,19 +128,21 @@ def _setup_gateway(hass, hass_config, host, key, allow_tradfri_groups):
     """Create a gateway."""
     from pytradfri import Gateway, RequestError
     try:
-        from pytradfri.api.aiocoap_api import api_factory
+        from pytradfri.api.aiocoap_api import APIFactory
     except ImportError:
         _LOGGER.exception("Looks like something isn't installed!")
         return False
 
     try:
-        api = yield from api_factory(host, key, loop=hass.loop)
+        factory = APIFactory(host, psk_id=GATEWAY_IDENTITY, psk=key,
+                             loop=hass.loop)
+        api = factory.request
+        gateway = Gateway()
+        gateway_info_result = yield from api(gateway.get_gateway_info())
     except RequestError:
         _LOGGER.exception("Tradfri setup failed.")
         return False
 
-    gateway = Gateway()
-    gateway_info_result = yield from api(gateway.get_gateway_info())
     gateway_id = gateway_info_result.id
     hass.data.setdefault(KEY_API, {})
     hass.data.setdefault(KEY_GATEWAY, {})
diff --git a/requirements_all.txt b/requirements_all.txt
index 0ae8a3d715c..526aecd9399 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -893,7 +893,7 @@ pythonwhois==2.4.3
 pytrackr==0.0.5
 
 # homeassistant.components.tradfri
-pytradfri==3.0.2
+pytradfri==4.0.1
 
 # homeassistant.components.device_tracker.unifi
 pyunifi==2.13