cold

sleeper_agents_aom_engine/reporter/incoming/sender.py

@@ -0,0 +1,131 @@
+import os
+import sys
+import requests
+import traceback
+import time
+import io
+import binascii
+import logging
+import matplotlib.pyplot as plt
+import numpy as np
+import datetime
+import re
+
+rootDirectory = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+sys.path.append(rootDirectory)
+from library.config import glob_the_configs
+from library.prom_api import PromAPI
+
+HOSTNAME = None
+
+
+class SlackSender:
+   def __init__(self, token, log):
+      self.token = token
+      self.alertList = []
+      self.log = log
+      try:
+         self.alertList = glob_the_configs(rootDirectory, \
+            rootDirectory + "/AoM_Configs/alert_routing_lookup", \
+            'http://consul.service.consul:8500', '127.0.0.1', log)
+      except Exception:
+         log.error("Failed to load config files: {}".format(traceback.format_exc()))
+
+   def respond(self, channel, alertId, interval, step):
+      self.log.debug("incomming message: channel: {} alert ID: {} interval: {} step: {}".format(channel, alertId, interval, step))
+      matchingAlert = next((alert for alert in self.alertList if alert['id'] == alertId), None)
+      if not matchingAlert is None:
+         query_args = {
+            'interval' : matchingAlert['interval'],
+            'start_time' : matchingAlert['start_time'],
+            'end_time' : matchingAlert['end_time'],
+            'query' : matchingAlert['query'],
+         }
+         prom_api = PromAPI(endpoint=matchingAlert['prometheus_url'])
+         if interval :
+            try:
+               dur = parse_go_duration(interval)
+               if dur < 60 :
+                  dur = 60
+               end = 0
+               start = -1 * dur
+               query_args['start_time'] = start
+               query_args['end_time'] = end
+            except Exception:
+               pass
+         if step :
+            query_args['interval'] = step
+         self.log.debug("QUERY_ARGS {} FROM {} {}".format(query_args, interval, step))
+         ret = prom_api.query_range(
+            query=query_args['query'],
+            start=query_args['start_time'],
+            end=query_args['end_time'],
+            duration=query_args['interval'])
+         if 'status' in ret and ret['status'] == 'success' and 'data' in ret and 'result' in ret['data'] and len(ret['data']['result']) > 0 and 'values' in ret['data']['result'][0] and ret['data']['result'][0]['values'] is not None and len(ret['data']['result'][0]['values']) > 0:
+            resultsForGraph = {}
+            for row in ret['data']['result'][0]['values']:
+               resultsForGraph[row[0]] = row[1]
+            finalResults = {}
+            for res in resultsForGraph:
+               finalResults[time.strftime('%H:%M:%S', time.localtime(res))] = float(resultsForGraph[res])
+            plt.clf()
+            plt.plot(list(finalResults.keys()), list(finalResults.values()))
+            plt.suptitle(alertId + " (all times UTC)")
+            if len(finalResults.keys()) > 5:
+                tickTuples = [(index, x) for index, x in enumerate(finalResults.keys()) if index % int(len(finalResults.keys()) / 5) == 0]
+                tickList = []
+                for pair in tickTuples:
+                    tickList.append(pair[1])
+                plt.xticks(ticks = tickList, rotation='vertical')
+            else:
+                plt.xticks(rotation='vertical')
+            plt.ylim(bottom = 0)
+            plt.subplots_adjust(bottom=0.2)
+            pngData = io.BytesIO()
+            fig = plt.gcf()
+            fig.savefig(pngData, format = 'png')
+            self.sendGraph(channel, pngData)
+         else:
+            self.log.debug("didn't meet criteria")
+            self.sendQueryResults(channel, ret)
+      else:
+         self.sendQueryResults(channel, "Sorry, I couldn't find a matching alert with ID {}".format(alertId))
+
+   def sendQueryResults(self, channelId, queryResults):
+      response = requests.post('https://slack.com/api/chat.postMessage',
+         headers = {
+            'Authorization': "Bearer " + self.token,
+            'Content-Type': 'application/json; charset=utf-8'
+         },
+         json = { 'text': queryResults, 'channel': channelId }
+      )
+      self.log.debug("slack response: {}".format(response.text))
+
+   def sendGraph(self, channelId, rawData):
+      request = requests.Request('POST', 'https://slack.com/api/files.upload',
+         data = { 'token': self.token, 'filetype': 'png', 'channels': channelId },
+         files = { 'file': ('graph.png', rawData.getvalue(), 'image/png')}
+      ).prepare()
+      self.log.debug("headers to send to Slack: {}".format('\r\n'.join('{}: {}'.format(k, v) for k, v in request.headers.items())))
+      self.log.debug("body to send to Slack: {}".format(len(request.body)))
+      response = requests.Session().send(request)
+      self.log.debug("slack response: {}".format(response.text))
+
+   def setAlertList(self, newAlertList):
+      # TODO can rework to be a dictionary for faster lookup if necessary
+      self.alertList = newAlertList
+
+def parse_go_duration(duration) :
+   duration = str(duration)
+   e = Exception("invalid duration "+duration)
+   if not re.match("^[0-9][0-9]*[a-z]$", duration) :
+      raise e
+   unit = duration[-1:]
+   n = int(duration[:-1])
+   if unit == "s" :
+      return n
+   if unit == "m" :
+      return 60*n
+   if unit == "h" :
+      return 60*60*n
+   raise e