Working on igd

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
-from homeassistant.helpers.typing import ConfigType, HomeAssistantType
+from homeassistant.helpers.typing import HomeAssistantType
-from homeassistant.util import get_local_ip
 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 udn not in hass.data[DOMAIN]['devices']:
-    if igd_device is None:
         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

homeassistant/components/igd/config_flow.py

@@ -21,6 +21,7 @@ def ensure_domain_data(hass):
         'active': False,
         'port_forward': False,
         'sensors': False,
+        'ports': {},
     })
 
 

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')

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()

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(device, IN),
-        KBytePerSecondIGDSensor(igd_device, OUT),
+        KBytePerSecondIGDSensor(device, OUT),
-        PacketsPerSecondIGDSensor(igd_device, IN),
+        PacketsPerSecondIGDSensor(device, IN),
-        PacketsPerSecondIGDSensor(igd_device, OUT),
+        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:

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 homeassistant.components.igd.device.Device
-    mock_igd_device = MagicMock()
+    mock_device = MagicMock()
-    mock_igd_device.udn = udn
+    mock_device.udn = udn
-    mock_igd_device.async_add_port_mapping.return_value = mock_coro()
+    mock_device.async_add_port_mappings.return_value = mock_coro()
-    mock_igd_device.async_delete_port_mapping.return_value = mock_coro()
+    mock_device.async_delete_port_mappings.return_value = mock_coro()
-    with patch.object(igd, '_async_create_igd_device') as mock_create_device:
+    with patch.object(Device, 'async_create_device') as mock_create_device:
         mock_create_device.return_value = mock_coro(
-            return_value=mock_igd_device)
+            return_value=mock_device)
-        with patch('homeassistant.components.igd.get_local_ip',
+        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 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_device = MockDevice(udn)
-    mock_igd_device.udn = udn
+    with patch.object(Device, 'async_create_device') as mock_create_device:
-    mock_igd_device.async_add_port_mapping.return_value = mock_coro()
+        mock_create_device.return_value = mock_coro(return_value=mock_device)
-    mock_igd_device.async_delete_port_mapping.return_value = mock_coro()
+        with patch('homeassistant.components.igd.device.get_local_ip',
-    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',
                    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
+
+            # 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]