abuseipdb.py

import logging
from datetime import timedelta

import voluptuous as vol

from homeassistant.const import CONF_IP_ADDRESS, CONF_API_KEY

from homeassistant.components.sensor import PLATFORM_SCHEMA

from homeassistant.helpers.entity import Entity
import homeassistant.helpers.config_validation as cv
from homeassistant.util import Throttle

PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
    vol.Required(CONF_IP_ADDRESS): vol.All(cv.ensure_list, [cv.string]),
    vol.Required(CONF_API_KEY): cv.string,
})


class Abuseipdb:

    URL = "https://www.abuseipdb.com/check/%s/json?key=%s"

    @staticmethod
    async def fetch(session, url):
        import json
        async with session.get(url) as response:
            return json.loads(await response.text())

    @staticmethod
    async def get_data(ip, api_key):
        import aiohttp
        async with aiohttp.ClientSession() as session:
            url = Abuseipdb.URL % (ip, api_key)
            html = await Abuseipdb.fetch(
                session, url)
            return html


_LOGGER = logging.getLogger(__name__)

SCAN_INTERVAL = timedelta(minutes=5)


async def async_setup_platform(
        hass, config, async_add_entities, discovery_info=None):
    """Set up the asuswrt sensors."""
    entities = []
    for ip_adress in config[CONF_IP_ADDRESS]:
        data = await Abuseipdb.get_data(ip_adress, config[CONF_API_KEY])
        entities.append(
            AbuseipdbSensor(data, ip_adress, config[CONF_API_KEY]),
        )
    async_add_entities(entities)


class AbuseipdbSensor(Entity):
    """Representation of a kris scale sensor."""

    def __init__(self, data, ip, api_key):
        """Init the sensor."""
        self._unit = "Abuse reports"
        self._attributes = {}
        self._state = 0
        if data:
            self._state = len(data)
            self.filter_attributes(data)
        self._name = "%s abuseipdb reports" % ip
        self.api_key = api_key
        self.ip = ip

    @property
    def unit_of_measurement(self):
        """Return the unit of measurement."""
        return self._unit

    @property
    def state(self):
        """Property for the state attributes."""
        return self._state

    @property
    def name(self):
        """Name property for sensor."""
        return self._name

    @property
    def state_attributes(self):
        """Property for the state attributes."""
        return self._attributes

    def filter_attributes(self, data):
        """Filter out data for this Sensor."""
        self._attributes = {
            "Category": data[0]['category'],
            "Country": data[0]['country'],
            "Confidence": data[0]['abuseConfidenceScore']
        }

    @Throttle(SCAN_INTERVAL)
    async def async_update(self):
        """Fetch new state data for the sensor."""
        data = await Abuseipdb.get_data(self.ip, self.api_key)
        if data:
            self._state = len(data)
            self.filter_attributes(data)