From 789929d4457fd9bff7710d46171000a4b6c4e3aa Mon Sep 17 00:00:00 2001
From: cgtobi <cgtobi@users.noreply.github.com>
Date: Sat, 21 Oct 2017 20:45:53 +0200
Subject: [PATCH] Add support for multiple disks to be monitored. (#9977)

* Add support for multiple disks.

* Fix LINT error.

* Make disk config optional to not break existing installations.

* Change state handling as per request by @fabaff.
---
 homeassistant/components/sensor/hddtemp.py | 49 +++++++++++++++-------
 1 file changed, 34 insertions(+), 15 deletions(-)

diff --git a/homeassistant/components/sensor/hddtemp.py b/homeassistant/components/sensor/hddtemp.py
index cd84bd8f9a8..fef84091b18 100644
--- a/homeassistant/components/sensor/hddtemp.py
+++ b/homeassistant/components/sensor/hddtemp.py
@@ -22,6 +22,8 @@ _LOGGER = logging.getLogger(__name__)
 ATTR_DEVICE = 'device'
 ATTR_MODEL = 'model'
 
+CONF_DISKS = 'disks'
+
 DEFAULT_HOST = 'localhost'
 DEFAULT_PORT = 7634
 DEFAULT_NAME = 'HD Temperature'
@@ -30,6 +32,8 @@ DEFAULT_TIMEOUT = 5
 SCAN_INTERVAL = timedelta(minutes=1)
 
 PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
+    vol.Optional(CONF_DISKS, default=[]):
+        vol.All(cv.ensure_list, [cv.string]),
     vol.Optional(CONF_HOST, default=DEFAULT_HOST): cv.string,
     vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
     vol.Optional(CONF_NAME, default=DEFAULT_NAME): cv.string,
@@ -41,25 +45,37 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
     name = config.get(CONF_NAME)
     host = config.get(CONF_HOST)
     port = config.get(CONF_PORT)
+    disks = config.get(CONF_DISKS)
 
-    hddtemp = HddTempData(host, port)
-    hddtemp.update()
-
-    if hddtemp.data is None:
+    try:
+        hddtemp = HddTempData(host, port)
+        hddtemp.update()
+    except RuntimeError:
         _LOGGER.error("Unable to fetch the data from %s:%s", host, port)
         return False
 
-    add_devices([HddTempSensor(name, hddtemp)], True)
+    if not disks:
+        disks = [next(iter(hddtemp.data)).split('|')[0]]
+
+    dev = []
+    for disk in disks:
+        if disk in hddtemp.data:
+            dev.append(HddTempSensor(name, disk, hddtemp))
+        else:
+            continue
+
+    add_devices(dev, True)
 
 
 class HddTempSensor(Entity):
     """Representation of a HDDTemp sensor."""
 
-    def __init__(self, name, hddtemp):
+    def __init__(self, name, disk, hddtemp):
         """Initialize a HDDTemp sensor."""
         self.hddtemp = hddtemp
-        self._name = name
-        self._state = False
+        self.disk = disk
+        self._name = '{} {}'.format(name, disk)
+        self._state = None
         self._details = None
 
     @property
@@ -75,7 +91,7 @@ class HddTempSensor(Entity):
     @property
     def unit_of_measurement(self):
         """Return the unit the value is expressed in."""
-        if self._details[4] == 'C':
+        if self._details[3] == 'C':
             return TEMP_CELSIUS
         return TEMP_FAHRENHEIT
 
@@ -83,17 +99,17 @@ class HddTempSensor(Entity):
     def device_state_attributes(self):
         """Return the state attributes of the sensor."""
         return {
-            ATTR_DEVICE: self._details[1],
-            ATTR_MODEL: self._details[2],
+            ATTR_DEVICE: self._details[0],
+            ATTR_MODEL: self._details[1],
         }
 
     def update(self):
         """Get the latest data from HDDTemp daemon and updates the state."""
         self.hddtemp.update()
 
-        if self.hddtemp.data is not None:
-            self._details = self.hddtemp.data.split('|')
-            self._state = self._details[3]
+        if self.disk in self.hddtemp.data:
+            self._details = self.hddtemp.data[self.disk].split('|')
+            self._state = self._details[2]
         else:
             self._state = STATE_UNKNOWN
 
@@ -112,7 +128,10 @@ class HddTempData(object):
         try:
             connection = Telnet(
                 host=self.host, port=self.port, timeout=DEFAULT_TIMEOUT)
-            self.data = connection.read_all().decode('ascii')
+            data = connection.read_all().decode(
+                'ascii').lstrip('|').rstrip('|').split('||')
+            self.data = {data[i].split('|')[0]: data[i]
+                         for i in range(0, len(data), 1)}
         except ConnectionRefusedError:
             _LOGGER.error(
                 "HDDTemp is not available at %s:%s", self.host, self.port)