lightbulb.py

from checks import AgentCheck
from datadog import initialize, api

import requests
import simplejson as json
from pyHS100 import SmartBulb
import time
from datetime import datetime, timedelta

from threading import Thread

OK = "Ok"
CRITICAL = "Alert"
WARNING = "Warn"
OUTAGE = "OUTAGE"

GREEN = {
    u'color_temp': 0,
    u'hue': 97,
    u'saturation': 93,
    u'on_off': 1,
    u'brightness': 40,
    u'mode': u'normal'
}

RED = {
    u'color_temp': 0,
    u'hue': 350,
    u'saturation': 94,
    u'on_off': 1,
    u'brightness': 40,
    u'mode': u'normal'
}

ORANGE = {
    u'color_temp': 0,
    u'hue': 37,
    u'saturation': 98,
    u'on_off': 1,
    u'brightness': 40,
    u'mode': u'normal'
}

class Bulb():
    def __init__(self, ip):
            self.bulb = SmartBulb(ip)
            self.action = None
            self.thread = Thread(target = self.run)
            self.thread.start()

    def run_action(self, action):
        self.action = getattr(self, action)

    def run(self):
        while True:
            time.sleep(1)
            if self.action is not None:
                self.action()

    def blink_red(self):
        self.make_red()
        self.turn_off()
        time.sleep(1)
        self.turn_on()
        time.sleep(1)

    def do_nothing(self):
        pass

    def graceful_turn_off(self):
        self.make_green()
        time.sleep(5)
        self.turn_off()
        self.run_action("do_nothing")

    def turn_off(self):
        self.bulb.state = "OFF"

    def turn_on(self):
        self.bulb.state = "ON"

    def make_red(self):
        self.bulb.set_light_state(RED)

    def make_green(self):
        self.bulb.set_light_state(GREEN)

    def make_orange(self):
        self.bulb.set_light_state(ORANGE)

class LightBulb(AgentCheck):

    def get_bulb(self, ip):
        if not hasattr(self, "cache"):
            self.cache = {}

        if not ip in self.cache:
            self.cache[ip] = [Bulb(ip), None]

        return self.cache[ip]

    def set_state(self, ip, state):
        self.cache[ip][1] = state

    def check(self, instance):

        monitor_filter = instance['monitor_filter']
        ip = instance['ip_addr']
        dd_api_key = instance['dd_api_key']
        dd_app_key = instance['dd_app_key']
        pd_api_key = instance['pd_api_key']
        pd_service_ids = instance['pd_service_ids']
        pd_subdomain = instance['pd_subdomain']

        bulb, state = self.get_bulb(ip)

        initialize(api_key=dd_api_key, app_key=dd_app_key)

        should_alert_on_monitors = self.should_alert_on_monitors(monitor_filter)
        should_alert_on_pd = self.should_alert_on_pd(
            pd_subdomain, pd_service_ids, pd_api_key
        )

        if should_alert_on_pd:
            self.set_state(ip, OUTAGE)
            bulb.run_action("blink_red")
        elif should_alert_on_monitors == CRITICAL:
            self.set_state(ip, CRITICAL)
            bulb.run_action("make_red")
        elif should_alert_on_monitors == WARNING:
            self.set_state(ip, WARNING)
            bulb.run_action("make_orange")
        elif state != GREEN:
            self.set_state(ip, GREEN)
            bulb.run_action("graceful_turn_off")

    def should_alert_on_monitors(self, monitor_filter):

        monitors = api.Monitor.get_all(with_downtimes=False, name=monitor_filter)
        self.log.info("API returned {} monitors".format(len(monitors)))
        for m in monitors:
            if m['overall_state'] == CRITICAL:
                self.log.info("Found at least one critical monitor")
                return CRITICAL
            if m['overall_state'] == WARNING:
                self.log.info("Found at least one warning monitor")
                return WARNING
        self.log.info("Didn't find any bad monitor")
        return OK

    def should_alert_on_pd(self, pd_subdomain, pd_service_ids, pd_api_key):

        since = datetime.utcnow() - timedelta(hours=1)
        since = since.isoformat() + "Z"
        services = ','.join(pd_service_ids)
        headers = {"Authorization": "Token token={}".format(pd_api_key)}
        url = "https://{}.pagerduty.com/api/v1/incidents".format(pd_subdomain)

        params = {
            "since": since,
            "service": services,
            "status": "triggered,acknowledged",
        }

        r = requests.get(url, params=params, headers=headers)
        total = r.json()['total']
        self.log.info("Found {} pagerduty incidents".format(total))
        return total > 0