From c2525782aa12a8a74219339ff3d418e623c2fb68 Mon Sep 17 00:00:00 2001
From: Adam Cooper <adam@genericstudent.com>
Date: Fri, 1 Dec 2017 16:36:15 +0000
Subject: [PATCH] Refactored WHOIS sensor to resolve assumed key errors
 (#10662)

* Refactored WHOIS sensor to resolve assumed key errors

Altered it to now set an attribute key and value only if the attribute
is present in the WHOIS response. This prevents assumed keys (registrar)
from raising a KeyError on WHOIS lookups that don't contain registrar
information (onet.pl, wp.pl, for example).

* Removed non-used self._data

* WHOIS sensor now creates a new local attributes dict and overrides

* Corrected typos, refactored error cases to clear state adn attributes

* Resolved double return and refactored error logging
---
 homeassistant/components/sensor/whois.py | 62 ++++++++++++++----------
 1 file changed, 36 insertions(+), 26 deletions(-)

diff --git a/homeassistant/components/sensor/whois.py b/homeassistant/components/sensor/whois.py
index 9f50a4c13db..771c4bc9d73 100644
--- a/homeassistant/components/sensor/whois.py
+++ b/homeassistant/components/sensor/whois.py
@@ -47,14 +47,13 @@ def setup_platform(hass, config, add_devices, discovery_info=None):
         if 'expiration_date' in get_whois(domain, normalized=True):
             add_devices([WhoisSensor(name, domain)], True)
         else:
-            _LOGGER.warning(
+            _LOGGER.error(
                 "WHOIS lookup for %s didn't contain expiration_date",
                 domain)
             return
     except WhoisException as ex:
-        _LOGGER.error("Exception %s occurred during WHOIS lookup for %s",
-                      ex,
-                      domain)
+        _LOGGER.error(
+            "Exception %s occurred during WHOIS lookup for %s", ex, domain)
         return
 
 
@@ -71,10 +70,7 @@ class WhoisSensor(Entity):
         self._domain = domain
 
         self._state = None
-        self._data = None
-        self._updated_date = None
-        self._expiration_date = None
-        self._name_servers = []
+        self._attributes = None
 
     @property
     def name(self):
@@ -99,38 +95,52 @@ class WhoisSensor(Entity):
     @property
     def device_state_attributes(self):
         """Get the more info attributes."""
-        if self._data:
-            updated_formatted = self._updated_date.isoformat()
-            expires_formatted = self._expiration_date.isoformat()
+        return self._attributes
 
-            return {
-                ATTR_NAME_SERVERS: ' '.join(self._name_servers),
-                ATTR_REGISTRAR: self._data['registrar'][0],
-                ATTR_UPDATED: updated_formatted,
-                ATTR_EXPIRES: expires_formatted,
-            }
+    def _empty_state_and_attributes(self):
+        """Empty the state and attributes on an error."""
+        self._state = None
+        self._attributes = None
 
     def update(self):
-        """Get the current WHOIS data for hostname."""
+        """Get the current WHOIS data for the domain."""
         from pythonwhois.shared import WhoisException
 
         try:
             response = self.whois(self._domain, normalized=True)
         except WhoisException as ex:
             _LOGGER.error("Exception %s occurred during WHOIS lookup", ex)
+            self._empty_state_and_attributes()
             return
 
         if response:
-            self._data = response
+            if 'expiration_date' not in response:
+                _LOGGER.error(
+                    "Failed to find expiration_date in whois lookup response. "
+                    "Did find: %s", ', '.join(response.keys()))
+                self._empty_state_and_attributes()
+                return
 
-            if self._data['nameservers']:
-                self._name_servers = self._data['nameservers']
+            if not response['expiration_date']:
+                _LOGGER.error("Whois response contains empty expiration_date")
+                self._empty_state_and_attributes()
+                return
 
-            if 'expiration_date' in self._data:
-                self._expiration_date = self._data['expiration_date'][0]
-            if 'updated_date' in self._data:
-                self._updated_date = self._data['updated_date'][0]
+            attrs = {}
 
-            time_delta = (self._expiration_date - self._expiration_date.now())
+            expiration_date = response['expiration_date'][0]
+            attrs[ATTR_EXPIRES] = expiration_date.isoformat()
 
+            if 'nameservers' in response:
+                attrs[ATTR_NAME_SERVERS] = ' '.join(response['nameservers'])
+
+            if 'updated_date' in response:
+                attrs[ATTR_UPDATED] = response['updated_date'][0].isoformat()
+
+            if 'registrar' in response:
+                attrs[ATTR_REGISTRAR] = response['registrar'][0]
+
+            time_delta = (expiration_date - expiration_date.now())
+
+            self._attributes = attrs
             self._state = time_delta.days