evlli/hass-core
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Working on igd

Browse source
This commit is contained in:
Steven Looman 2018-09-08 00:11:23 +02:00
parent f4a54e2483
commit 6433f2e2e3
6 changed files with 255 additions and 134 deletions
homeassistant/components
igd
__init__.pyconfig_flow.pyconst.pydevice.py
sensor
igd.py
tests/components/igd
test_init.py

128
homeassistant/components/igd/__init__.py
View file

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

1
homeassistant/components/igd/config_flow.py
View file

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

9
homeassistant/components/igd/const.py
View file

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

131
homeassistant/components/igd/device.py Normal file
View file

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

16
homeassistant/components/sensor/igd.py
View file

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

104
tests/components/igd/test_init.py
View file

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