From ce3ff13122258afbadaa3ff4daa2c3ea1094fdbb Mon Sep 17 00:00:00 2001 From: Eryn Wells Date: Mon, 5 Sep 2016 19:14:28 -0400 Subject: [PATCH] Pin manager, part 1 --- plugins/cookie.py | 138 +++++++++++++++++++++++++++++++++++++++++++--- rtmbot.py | 6 +- service.py | 61 ++++++++++++++++++++ 3 files changed, 195 insertions(+), 10 deletions(-) create mode 100644 service.py diff --git a/plugins/cookie.py b/plugins/cookie.py index a169d4d..992540f 100644 --- a/plugins/cookie.py +++ b/plugins/cookie.py @@ -2,26 +2,146 @@ # Pin management # Eryn Wells +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 diff --git a/rtmbot.py b/rtmbot.py index 77c7fa3..e946293 100755 --- a/rtmbot.py +++ b/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() diff --git a/service.py b/service.py new file mode 100644 index 0000000..636bf53 --- /dev/null +++ b/service.py @@ -0,0 +1,61 @@ +# service.py +# Eryn Wells + +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