diff --git a/securitybot/chat/slack.py b/securitybot/chat/slack.py index 6d9510c..1b9eaef 100644 --- a/securitybot/chat/slack.py +++ b/securitybot/chat/slack.py @@ -7,12 +7,16 @@ import logging from slackclient import SlackClient import json +import time from securitybot.user import User from securitybot.chat.chat import Chat, ChatException from typing import Any, Dict, List +RATE_LIMIT_SLEEP = 10 # sleep 10 seconds when rate limit +RATE_LIMIT_TRIES = 6 # maximum 6 tries (60s) when rate limit reached + class Slack(Chat): ''' A wrapper around the Slack API designed for Securitybot. @@ -37,7 +41,7 @@ def _validate(self): raise ChatException('Unable to connect to Slack API.') logging.info('Connection to Slack API successful!') - def _api_call(self, method, **kwargs): + def _api_call(self, method, rate_limit_retry=False, **kwargs): # type: (str, **Any) -> Dict[str, Any] ''' Performs a _validated_ Slack API call. After performing a normal API @@ -50,7 +54,19 @@ def _api_call(self, method, **kwargs): Returns: (dict): Parsed JSON from the response. ''' - response = self._slack.api_call(method, **kwargs) + cur_try = 0 + while True: + response = self._slack.api_call(method, **kwargs) + if cur_try > RATE_LIMIT_TRIES: + raise ChatException('Slack rate limit max tries reached.') + + if response.get('error', '').lower() == 'ratelimited' and rate_limit_retry: + logging.debug('Rate limiting reached. Sleeping {}.'.format(RATE_LIMIT_SLEEP)) + cur_try += 1 + time.sleep(RATE_LIMIT_SLEEP) + else: + break + if not ('ok' in response and response['ok']): if kwargs: logging.error('Bad Slack API request on {} with {}'.format(method, kwargs)) @@ -84,7 +100,20 @@ def get_users(self): } } ''' - return self._api_call('users.list')['members'] + members = [] + next_cursor = None + while True: + response = self._api_call('users.list', rate_limit_retry=True, cursor=next_cursor) + active_members = [m for m in response['members'] if not m.get('deleted')] + members.extend(active_members) + logging.debug('Fetched {} members'.format(len(members))) + + metadata = response.get('response_metadata') + next_cursor = metadata.get('next_cursor') if metadata else None + if not metadata or not metadata.get('next_cursor'): + break + + return members def get_messages(self): # type () -> List[Dict[str, Any]]