diff --git a/homeassistant/components/igd/__init__.py b/homeassistant/components/igd/__init__.py index 4f7bc9606ac..118cd82be12 100644 --- a/homeassistant/components/igd/__init__.py +++ b/homeassistant/components/igd/__init__.py @@ -6,7 +6,6 @@ https://home-assistant.io/components/igd/ """ import asyncio -from ipaddress import IPv4Address from ipaddress import ip_address import aiohttp @@ -17,26 +16,24 @@ from homeassistant.const import EVENT_HOMEASSISTANT_STOP from homeassistant.exceptions import PlatformNotReady from homeassistant.helpers import config_validation as cv from homeassistant.helpers import discovery -from homeassistant.helpers.aiohttp_client import async_get_clientsession -from homeassistant.helpers.typing import ConfigType, HomeAssistantType -from homeassistant.util import get_local_ip +from homeassistant.helpers.typing import ConfigType +from homeassistant.helpers.typing import HomeAssistantType from homeassistant.components.discovery import DOMAIN as DISCOVERY_DOMAIN from .const import ( CONF_ENABLE_PORT_MAPPING, CONF_ENABLE_SENSORS, - CONF_UDN, CONF_SSDP_DESCRIPTION + CONF_HASS, CONF_LOCAL_IP, CONF_PORTS, + CONF_UDN, CONF_SSDP_DESCRIPTION, ) from .const import DOMAIN from .const import LOGGER as _LOGGER from .config_flow import ensure_domain_data +from .device import Device REQUIREMENTS = ['async-upnp-client==0.12.4'] DEPENDENCIES = ['http'] -CONF_LOCAL_IP = 'local_ip' -CONF_PORTS = 'ports' - NOTIFICATION_ID = 'igd_notification' NOTIFICATION_TITLE = 'UPnP/IGD Setup' @@ -45,90 +42,16 @@ CONFIG_SCHEMA = vol.Schema({ vol.Optional(CONF_ENABLE_PORT_MAPPING, default=False): cv.boolean, vol.Optional(CONF_ENABLE_SENSORS, default=True): cv.boolean, vol.Optional(CONF_LOCAL_IP): vol.All(ip_address, cv.string), + vol.Optional(CONF_PORTS): + vol.Schema({vol.Any(CONF_HASS, cv.positive_int): cv.positive_int}) }), }, extra=vol.ALLOW_EXTRA) -async def _async_create_igd_device(hass: HomeAssistantType, - ssdp_description: str): - """Create IGD device.""" - # build requester - from async_upnp_client.aiohttp import AiohttpSessionRequester - session = async_get_clientsession(hass) - requester = AiohttpSessionRequester(session, True) - - # create upnp device - from async_upnp_client import UpnpFactory - factory = UpnpFactory(requester, disable_state_variable_validation=True) - try: - upnp_device = await factory.async_create_device(ssdp_description) - except (asyncio.TimeoutError, aiohttp.ClientError): - raise PlatformNotReady() - - # wrap with IgdDevice - from async_upnp_client.igd import IgdDevice - igd_device = IgdDevice(upnp_device, None) - return igd_device - - -def _store_device(hass: HomeAssistantType, udn, igd_device): - """Store an igd_device by udn.""" - if igd_device is not None: - hass.data[DOMAIN]['devices'][udn] = igd_device - elif udn in hass.data[DOMAIN]['devices']: - del hass.data[DOMAIN]['devices'][udn] - - -def _get_device(hass: HomeAssistantType, udn): - """Get an igd_device by udn.""" - return hass.data[DOMAIN]['devices'].get(udn) - - -async def _async_add_port_mapping(hass: HomeAssistantType, - igd_device, - local_ip=None): - """Create a port mapping.""" - # determine local ip, ensure sane IP - if local_ip is None: - local_ip = get_local_ip() - - if local_ip == '127.0.0.1': - _LOGGER.warning('Could not create port mapping, our IP is 127.0.0.1') - return - local_ip = IPv4Address(local_ip) - - # create port mapping - from async_upnp_client import UpnpError - port = hass.http.server_port - _LOGGER.debug('Creating port mapping %s:%s:%s (TCP)', port, local_ip, port) - try: - await igd_device.async_add_port_mapping(remote_host=None, - external_port=port, - protocol='TCP', - internal_port=port, - internal_client=local_ip, - enabled=True, - description="Home Assistant", - lease_duration=None) - except (asyncio.TimeoutError, aiohttp.ClientError, UpnpError): - _LOGGER.warning('Could not add port mapping') - - -async def _async_delete_port_mapping(hass: HomeAssistantType, igd_device): - """Remove a port mapping.""" - from async_upnp_client import UpnpError - port = hass.http.server_port - try: - await igd_device.async_delete_port_mapping(remote_host=None, - external_port=port, - protocol='TCP') - except (asyncio.TimeoutError, aiohttp.ClientError, UpnpError): - _LOGGER.warning('Could not delete port mapping') - - # config async def async_setup(hass: HomeAssistantType, config: ConfigType): """Register a port mapping for Home Assistant via UPnP.""" + _LOGGER.debug('async_setup: %s', config.get(DOMAIN)) ensure_domain_data(hass) # ensure sane config @@ -143,12 +66,22 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType): if CONF_LOCAL_IP in igd_config: hass.data[DOMAIN]['local_ip'] = igd_config[CONF_LOCAL_IP] + # determine ports + ports = {} + if CONF_PORTS in igd_config: + ports = igd_config[CONF_PORTS] + + if CONF_HASS in ports: + internal_port = hass.http.server_port + ports[internal_port] = ports[CONF_HASS] + del ports[CONF_HASS] + hass.data[DOMAIN]['auto_config'] = { 'active': True, 'port_forward': igd_config[CONF_ENABLE_PORT_MAPPING], + 'ports': ports, 'sensors': igd_config[CONF_ENABLE_SENSORS], } - _LOGGER.debug('Enabled auto_config: %s', hass.data[DOMAIN]['auto_config']) return True @@ -157,27 +90,31 @@ async def async_setup(hass: HomeAssistantType, config: ConfigType): async def async_setup_entry(hass: HomeAssistantType, config_entry: ConfigEntry): """Set up a bridge from a config entry.""" + _LOGGER.debug('async_setup_entry: %s', config_entry.data) ensure_domain_data(hass) data = config_entry.data # build IGD device ssdp_description = data[CONF_SSDP_DESCRIPTION] try: - igd_device = await _async_create_igd_device(hass, ssdp_description) + device = await Device.async_create_device(hass, ssdp_description) except (asyncio.TimeoutError, aiohttp.ClientError): raise PlatformNotReady() - _store_device(hass, igd_device.udn, igd_device) + hass.data[DOMAIN]['devices'][device.udn] = device # port mapping if data.get(CONF_ENABLE_PORT_MAPPING): local_ip = hass.data[DOMAIN].get('local_ip') - await _async_add_port_mapping(hass, igd_device, local_ip=local_ip) + ports = hass.data[DOMAIN]['auto_config']['ports'] + _LOGGER.debug('Enabling port mappings: %s', ports) + await device.async_add_port_mappings(ports, local_ip=local_ip) # sensors if data.get(CONF_ENABLE_SENSORS): + _LOGGER.debug('Enabling sensors') discovery_info = { - 'udn': data[CONF_UDN], + 'udn': device.udn, } hass_config = config_entry.data hass.async_create_task(discovery.async_load_platform( @@ -194,22 +131,25 @@ async def async_setup_entry(hass: HomeAssistantType, async def async_unload_entry(hass: HomeAssistantType, config_entry: ConfigEntry): """Unload a config entry.""" + _LOGGER.debug('async_unload_entry: %s', config_entry.data) data = config_entry.data udn = data[CONF_UDN] - igd_device = _get_device(hass, udn) - if igd_device is None: + if udn not in hass.data[DOMAIN]['devices']: return True + device = hass.data[DOMAIN]['devices'][udn] # port mapping if data.get(CONF_ENABLE_PORT_MAPPING): - await _async_delete_port_mapping(hass, igd_device) + _LOGGER.debug('Deleting port mappings') + await device.async_delete_port_mappings() # sensors for sensor in hass.data[DOMAIN]['sensors'].get(udn, []): + _LOGGER.debug('Deleting sensor: %s', sensor) await sensor.async_remove() # clear stored device - _store_device(hass, udn, None) + del hass.data[DOMAIN]['devices'][udn] return True diff --git a/homeassistant/components/igd/config_flow.py b/homeassistant/components/igd/config_flow.py index 6f0cec12b5b..7af6c683ed9 100644 --- a/homeassistant/components/igd/config_flow.py +++ b/homeassistant/components/igd/config_flow.py @@ -21,6 +21,7 @@ def ensure_domain_data(hass): 'active': False, 'port_forward': False, 'sensors': False, + 'ports': {}, }) diff --git a/homeassistant/components/igd/const.py b/homeassistant/components/igd/const.py index 47ff9f11074..8ba774447da 100644 --- a/homeassistant/components/igd/const.py +++ b/homeassistant/components/igd/const.py @@ -2,9 +2,12 @@ import logging -DOMAIN = 'igd' -LOGGER = logging.getLogger('homeassistant.components.igd') CONF_ENABLE_PORT_MAPPING = 'port_forward' CONF_ENABLE_SENSORS = 'sensors' -CONF_UDN = 'udn' +CONF_HASS = 'hass' +CONF_LOCAL_IP = 'local_ip' +CONF_PORTS = 'ports' CONF_SSDP_DESCRIPTION = 'ssdp_description' +CONF_UDN = 'udn' +DOMAIN = 'igd' +LOGGER = logging.getLogger('homeassistant.components.igd') diff --git a/homeassistant/components/igd/device.py b/homeassistant/components/igd/device.py new file mode 100644 index 00000000000..af16623e4b2 --- /dev/null +++ b/homeassistant/components/igd/device.py @@ -0,0 +1,131 @@ +"""Hass representation of an IGD.""" + +import asyncio +from ipaddress import IPv4Address + +import aiohttp + +from homeassistant.helpers.aiohttp_client import async_get_clientsession +from homeassistant.helpers.typing import HomeAssistantType +from homeassistant.util import get_local_ip + +from .const import LOGGER as _LOGGER + + +class Device: + """Hass representation of an IGD.""" + + def __init__(self, igd_device): + """Initializer.""" + self._igd_device = igd_device + self._mapped_ports = [] + + @classmethod + async def async_create_device(cls, + hass: HomeAssistantType, + ssdp_description: str): + """Create IGD device.""" + # build async_upnp_client requester + from async_upnp_client.aiohttp import AiohttpSessionRequester + session = async_get_clientsession(hass) + requester = AiohttpSessionRequester(session, True) + + # create async_upnp_client device + from async_upnp_client import UpnpFactory + factory = UpnpFactory(requester, + disable_state_variable_validation=True) + upnp_device = await factory.async_create_device(ssdp_description) + + # wrap with async_upnp_client IgdDevice + from async_upnp_client.igd import IgdDevice + igd_device = IgdDevice(upnp_device, None) + + return Device(igd_device) + + @property + def udn(self): + """Get the UDN.""" + return self._igd_device.udn + + @property + def name(self): + """Get the name.""" + return self._igd_device.name + + async def async_add_port_mappings(self, ports, local_ip=None): + """Add port mappings.""" + # determine local ip, ensure sane IP + if local_ip is None: + local_ip = get_local_ip() + + if local_ip == '127.0.0.1': + _LOGGER.error( + 'Could not create port mapping, our IP is 127.0.0.1') + local_ip = IPv4Address(local_ip) + + # create port mappings + for external_port, internal_port in ports.items(): + await self._async_add_port_mapping(external_port, + local_ip, + internal_port) + self._mapped_ports.append(external_port) + + async def _async_add_port_mapping(self, + external_port, + local_ip, + internal_port): + """Add a port mapping.""" + # create port mapping + from async_upnp_client import UpnpError + _LOGGER.info('Creating port mapping %s:%s:%s (TCP)', + external_port, local_ip, internal_port) + try: + await self._igd_device.async_add_port_mapping( + remote_host=None, + external_port=external_port, + protocol='TCP', + internal_port=internal_port, + internal_client=local_ip, + enabled=True, + description="Home Assistant", + lease_duration=None) + + self._mapped_ports.append(external_port) + except (asyncio.TimeoutError, aiohttp.ClientError, UpnpError): + _LOGGER.error('Could not add port mapping: %s:%s:%s', + external_port, local_ip, internal_port) + + async def async_delete_port_mappings(self): + """Remove a port mapping.""" + for port in self._mapped_ports: + await self._async_delete_port_mapping(port) + + async def _async_delete_port_mapping(self, external_port): + """Remove a port mapping.""" + from async_upnp_client import UpnpError + try: + await self._igd_device.async_delete_port_mapping( + remote_host=None, + external_port=external_port, + protocol='TCP') + + self._mapped_ports.remove(external_port) + except (asyncio.TimeoutError, aiohttp.ClientError, UpnpError): + _LOGGER.error('Could not delete port mapping') + + async def async_get_total_bytes_received(self): + """Get total bytes received.""" + return await self._igd_device.async_get_total_bytes_received() + + async def async_get_total_bytes_sent(self): + """Get total bytes sent.""" + return await self._igd_device.async_get_total_bytes_sent() + + async def async_get_total_packets_received(self): + """Get total packets received.""" + # pylint: disable=invalid-name + return await self._igd_device.async_get_total_packets_received() + + async def async_get_total_packets_sent(self): + """Get total packets sent.""" + return await self._igd_device.async_get_total_packets_sent() diff --git a/homeassistant/components/sensor/igd.py b/homeassistant/components/sensor/igd.py index 908e465aedb..ff93af89f34 100644 --- a/homeassistant/components/sensor/igd.py +++ b/homeassistant/components/sensor/igd.py @@ -14,7 +14,7 @@ from homeassistant.helpers.entity import Entity _LOGGER = logging.getLogger(__name__) -DEPENDENCIES = ['igd', 'history'] +DEPENDENCIES = ['igd'] BYTES_RECEIVED = 'bytes_received' BYTES_SENT = 'bytes_sent' @@ -52,18 +52,18 @@ async def async_setup_platform(hass, config, async_add_devices, return udn = discovery_info['udn'] - igd_device = hass.data[DOMAIN]['devices'][udn] + device = hass.data[DOMAIN]['devices'][udn] # raw sensors + per-second sensors sensors = [ - RawIGDSensor(igd_device, name, sensor_type) + RawIGDSensor(device, name, sensor_type) for name, sensor_type in SENSOR_TYPES.items() ] sensors += [ - KBytePerSecondIGDSensor(igd_device, IN), - KBytePerSecondIGDSensor(igd_device, OUT), - PacketsPerSecondIGDSensor(igd_device, IN), - PacketsPerSecondIGDSensor(igd_device, OUT), + KBytePerSecondIGDSensor(device, IN), + KBytePerSecondIGDSensor(device, OUT), + PacketsPerSecondIGDSensor(device, IN), + PacketsPerSecondIGDSensor(device, OUT), ] hass.data[DOMAIN]['sensors'][udn] = sensors async_add_devices(sensors, True) @@ -108,7 +108,6 @@ class RawIGDSensor(Entity): async def async_update(self): """Get the latest information from the IGD.""" - _LOGGER.debug('%s: async_update', self) if self._type_name == BYTES_RECEIVED: self._state = await self._device.async_get_total_bytes_received() elif self._type_name == BYTES_SENT: @@ -171,7 +170,6 @@ class PerSecondIGDSensor(Entity): async def async_update(self): """Get the latest information from the IGD.""" - _LOGGER.debug('%s: async_update', self) new_value = await self._async_fetch_value() if self._last_value is None: diff --git a/tests/components/igd/test_init.py b/tests/components/igd/test_init.py index cd3f5324bb4..336935710d8 100644 --- a/tests/components/igd/test_init.py +++ b/tests/components/igd/test_init.py @@ -1,15 +1,51 @@ """Test IGD setup process.""" +from ipaddress import ip_address from unittest.mock import patch, MagicMock from homeassistant.setup import async_setup_component from homeassistant.components import igd +from homeassistant.components.igd.device import Device from homeassistant.const import EVENT_HOMEASSISTANT_STOP from tests.common import MockConfigEntry from tests.common import mock_coro +class MockDevice(Device): + """Mock device for Device.""" + + def __init__(self, udn): + """Initializer.""" + super().__init__(None) + self._udn = udn + self.added_port_mappings = [] + self.removed_port_mappings = [] + + @classmethod + async def async_create_device(cls, hass, ssdp_description): + """Return self.""" + return cls() + + @property + def udn(self): + """Get the UDN.""" + return self._udn + + async def _async_add_port_mapping(self, + external_port, + local_ip, + internal_port): + """Add a port mapping.""" + entry = [external_port, local_ip, internal_port] + self.added_port_mappings.append(entry) + + async def _async_delete_port_mapping(self, external_port): + """Remove a port mapping.""" + entry = external_port + self.removed_port_mappings.append(entry) + + async def test_async_setup_no_auto_config(hass): """Test async_setup.""" # setup component, enable auto_config @@ -18,6 +54,7 @@ async def test_async_setup_no_auto_config(hass): assert hass.data[igd.DOMAIN]['auto_config'] == { 'active': False, 'port_forward': False, + 'ports': {}, 'sensors': False, } @@ -30,6 +67,7 @@ async def test_async_setup_auto_config(hass): assert hass.data[igd.DOMAIN]['auto_config'] == { 'active': True, 'port_forward': False, + 'ports': {}, 'sensors': True, } @@ -38,12 +76,13 @@ async def test_async_setup_auto_config_port_forward(hass): """Test async_setup.""" # setup component, enable auto_config await async_setup_component(hass, 'igd', { - 'igd': {'port_forward': True}, + 'igd': {'port_forward': True, 'ports': {8123: 8123}}, 'discovery': {}}) assert hass.data[igd.DOMAIN]['auto_config'] == { 'active': True, 'port_forward': True, + 'ports': {8123: 8123}, 'sensors': True, } @@ -58,6 +97,7 @@ async def test_async_setup_auto_config_no_sensors(hass): assert hass.data[igd.DOMAIN]['auto_config'] == { 'active': True, 'port_forward': False, + 'ports': {}, 'sensors': False, } @@ -75,27 +115,30 @@ async def test_async_setup_entry_default(hass): # ensure hass.http is available await async_setup_component(hass, 'igd') - # mock async_upnp_client.igd.IgdDevice - mock_igd_device = MagicMock() - mock_igd_device.udn = udn - mock_igd_device.async_add_port_mapping.return_value = mock_coro() - mock_igd_device.async_delete_port_mapping.return_value = mock_coro() - with patch.object(igd, '_async_create_igd_device') as mock_create_device: + # mock homeassistant.components.igd.device.Device + mock_device = MagicMock() + mock_device.udn = udn + mock_device.async_add_port_mappings.return_value = mock_coro() + mock_device.async_delete_port_mappings.return_value = mock_coro() + with patch.object(Device, 'async_create_device') as mock_create_device: mock_create_device.return_value = mock_coro( - return_value=mock_igd_device) - with patch('homeassistant.components.igd.get_local_ip', + return_value=mock_device) + with patch('homeassistant.components.igd.device.get_local_ip', return_value='192.168.1.10'): assert await igd.async_setup_entry(hass, entry) is True # ensure device is stored/used - assert hass.data[igd.DOMAIN]['devices'][udn] == mock_igd_device + assert hass.data[igd.DOMAIN]['devices'][udn] == mock_device - hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) - await hass.async_block_till_done() + hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) + await hass.async_block_till_done() + # ensure cleaned up assert udn not in hass.data[igd.DOMAIN]['devices'] - assert len(mock_igd_device.async_add_port_mapping.mock_calls) == 0 - assert len(mock_igd_device.async_delete_port_mapping.mock_calls) == 0 + + # ensure no port-mapping-methods called + assert len(mock_device.async_add_port_mappings.mock_calls) == 0 + assert len(mock_device.async_delete_port_mappings.mock_calls) == 0 async def test_async_setup_entry_port_forward(hass): @@ -109,25 +152,30 @@ async def test_async_setup_entry_port_forward(hass): }) # ensure hass.http is available - await async_setup_component(hass, 'igd') + await async_setup_component(hass, 'igd', { + 'igd': {'port_forward': True, 'ports': {8123: 8123}}, + 'discovery': {}}) - mock_igd_device = MagicMock() - mock_igd_device.udn = udn - mock_igd_device.async_add_port_mapping.return_value = mock_coro() - mock_igd_device.async_delete_port_mapping.return_value = mock_coro() - with patch.object(igd, '_async_create_igd_device') as mock_create_device: - mock_create_device.return_value = mock_coro( - return_value=mock_igd_device) - with patch('homeassistant.components.igd.get_local_ip', + mock_device = MockDevice(udn) + with patch.object(Device, 'async_create_device') as mock_create_device: + mock_create_device.return_value = mock_coro(return_value=mock_device) + with patch('homeassistant.components.igd.device.get_local_ip', return_value='192.168.1.10'): assert await igd.async_setup_entry(hass, entry) is True # ensure device is stored/used - assert hass.data[igd.DOMAIN]['devices'][udn] == mock_igd_device + assert hass.data[igd.DOMAIN]['devices'][udn] == mock_device - hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) - await hass.async_block_till_done() + # ensure add-port-mapping-methods called + assert mock_device.added_port_mappings == [ + [8123, ip_address('192.168.1.10'), 8123] + ] + hass.bus.async_fire(EVENT_HOMEASSISTANT_STOP) + await hass.async_block_till_done() + + # ensure cleaned up assert udn not in hass.data[igd.DOMAIN]['devices'] - assert len(mock_igd_device.async_add_port_mapping.mock_calls) > 0 - assert len(mock_igd_device.async_delete_port_mapping.mock_calls) > 0 + + # ensure delete-port-mapping-methods called + assert mock_device.removed_port_mappings == [8123]