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
pull/10880/merge
Adam Cooper 2017-12-01 16:36:15 +00:00 committed by Adam Mills
parent bc4de4e769
commit c2525782aa
1 changed files with 36 additions and 26 deletions

View File

@ -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