Audit Log Support

config.example.yaml

@@ -20,3 +20,4 @@ bot:
     - rowboat.plugins.sql
     - rowboat.plugins.internal
     - rowboat.plugins.stats
+    - rowboat.plugins.audit_log

requirements.txt

@@ -7,7 +7,7 @@ contextlib2==0.5.4
 cssselect==1.0.1
 dill==0.2.6
 -e git+https://github.com/b1naryth1ef/disco.git@452574ff20cd1588d975f04a5af9ee41fdc76953#egg=disco-py
-# disco-py==0.0.11-rc.8
+# disco-py==0.0.11
 Distance==0.1.3
 emoji==0.3.9
 enum34==1.1.6

rowboat/models/guild.py

@@ -31,6 +31,7 @@ class Guild(BaseModel):
     region = TextField(null=True)
 
     last_ban_sync = DateTimeField(null=True)
+    next_audit_log_sync = DateTimeField(null=True)
 
     # Rowboat specific data
     config = BinaryJSONField(null=True)

rowboat/models/migrations/0013_add_guild_next_audit_log_sync.py

@@ -0,0 +1,7 @@
+from rowboat.models.migrations import Migrate
+from rowboat.models.guild import Guild
+
+
+@Migrate.only_if(Migrate.missing, Guild, 'next_audit_log_sync')
+def add_guild_next_audit_log_sync(m):
+    m.add_columns(Guild, Guild.next_audit_log_sync)

rowboat/models/migrations/__init__.py

@@ -143,4 +143,5 @@ def deco(func):
             cls([lambda c: True], func).run()
         return deco
 
-init_db()
+from rowboat import ENV
+init_db(ENV)

rowboat/plugins/audit_log.py

@@ -0,0 +1,86 @@
+from datetime import datetime, timedelta
+
+import gevent
+from holster.emitter import Emitter
+
+from rowboat.redis import rdb
+from rowboat.plugins import BasePlugin as Plugin
+from rowboat.models.guild import Guild
+
+LAST_ENTRY_KEY = 'guild:laleik:{}'
+
+
+class AuditLogPlugin(Plugin):
+    global_plugin = True
+
+    def load(self, ctx):
+        super(AuditLogPlugin, self).load(ctx)
+
+        self.emitter = Emitter(gevent.spawn)
+
+    @Plugin.command('poll', '<guild_id:snowflake>', level=-1, group='auditlogs')
+    def command_poll(self, event, guild_id):
+        Guild.update(
+            next_audit_log_sync=datetime.utcnow()
+        ).where(
+            (Guild.guild_id == guild_id)
+        ).execute()
+        event.msg.reply('Ok, queued that guild for audit log sync')
+
+    @Plugin.schedule(10, init=True)
+    def poll_audit_log(self):
+        to_poll = filter(bool, (
+            self.client.state.guilds.get(guild.guild_id)
+            for guild in Guild.select().where(
+                (Guild.enabled == 1) &
+                (
+                    (Guild.next_audit_log_sync < datetime.utcnow()) |
+                    (Guild.next_audit_log_sync >> None)
+                )
+            )
+        ))
+
+        for guild in to_poll:
+            entries = self.poll_guild(guild)
+            self.log.info('Polled audit logs for guild %s (%s), %s entries', guild.id, guild.name, entries)
+
+            if not entries or entries > 100:
+                next_sync = datetime.utcnow() + timedelta(seconds=120)
+            else:
+                next_sync = datetime.utcnow() + timedelta(seconds=360)
+
+            Guild.update(
+                next_audit_log_sync=next_sync
+            ).where(
+                (Guild.guild_id == guild.id)
+            ).execute()
+
+    def poll_guild(self, guild):
+        last_entry_id = int(rdb.get(LAST_ENTRY_KEY.format(guild.id)) or 0)
+
+        # If we haven't polled this guild before (or it has no entries), attempt
+        #  to cache the last entry id so we can start emitting audit log entries
+        if not last_entry_id:
+            entries = guild.get_audit_log_entries(limit=1)
+            if entries:
+                rdb.set(LAST_ENTRY_KEY.format(guild.id), entries[0].id)
+            return 0
+
+        first_entry = None
+        entries = 0
+
+        # Iterate over the paginator
+        for entry in guild.audit_log_iter():
+            # Make sure to set the first entry
+            if not first_entry:
+                first_entry = entry
+
+            # Break if we've hit the last entry
+            if entry.id == last_entry_id:
+                break
+
+            entries += 1
+            self.emitter.emit(entry.action_type, entry)
+
+        rdb.set(LAST_ENTRY_KEY.format(guild.id), first_entry.id)
+        return entries