cold

sleeper_agents_aom_engine/reporter/incoming/main.py

@@ -0,0 +1,14 @@
+import os
+import logging
+from receiver import SlackReceiver
+from sender import SlackSender
+
+if __name__ == "__main__":
+	log = logging.getLogger()
+	log.setLevel(logging.DEBUG)
+	log.addHandler(logging.StreamHandler())
+
+	slack_token = os.environ["SLACK_API_TOKEN"]
+	sender = SlackSender(slack_token, log)
+	receiver = SlackReceiver(slack_token, log, sender.respond)
+	receiver.start()

sleeper_agents_aom_engine/reporter/incoming/message.py

@@ -0,0 +1,12 @@
+class SlackMessage() :
+   def __init__(self, payload) :
+      data = payload['data']
+      for i in data :
+         try :
+            setattr(self, i, getattr(data, i))
+         except Exception :
+            setattr(self, i, data[i])
+      attrs = dir(self)
+      assert "text" in attrs, "no text in message"
+      assert "user" in attrs, "no user in message"
+      assert "channel" in attrs, "no channel in message"

sleeper_agents_aom_engine/reporter/incoming/receiver.py

@@ -0,0 +1,52 @@
+import slack
+import os
+import re
+import ssl as ssl_lib
+import certifi
+
+from message import SlackMessage
+
+class SlackReceiver() :
+   def __init__(self, token, log, callback) :
+      self.token = token
+      self.ssl_context = ssl_lib.create_default_context(cafile=certifi.where())
+      self.callback = callback
+      self.log = log
+
+   def start(self) :
+      self.rtm_client = slack.RTMClient(token=self.token, ssl=self.ssl_context)
+      @slack.RTMClient.run_on(event="message")
+      def receive(**payload) :
+         msg = self.parse(payload)
+         if msg is not None:
+            self.receive(msg)
+      print("Starting")
+      self.rtm_client.start()
+
+   def parse(self, payload) :
+      self.log.debug("slack message received: {}".format(payload))
+      if 'data' in payload and 'bot_id' in payload['data'] and payload['data']['bot_id'] == 'BNYAX72BB':
+         # it's the bot's response, ignore it
+         return None
+      if 'data' in payload and 'user' in payload['data'] and payload['data']['user'] == 'UNS0QKMMY':
+         # it's the bot uploading files, ignore it
+         return None
+      if 'data' in payload and (
+            ('text' in payload['data'] and '<@UNS0QKMMY>' not in payload['data']['text']) 
+            and ('channel' in payload['data'] and not payload['data']['channel'].startswith('DP'))
+         ):
+         # message in a channel and the bot wasn't pinged, or was not a direct message - ignore
+         self.log.debug("received message, but I wasn't pinged or DM'ed")
+         return None
+      if 'data' in payload and 'text' in payload['data'] and '<@UNS0QKMMY>' in payload['data']['text']:
+         # remove the ping text
+         payload['data']['text'] = re.sub('\\<@UNS0QKMMY\\>', '', payload['data']['text'])
+      return SlackMessage(payload)
+
+   def receive(self, msg) :
+      text = msg.text.split()
+      channel = msg.channel
+      ID = text[0] if len(text) > 0 else None
+      interval = text[1] if len(text) > 1 else None
+      step = text[2] if len(text) > 2 else None
+      self.callback(channel, ID, interval, step)

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

sleeper_agents_aom_engine/reporter/incoming/test.sh

@@ -0,0 +1,14 @@
+#! /bin/bash
+
+SLACK_API_TOKEN=${SLACK_API_TOKEN:-}
+if [ -z "$SLACK_API_TOKEN" ]; then
+   SLACK_API_TOKEN=$(security find-generic-password -a $USER -s SLACK_API_TOKEN -w 2> /dev/null)
+   if [ -z "$SLACK_API_TOKEN" ]; then
+      read -s -p "Enter SLACK_API_TOKEN" SLACK_API_TOKEN
+      echo ""
+   fi
+fi
+security delete-generic-password -a $USER -s SLACK_API_TOKEN 2> /dev/null 1>&2
+security add-generic-password -a $USER -s SLACK_API_TOKEN -w $SLACK_API_TOKEN 1>&2
+
+SLACK_BOT_TOKEN=$SLACK_API_TOKEN python3 ./main.py

sleeper_agents_aom_engine/reporter/incoming/test_message.py

@@ -0,0 +1,16 @@
+import unittest
+
+from message import SlackMessage
+
+class Test_SlackMessage(unittest.TestCase) :
+   def test_basic(self) :
+      def recv(*args, **payload) :
+         return SlackMessage(payload)
+      try :
+         recv(data={"hello":"world"})
+         self.fail("did not raise exception at bad slack message payload")
+      except Exception :
+         pass
+
+if __name__ == "__main__" :
+   unittest.main()

sleeper_agents_aom_engine/reporter/incoming/test_receiver.py

@@ -0,0 +1,14 @@
+import unittest
+
+from receiver import SlackReceiver
+from message import SlackMessage
+
+class Test_SlackReceiver(unittest.TestCase) :
+   def test_basic(self) :
+      r = SlackReceiver("token", print)
+      msg = r.parse(data={"user":"u", "text":"text", "channel":"channel"})
+      self.assertTrue(isinstance(msg, SlackMessage))
+      r.receive(msg)
+
+if __name__ == "__main__" :
+   unittest.main()

sleeper_agents_aom_engine/reporter/incoming/test_sender.py

@@ -0,0 +1,35 @@
+import unittest
+import sender
+
+class Test_Compute_Relative_Time(unittest.TestCase) :
+   def test_basic(self) :
+      class MockDateTime() :
+         t = 1000000
+         def __init__(self) :
+            self.datetime = self
+         def fromtimestamp(self, t) :
+            self.t = t
+         def timestamp(self) :
+            return self.t
+         def utcnow(self) :
+            return self
+      mock_date_time = MockDateTime()
+      was = sender.datetime
+      sender.datetime = mock_date_time
+      self.case("1s", 1)
+      self.case("5s", 5)
+      self.case("0s", 0)
+      self.case("0m", 0)
+      self.case("1m", 60)
+      self.case("9m", 9*60)
+      self.case("0h", 0)
+      self.case("1h", 60*60)
+      self.case("9h", 9*60*60)
+      sender.datetime = was
+
+   def case(self, duration, expected) :
+      seconds = sender.parse_go_duration(duration)
+      self.assertEqual(seconds, expected)
+
+if __name__ == "__main__" :
+   unittest.main()