Pin manager, part 1

plugins/cookie.py

@@ -2,26 +2,146 @@
 # Pin management
 # Eryn Wells <eryn@erynwells.me>
 
+import json
 import logging
+import random
+
+import requests
+
+from service import slack
 
 LOGGER = logging.getLogger('cookie')
+MAX_PINS = 100
+CHANNELS = []
+
+outputs = []
+
+class Channel(object):
+    def __init__(self, json):
+        self.ident = json['id']
+        self.name = json['name']
+        self.pins = None
+
+    @property
+    def pretty_name(self):
+        return '#' + self.name
+
+    @property
+    def pin_file(self):
+        return 'pins.{}.json'.format(self.name)
+
+    @property
+    def saved_pins(self):
+        try:
+            with open(self.pin_file, 'r') as f:
+                obj = json.load(f)
+                return obj
+        except FileNotFoundError:
+            return None
+
+    @property
+    def oldest_pin(self):
+        if not self.pins or len(self.pins) == 0:
+            return None
+        return self.pins[-1]
+
+    def fetch_pins(self):
+        pins = slack.pins(self.ident)
+        # Newest (highest value timestamp) sorted at the beginning.
+        pins.sort(key=lambda it: it['created'], reverse=True)
+        self.pins = pins
+        return pins
+
+    def unpin_oldest_if_needed(self):
+        if not _should_unpin(self.pins):
+            return
+        oldest_pin = self.oldest_pin
+        LOGGER.info('Writing pin to {}'.format(self.pin_file))
+        self.save_pin(oldest_pin)
+        if oldest_pin['type'] == 'message':
+            LOGGER.info('Unpinning oldest message: "{}"'.format(oldest_pin['message']))
+        removed = self.remove_pin(oldest_pin)
+        return oldest_pin
+
+    def remove_pin(self, pin):
+        result = slack.remove_pin(pin)
+        return result
+
+    def save_pin(self, pin):
+        pins = self.saved_pins
+        if pins is None:
+            pins = [] 
+        filtered_pins = list(filter(lambda p: p['created'] == pin['created'], pins))
+        if len(filtered_pins) > 0:
+            LOGGER.info('Message already pinned; skipping')
+        pins.append(pin)
+        self.write_pins(pins)
+
+    def write_pins(self, pins):
+        with open(self.pin_file, 'w') as f:
+            json.dump(pins, f, indent=2)
+
+def _channels():
+    channels = slack.channels()
+    if not channels:
+        return None
+    channels = [c for c in channels if c['is_member']]
+    return channels
+
+def _should_unpin(pins):
+    return len(pins) >= MAX_PINS
+
+#
+# RTM
+#
 
 def process_hello(data):
-    # TODO: Get list of channels I'm in.
-    # TODO: Query for pins from each channel.
-    # TODO: Determine if a message should be unpinned to make room for next pin.
     LOGGER.info('Hello!')
+    channels = _channels()
+    LOGGER.info('I am in these channels: {}'.format(', '.join(['#' + c['name'] for c in channels])))
+    for c in channels:
+        ch = Channel(c)
+        CHANNELS.append(ch)
+        ch.fetch_pins()
+        LOGGER.info('  {} has {} pins.'.format(ch.pretty_name, len(ch.pins)))
+        ch.unpin_oldest_if_needed()
 
 def process_channel_joined(data):
-    # TODO: Query for pins from this channel.
-    # TODO: Determine if a message should be unpinned to make room for next pin.
     LOGGER.info('Joined #{}'.format(data['channel']['name']))
+    ch = Channel(data['channel'])
+    CHANNELS.append(ch)
+    ch.fetch_pins()
+    ch.unpin_oldest_if_needed()
 
 def process_pin_added(data):
-    # TODO: Unpin oldest message if needed.
-    # TODO: Write that pin to a file.
     LOGGER.info('Pin added')
+    ch.fetch_pins()
+    ch.unpin_oldest_if_needed()
 
 def process_message(data):
-    # TODO: !cookie
-    LOGGER.info('Received message')
+    try:
+        text = data['text'].strip()
+    except KeyError:
+        # TODO: Make this better.
+        return
+    LOGGER.info('Received message: {}'.format(text))
+
+    if text == '!lore':
+        try:
+            ch = list(filter(lambda c: c.ident == data['channel'], CHANNELS))[0]
+            random_pin = _lore(ch)
+            if random_pin:
+                outputs.append([data['channel'], random_pin])
+        except IndexError:
+            pass
+
+#
+# Private
+#
+
+def _lore(channel):
+    pins = channel.saved_pins
+    random_pin = random.choice(pins)
+    if random_pin['type'] == 'message':
+        return random_pin['message']['permalink']
+    return random_pin

rtmbot.py

@@ -4,6 +4,7 @@ from argparse import ArgumentParser
 
 import yaml
 from rtmbot import RtmBot
+import service
 
 
 def parse_args():
@@ -18,7 +19,10 @@ def parse_args():
 
 # load args with config path
 args = parse_args()
-config = yaml.load(open(args.config or 'rtmbot.conf', 'r'))
+config = None
+with open(args.config or 'rtmbot.conf', 'r') as f:
+    config = yaml.load(f)
+service.slack = service.SlackService(config.get('SLACK_TOKEN'))
 bot = RtmBot(config)
 try:
     bot.start()

service.py

@@ -0,0 +1,61 @@
+# service.py
+# Eryn Wells <eryn@erynwells.me>
+
+import requests
+
+slack = None
+
+class SlackService(object):
+    '''Handles requests at the Slack API.'''
+
+    _API_BASE = 'https://slack.com/api'
+
+    def __init__(self, token):
+        self.token = token
+
+    #
+    # Endpoints
+    #
+
+    def channels(self):
+        params = self.__params()
+        r = requests.get(self.__url('channels.list'), params=params)
+        json = self.__extract_json(r)
+        return json['channels'] if json else None
+
+    def pins(self, channel):
+        params = self.__params(channel=channel)
+        r = requests.get(self.__url('pins.list'), params=params)
+        json = self.__extract_json(r)
+        return json['items'] if json else None
+
+    def remove_pin(self, pin):
+        params = self.__params()
+        if pin['type'] == 'message':
+            params['channel'] = pin['channel']
+            params['timestamp'] = pin['message']['ts']
+        r = requests.get(self.__url('pins.remove'), params=params)
+        json = self.__extract_json(r)
+        return json is not None
+
+    #
+    # Private
+    #
+
+    def __url(self, verb):
+        return SlackService._API_BASE + '/' + verb
+
+    def __params(self, **kwargs):
+        self.__append_token_param(kwargs)
+        return kwargs
+
+    def __append_token_param(self, params):
+        params['token'] = self.token
+
+    def __extract_json(self, response, code=200):
+        if response.status_code != code:
+            return None
+        json = response.json()
+        if not json['ok']:
+            return None
+        return json