From f303c8a6ba964d9aae541a18e950e03aa8d8e2e1 Mon Sep 17 00:00:00 2001
From: jweigele <jweigele@local>
Date: Fri, 18 Jun 2021 19:00:18 -0700
Subject: [PATCH] Horn logging with aioprometheus

---
 grahbot.py | 49 ++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 3 deletions(-)

diff --git a/grahbot.py b/grahbot.py
index 7ec8221..693251a 100755
--- a/grahbot.py
+++ b/grahbot.py
@@ -12,10 +12,11 @@ import json
 import traceback
 import os
 import logging
+from aioprometheus import Counter, Service, Gauge
 log = logging.getLogger(__name__)
 log.setLevel(logging.DEBUG)
 ch = logging.StreamHandler()
-ch.setFormatter(logging.Formatter('|%(levelname)s|%(asctime)s|%(module)s| %(message)s'))
+ch.setFormatter(logging.Formatter('|%(levelname)s|%(asctime)s|%(module)s|%(lineno)d %(message)s'))
 log.addHandler(ch)
 
 from collections import defaultdict, OrderedDict
@@ -131,12 +132,41 @@ class HornClient(object):
             #exchange.publish(msg, routing_key=self.rabbit_config['exchange'])
 
 
+class HornProm(object):
+    def __init__(self):
+        self.horn_counter = Counter(
+            "grahbot_horn_count", "Number of horns submitted."
+        )
+        self.horn_voice = Gauge(
+            "grahbot_horn_voice_channels", "Channels actively joined by Grahbot"
+        )
+
+    async def inc_horn(self, data):
+        log.debug('Prom: Incrementing horn counter with data {}'.format(data))
+        self.horn_counter.inc(labels=data)
+
+    async def prom_voice_channel(self, channel, value):
+        log.debug('Prom: Setting voice channel {} to {}'.format(channel, value))
+        self.horn_voice.set(labels={'guild': channel.guild.name, 'channel': channel.name}, value=value)
+
+    async def prom_connect(self):
+        log.debug('Entered the prom_connect method')
+        svr = Service()
+        log.debug('Instantiated service')
+        svr.register(self.horn_counter)
+        svr.register(self.horn_voice)
+        
+        await svr.start(addr="0.0.0.0", port=42069)
+        log.debug('Started prom server?')
+        while (True):
+            await asyncio.sleep(1)
+
 class GrahDiscordException(Exception):
     pass
 
 
 
-class GrahDiscordBot(discord.Client, HornClient):
+class GrahDiscordBot(discord.Client, HornClient, HornProm):
     AUDIO_FILES = ['mp3', '.ogg']
     CMD_MARKER = '!'
     CMD_PREFIX = 'command_'
@@ -164,6 +194,10 @@ class GrahDiscordBot(discord.Client, HornClient):
         #self.horn_client = HornClient('config-piege.json')
         HornClient.__init__(self, config)
         self.rabbit = self.loop.create_task(self.rabbit_connect())
+
+        HornProm.__init__(self)
+        self.horn_prom = self.loop.create_task(self.prom_connect())
+        
         discord.Client.__init__(self, loop=self.loop, intents=intents)
         self.loop_forever()
 
@@ -192,6 +226,7 @@ class GrahDiscordBot(discord.Client, HornClient):
     async def horn(self, obj, properties):
         try:
             log.info('Horn! {}'.format(obj))
+
             # exclusive access on each horn or not?
             if 'exclusive' in obj:
                 exclusive = obj['exclusive']
@@ -205,6 +240,8 @@ class GrahDiscordBot(discord.Client, HornClient):
                 return None
             else:
                 sample_name = '{}/{}'.format(self.config['airhorn_directory'], obj['sample_name'])
+            await self.inc_horn({'sample_name': obj['sample_name'], 'guild': discord.utils.get(self.guilds, id=obj['guild']).name, 'source': obj['source'], 'remote_ip': obj['remote_ip']})
+            
         except:
             traceback.print_exc()
             log.debug('Error object was: {}'.format(obj))
@@ -407,7 +444,9 @@ class GrahDiscordBot(discord.Client, HornClient):
         await self.msg_join_voice_channel(' '.join(message.content.split(' ')[1:]), guild=guild)
 
     async def command_play(self, message, guild):
-        log.info('Was instructed to play {}'.format(' '.join(message.content.split(' ')[1:])))
+        horn = ' '.join(message.content.split(' ')[1:])
+        log.info('Was instructed to play {}'.format(horn))
+        await self.inc_horn({'source': 'discord', 'guild': guild.name, 'sample_name': horn})
         await self.msg_play_sound(name=' '.join(message.content.split(' ')[1:]), guild=guild)
 
     async def command_list(self, message, guild):
@@ -480,6 +519,10 @@ class GrahDiscordBot(discord.Client, HornClient):
     async def on_voice_state_update(self, member, before, after):
         if self.user.id == member.id:
             log.debug('Something happened to me in voice! Member {} before {} after {}'.format(member, before, after))
+            if before.channel is not None:
+                await self.prom_voice_channel(before.channel, 0)
+            if after.channel is not None:
+                await self.prom_voice_channel(after.channel, 1)
             if before.channel is not None and after.channel is not None and before.channel != after.channel:
                 log.debug('Before != after and channels differ, setting the VC appropriately')
                 # try to get the already established VC and resave
-- 
2.30.2