logger.info('Grabbing messages since ' + str(t_oldest) + ' through ' + str(end_time)) else: logger.info('Nothing to grab between ' + str(t_oldest) + ' through ' + str(end_time)) while (t_oldest < end_time) and (t_oldest < channel_data['lastmessage']): logger.info('') t_latest = t_oldest + one_day - datetime.timedelta(microseconds=1) logger.info('start: ' + get_rocketchat_timestamp(t_oldest)) history_data_obj = {} if (channel_data['type'] == 'channels'): history_data_obj = rocket.channels_history(channel_id, count=count_max, include='true', latest=get_rocketchat_timestamp(t_latest), oldest=get_rocketchat_timestamp(t_oldest)) elif (channel_data['type'] == 'ims'): history_data_obj = rocket.im_history(channel_id, count=count_max, include='true', latest=get_rocketchat_timestamp(t_latest), oldest=get_rocketchat_timestamp(t_oldest)) elif (channel_data['type'] == 'groups'): history_data_obj = rocket.groups_history(channel_id, count=count_max, include='true', latest=get_rocketchat_timestamp(t_latest), oldest=get_rocketchat_timestamp(t_oldest)) history_data = history_data_obj.json() history_data_text = history_data_obj.text num_messages = len(history_data['messages']) logger.info('Messages found: ' + str(num_messages)) if num_messages > 0: with open(output_dir + t_oldest.strftime('%Y-%m-%d') + '-' + channel_data['name'] + '.json', 'wb') as f: f.write(history_data_text.encode('utf-8').strip()) sleep(polite_pause) elif num_messages > count_max: logger.error('Too many messages for this room today. SKIPPING.') logger.info('end: ' + get_rocketchat_timestamp(t_latest))
class TestGroups(unittest.TestCase): def setUp(self): self.rocket = RocketChat() self.user = '******' self.password = '******' self.email = '*****@*****.**' self.rocket.users_register(email=self.email, name=self.user, password=self.password, username=self.user) self.rocket = RocketChat(self.user, self.password) testuser = self.rocket.users_info(username='******').json() if not testuser.get('success'): testuser = self.rocket.users_create('*****@*****.**', 'testuser1', 'password', 'testuser1').json() if not testuser.get('success'): self.fail("can't create test user") self.testuser_id = testuser.get('user').get('_id') self.test_group_name = str(uuid.uuid1()) self.test_group_id = self.rocket.groups_create( self.test_group_name).json().get('group').get('_id') def tearDown(self): self.rocket.users_delete(self.testuser_id) def test_groups_list_all(self): groups_list = self.rocket.groups_list_all().json() self.assertTrue(groups_list.get('success')) self.assertIn('groups', groups_list) def test_groups_list(self): groups_list = self.rocket.groups_list().json() self.assertTrue(groups_list.get('success')) self.assertIn('groups', groups_list) def test_groups_info(self): groups_info_by_id = self.rocket.groups_info( room_id=self.test_group_id).json() self.assertTrue(groups_info_by_id.get('success')) self.assertIn('group', groups_info_by_id) self.assertEqual( groups_info_by_id.get('group').get('_id'), self.test_group_id) groups_info_by_name = self.rocket.groups_info( room_name=self.test_group_name).json() self.assertTrue(groups_info_by_name.get('success')) self.assertIn('group', groups_info_by_name) self.assertEqual( groups_info_by_name.get('group').get('_id'), self.test_group_id) with self.assertRaises(RocketMissingParamException): self.rocket.groups_info() def test_groups_history(self): groups_history = self.rocket.groups_history( room_id=self.test_group_id).json() self.assertTrue(groups_history.get('success')) self.assertIn('messages', groups_history) def test_groups_add_and_remove_moderator(self): me = self.rocket.me().json() groups_add_moderator = self.rocket.groups_add_moderator( self.test_group_id, me.get('_id')).json() self.assertTrue(groups_add_moderator.get('success')) groups_remove_moderator = self.rocket.groups_remove_moderator( self.test_group_id, me.get('_id')).json() self.assertTrue(groups_remove_moderator.get('success')) def test_groups_add_and_remove_owner(self): self.rocket.groups_invite(self.test_group_id, self.testuser_id) groups_add_owner = self.rocket.groups_add_owner( self.test_group_id, user_id=self.testuser_id).json() self.assertTrue(groups_add_owner.get('success'), groups_add_owner.get('error')) groups_remove_owner = self.rocket.groups_remove_owner( self.test_group_id, user_id=self.testuser_id).json() self.assertTrue(groups_remove_owner.get('success'), groups_remove_owner.get('error')) def test_groups_archive_unarchive(self): groups_archive = self.rocket.groups_archive(self.test_group_id).json() self.assertTrue(groups_archive.get('success')) groups_unarchive = self.rocket.groups_unarchive( self.test_group_id).json() self.assertTrue(groups_unarchive.get('success')) def test_groups_close_open(self): groups_close = self.rocket.groups_close(self.test_group_id).json() self.assertTrue(groups_close.get('success')) groups_open = self.rocket.groups_open(self.test_group_id).json() self.assertTrue(groups_open.get('success')) def test_groups_create_delete(self): name = str(uuid.uuid1()) groups_create = self.rocket.groups_create(name).json() self.assertTrue(groups_create.get('success')) self.assertEqual(name, groups_create.get('group').get('name')) groups_delete = self.rocket.groups_delete(group=name).json() self.assertTrue(groups_delete.get('success')) groups_create = self.rocket.groups_create(name).json() self.assertTrue(groups_create.get('success')) room_id = groups_create.get('group').get('_id') groups_delete = self.rocket.groups_delete(room_id=room_id).json() self.assertTrue(groups_delete.get('success')) with self.assertRaises(RocketMissingParamException): self.rocket.groups_delete() def test_groups_get_integrations(self): groups_get_integrations = self.rocket.groups_get_integrations( room_id=self.test_group_id).json() self.assertTrue(groups_get_integrations.get('success')) def test_groups_invite(self): groups_invite = self.rocket.groups_invite(self.test_group_id, self.testuser_id).json() self.assertTrue(groups_invite.get('success')) def test_groups_kick(self): id_group_created = self.rocket.groups_create(str( uuid.uuid1())).json().get('group').get('_id') groups_invite = self.rocket.groups_invite(id_group_created, self.testuser_id).json() self.assertTrue(groups_invite.get('success')) groups_kick = self.rocket.groups_kick(id_group_created, self.testuser_id).json() self.assertTrue(groups_kick.get('success')) def test_groups_leave(self): groups_leave = self.rocket.groups_leave(self.test_group_id).json() self.assertFalse(groups_leave.get('success')) self.assertEqual(groups_leave.get('errorType'), 'error-you-are-last-owner') name = str(uuid.uuid1()) groups_create = self.rocket.groups_create(name).json() self.rocket.groups_invite( room_id=groups_create.get('group').get('_id'), user_id=self.testuser_id) self.rocket.groups_add_owner(groups_create.get('group').get('_id'), user_id=self.testuser_id).json() groups_leave = self.rocket.groups_leave( groups_create.get('group').get('_id')).json() self.assertTrue(groups_leave.get('success')) def test_groups_rename(self): name = str(uuid.uuid1()) name2 = str(uuid.uuid1()) groups_create = self.rocket.groups_create(name).json() groups_rename = self.rocket.groups_rename( room_id=groups_create.get('group').get('_id'), name=name2).json() self.assertTrue(groups_rename.get('success')) self.assertEqual(groups_rename.get('group').get('name'), name2) def test_groups_set_description(self): description = str(uuid.uuid1()) groups_set_description = self.rocket.groups_set_description( self.test_group_id, description).json() self.assertTrue(groups_set_description.get('success')) self.assertEqual(groups_set_description.get('description'), description, 'Description does not match') def test_groups_set_read_only(self): groups_set_read_only = self.rocket.groups_set_read_only( self.test_group_id, True).json() self.assertTrue(groups_set_read_only.get('success')) groups_set_read_only = self.rocket.groups_set_read_only( self.test_group_id, False).json() self.assertTrue(groups_set_read_only.get('success')) def test_groups_set_topic(self): topic = str(uuid.uuid1()) groups_set_topic = self.rocket.groups_set_topic( self.test_group_id, topic).json() self.assertTrue(groups_set_topic.get('success')) self.assertEqual(groups_set_topic.get('topic'), topic, 'Topic does not match') def test_groups_set_type(self): name = str(uuid.uuid1()) groups_create = self.rocket.groups_create(name).json() self.assertTrue(groups_create.get('success')) groups_set_type = self.rocket.groups_set_type( groups_create.get('group').get('_id'), 'c').json() self.assertTrue(groups_set_type.get('success')) self.assertTrue(groups_set_type.get('group').get('t'), 'p') groups_set_type = self.rocket.groups_set_type( groups_create.get('group').get('_id'), 'p').json() # should fail because this is no more a group self.assertFalse(groups_set_type.get('success')) def test_groups_members(self): groups_members = self.rocket.groups_members( room_id=self.test_group_id).json() self.assertTrue(groups_members.get('success')) groups_members = self.rocket.groups_members( group=self.test_group_name).json() self.assertTrue(groups_members.get('success')) with self.assertRaises(RocketMissingParamException): self.rocket.groups_members() def test_groups_roles(self): name = str(uuid.uuid1()) groups_create = self.rocket.groups_create(name).json() self.assertTrue(groups_create.get('success')) groups_roles = self.rocket.groups_roles( room_id=groups_create.get('group').get('_id')).json() self.assertTrue(groups_roles.get('success')) self.assertIsNotNone(groups_roles.get('roles')) groups_roles = self.rocket.groups_roles(room_name=name).json() self.assertTrue(groups_roles.get('success')) self.assertIsNotNone(groups_roles.get('roles')) with self.assertRaises(RocketMissingParamException): self.rocket.groups_roles() def test_groups_files(self): name = str(uuid.uuid1()) groups_create = self.rocket.groups_create(name).json() self.assertTrue(groups_create.get('success')) groups_files = self.rocket.groups_files( room_id=groups_create.get('group').get('_id')).json() self.assertTrue(groups_files.get('success')) groups_files = self.rocket.groups_files(room_name=name).json() self.assertTrue(groups_files.get('success')) with self.assertRaises(RocketMissingParamException): self.rocket.groups_files()
class RocketChatBot(object): USE_BUTTONS = False def __init__(self, botname, passwd, server, command_character=None): self.botname = botname self.api = RocketChat(user=botname, password=passwd, server_url=server) self.commands = [([ 'echo', ], self.echo)] self.auto_answers = [] self.direct_answers = [] self.unknown_command = [ 'command not found', ] self.handle_unknown = None self.lastts = {} self.command_character = command_character self.conversations = {} self.button_variants = {} self.personal_ids: set = set() def echo(self, message: Message): self.send_message('@' + message.user.username + ' : ' + message.text, message.chat.id) def get_status(self, auser): return self.api.users_get_presence(username=auser) def send_photo(self, chat_id, file, *args, **kwargs): self.api.rooms_upload(rid=chat_id, file=file, *args, **kwargs) def send_chat_action(self, *args, **kwargs): # No such methods in Rocket.Chat pass def post_new_message(self, obj): response = self.api.call_api_post("chat.postMessage", **obj) print(response) def send_message(self, channel_id, msg, reply_markup=None): attachments, msg = self.get_attachments(reply_markup, channel_id, msg) # pprint(attachments) response = self.api.chat_post_message(channel=channel_id, text=msg, attachments=attachments) if response.status_code // 100 > 2: logging.error(response.json().get( "error", "Error in sending message: {}".format(response.text))) def get_attachments(self, reply_markup, user_id, msg): if reply_markup is None: return None, msg if self.USE_BUTTONS: return reply_markup.to_json(), msg else: msg += reply_markup.get_text() self.button_variants.update({user_id: reply_markup.get_variants()}) return None, msg def reply_to(self, message: Message, *args, **kwargs): return self.send_message(message.chat.id, *args, **kwargs) def register_next_step_handler(self, message: Message, f: Callable, *args, **kwargs): self.conversations.update({message.from_user.id: (f, args, kwargs)}) def get_user(self, username: str): return self.api.users_info(username=username) def add_dm_handler(self, command, action): self.commands.append((command, action)) def add_auto_answer(self, triggers, answers): self.auto_answers.append((triggers, answers)) def add_direct_answer(self, triggers, answers): self.direct_answers.append((triggers, answers)) def set_unknown_handler(self, action): self.handle_unknown = action def handle_command_character_message(self, message, channel_id): msg = message['msg'].lstrip(self.command_character) command = msg.split()[0].lower() arguments = " ".join(msg.split()[1:]) user = message['u']['username'] attachments = message['attachments'] pass_message = Message(message_id=message["_id"], text=msg, chat=Chat(chat_id=channel_id), user=User.from_message(message), attachments=attachments, json=message) for cmd_list in self.commands: if command.lower() in cmd_list[0]: cmd_list[1](pass_message) return if not self.handle_auto_answer(message, self.direct_answers, channel_id): self.send_message('@' + user + ' :' + choice(self.unknown_command), channel_id) def handle_direct_message(self, message, channel_id): msg = message['msg'].partition('@' + self.botname)[2].strip() if message["msg"].startswith('@' + self.botname) \ else message["msg"].strip() if len(msg) > 0: command = msg.split()[0].lower() # arguments = " ".join(msg.split()[1:]) user = User.from_message(message) attachments = message['attachments'] pass_message = Message(message_id=message["_id"], text=msg, chat=Chat(chat_id=channel_id), user=user, attachments=attachments, json=message) conversation = self.conversations.get(user.id) variants = self.button_variants.get(channel_id) pass_message.text = variants.get( pass_message.text, pass_message.text) if variants else pass_message.text if conversation is not None: # Зарегистрирован следующий шаг f, args, kwargs = conversation self.conversations.pop(user.id) f(pass_message, *args, **kwargs) else: # Следующий шаг не найден, обработка как обычно for cmd_list in self.commands: if command.lower() in cmd_list[0]: cmd_list[1](pass_message) return if not self.handle_auto_answer(message, self.direct_answers, channel_id): if self.handle_unknown is not None: self.handle_unknown(pass_message) else: self.send_message( '@' + user.username + ' :' + choice(self.unknown_command), channel_id) else: user = User.from_message(message) attachments = message['attachments'] pass_message = Message(message_id=message["_id"], text=msg, chat=Chat(chat_id=channel_id), user=user, attachments=attachments, json=message) self.handle_unknown(pass_message) def handle_auto_answer(self, message, answers, channel_id): for kind in answers: for k in kind[0]: if k in message['msg'].lower(): self.send_message( choice(kind[1]) + ' @' + message['u']['username'], channel_id) return True return False def handle_messages(self, messages, channel_id): for message in messages['messages']: if message['u']['username'] != self.botname: pprint(message) pprint(channel_id) if message['u']['username'] == 'rocket.cat': pass # continue Thread(target=self.handle_direct_message, args=(message, channel_id)).start() # if message['msg'].startswith('@' + self.botname) or channel_id in self.personal_ids: # Thread(target=self.handle_direct_message, args=(message, channel_id)).start() # elif self.command_character is not None and message['msg'].startswith(self.command_character): # Thread(target=self.handle_command_character_message, args=(message, channel_id)).start() # elif 'mentions' not in message or message.get('mentions') == []: # Thread(target=self.handle_auto_answer, args=(message, self.auto_answers, channel_id)).start() def load_ts(self, channel_id, messages): if len(messages) > 0: self.lastts[channel_id] = messages[0]['ts'] else: self.lastts[channel_id] = '' def load_channel_ts(self, channel_id): self.load_ts(channel_id, self.api.channels_history(channel_id).json()['messages']) def load_group_ts(self, channel_id): self.load_ts(channel_id, self.api.groups_history(channel_id).json()['messages']) def load_im_ts(self, channel_id): response = self.api.im_history(channel_id).json() if response.get('success'): self.load_ts(channel_id, self.api.im_history(channel_id).json()['messages']) def process_messages(self, messages, channel_id): try: if "success" in messages: if messages['success'] == False: raise RuntimeError(messages['error']) if len(messages['messages']) > 0: self.lastts[channel_id] = messages['messages'][0]['ts'] self.handle_messages(messages, channel_id) except Exception as e: pprint(e) def process_channel(self, channel_id): if channel_id not in self.lastts: self.lastts[channel_id] = '' self.process_messages( self.api.channels_history(channel_id, oldest=self.lastts[channel_id]).json(), channel_id) def process_group(self, channel_id): if channel_id not in self.lastts: self.lastts[channel_id] = '' self.process_messages( self.api.groups_history(channel_id, oldest=self.lastts[channel_id]).json(), channel_id) def process_im(self, channel_id): if channel_id not in self.lastts: self.lastts[channel_id] = '' self.process_messages( self.api.im_history(channel_id, oldest=self.lastts[channel_id]).json(), channel_id) def run(self): for channel in self.api.channels_list_joined().json().get('channels'): self.load_channel_ts(channel.get('_id')) for group in self.api.groups_list().json().get('groups'): self.load_group_ts(group.get('_id')) for im in self.api.im_list().json().get('ims'): self.load_im_ts(im.get('_id')) while 1: for channel in self.api.channels_list_joined().json().get( 'channels'): Thread(target=self.process_channel, args=(channel.get('_id'), )).start() for group in self.api.groups_list().json().get('groups'): Thread(target=self.process_group, args=(group.get('_id'), )).start() for im in self.api.im_list().json().get('ims'): if self.botname in im.get("usernames"): self.personal_ids.add(im.get('_id')) Thread(target=self.process_im, args=(im.get('_id'), )).start() sleep(1)
def main(): """Main export process""" # args argparser_main = argparse.ArgumentParser() argparser_main.add_argument('configfile', help='Location of configuration file') argparser_main.add_argument('-s', '--datestart', help='Datetime to use for global starting point ' + \ 'e.g. 2016-01-01 (implied T00:00:00.000Z)') argparser_main.add_argument('-e', '--dateend', help='Datetime to use for global ending point ' + \ 'e.g. 2016-01-01 (implied T23:59:59.999Z)') argparser_main.add_argument('-r', '--readonlystate', help='Do not create or update history state file.', action="store_true") args = argparser_main.parse_args() start_time = (datetime .datetime .strptime(args.datestart, SHORT_DATE_FORMAT) .replace(hour=0, minute=0, second=0, microsecond=0) if args.datestart else None) end_time = (datetime .datetime .strptime(args.dateend, SHORT_DATE_FORMAT) .replace(hour=23, minute=59, second=59, microsecond=999999) if args.dateend \ else YESTERDAY.replace(hour=23, minute=59, second=59, microsecond=999999)) # config config_main = configparser.ConfigParser() config_main.read(args.configfile) polite_pause = int(config_main['rc-api']['pause_seconds']) count_max = int(config_main['rc-api']['max_msg_count_per_day']) output_dir = config_main['files']['history_output_dir'] state_file = config_main['files']['history_statefile'] rc_user = config_main['rc-api']['user'] rc_pass = config_main['rc-api']['pass'] rc_server = config_main['rc-api']['server'] # logging logger = logging.getLogger('export-history') logger.setLevel(logging.DEBUG) fh = logging.FileHandler('export-history.log') fh.setLevel(logging.DEBUG) ch = logging.StreamHandler() ch.setLevel(logging.DEBUG) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') fh.setFormatter(formatter) ch.setFormatter(formatter) logger.addHandler(fh) logger.addHandler(ch) logger.propagate = False room_state = {} logger.info('BEGIN execution at %s', str(datetime.datetime.today())) logger.debug('Command line arguments: %s', pprint.pformat(args)) if args.readonlystate: logger.info('Running in readonly state mode. No state file updates.') if os.path.isfile(state_file): logger.debug('LOAD state from %s', state_file) sf = open(state_file, 'rb') room_state = pickle.load(sf) sf.close() logger.debug('\n%s', pprint.pformat(room_state)) schema_version = 1.0 if '_meta' not in room_state else room_state['_meta']['schema_version'] if schema_version < VERSION: upgrade_state_schema(room_state, schema_version, logger) else: logger.debug('No state file at %s, so state will be created', state_file) room_state = {'_meta': {'schema_version': VERSION}} logger.debug('Initialize rocket.chat API connection') rocket = RocketChat(rc_user, rc_pass, server_url=rc_server) sleep(polite_pause) logger.debug('LOAD / UPDATE room state') assemble_state(room_state, rocket.channels_list_joined().json(), 'channels') sleep(polite_pause) assemble_state(room_state, rocket.im_list().json(), 'ims') sleep(polite_pause) assemble_state(room_state, rocket.groups_list().json(), 'groups') sleep(polite_pause) for channel_id, channel_data in room_state.items(): if channel_id != '_meta': # skip state metadata which is not a channel logger.info('------------------------') logger.info('Processing room: ' + channel_id + ' - ' + channel_data['name']) logger.debug('Global start time: %s', str(start_time)) logger.debug('Global end time: %s', str(end_time)) logger.debug('Room start ts: %s', str(channel_data['begintime'])) logger.debug('Last message: %s', str(channel_data['lastmessage'])) logger.debug('Last saved: %s ', str(channel_data['lastsaved'])) if start_time is not None: # use globally specified start time but if the start time # is before the channel existed, fast-forward to its creation t_oldest = channel_data['begintime'] if channel_data['begintime'] > start_time \ else start_time elif channel_data['lastsaved'] != NULL_DATE: # no global override for start time, so use a tick after # the last saved date if it exists t_oldest = channel_data['lastsaved'] + datetime.timedelta(microseconds=1) else: # nothing specified at all so use the beginning time of the channel t_oldest = channel_data['begintime'] t_latest = NULL_DATE if (t_oldest < end_time) and (t_oldest < channel_data['lastmessage']): logger.info('Grabbing messages since ' + str(t_oldest) + ' through ' + str(end_time)) else: logger.info('Nothing to grab between ' + str(t_oldest) + ' through ' + str(end_time)) while (t_oldest < end_time) and (t_oldest < channel_data['lastmessage']): logger.info('') t_latest = t_oldest + ONE_DAY - datetime.timedelta(microseconds=1) logger.info('start: %s', get_rocketchat_timestamp(t_oldest)) history_data_obj = {} retry_flag = True retry_count = 0 while retry_flag: retry_count += 1 logger.debug('invoking API to get messages (attempt %d)', retry_count) if channel_data['type'] == 'channels': history_data_obj = rocket.channels_history( channel_id, count=count_max, include='true', latest=get_rocketchat_timestamp(t_latest), oldest=get_rocketchat_timestamp(t_oldest)) elif channel_data['type'] == 'ims': history_data_obj = rocket.im_history( channel_id, count=count_max, include='true', latest=get_rocketchat_timestamp(t_latest), oldest=get_rocketchat_timestamp(t_oldest)) elif channel_data['type'] == 'groups': history_data_obj = rocket.groups_history( channel_id, count=count_max, include='true', latest=get_rocketchat_timestamp(t_latest), oldest=get_rocketchat_timestamp(t_oldest)) history_data = history_data_obj.json() history_data_text = history_data_obj.text if not history_data['success']: error_text = history_data['error'] logger.error('Error response from API endpoint: %s', error_text) if 'error-too-many-requests' in error_text: seconds_search = re.search(r'must wait (\d+) seconds', error_text, re.IGNORECASE) if seconds_search: seconds_to_wait = int(seconds_search.group(1)) if seconds_to_wait < 300: polite_pause += seconds_to_wait \ if seconds_to_wait < polite_pause \ else polite_pause logger.error('Attempting handle API rate limit error by \ sleeping for %d and updating polite_pause \ to %d for the duration of this execution', seconds_to_wait, polite_pause) sleep(seconds_to_wait) else: raise Exception('Unresonable amount of time to wait ' + 'for API rate limit') else: raise Exception('Can not parse too-many-requests error message') else: raise Exception('Untrapped error response from history API: ' + '{error_text}' .format(error_text=error_text)) else: retry_flag = False num_messages = len(history_data['messages']) logger.info('Messages found: %s', str(num_messages)) if num_messages > 0: with open(output_dir + t_oldest.strftime('%Y-%m-%d') + '-' + channel_data['name'] + '.json', 'wb') as f: f.write(history_data_text.encode('utf-8').strip()) elif num_messages > count_max: logger.error('Too many messages for this room today. SKIPPING.') logger.info('end: %s', get_rocketchat_timestamp(t_latest)) logger.info('') t_oldest += ONE_DAY sleep(polite_pause) logger.info('------------------------\n') # I am changing what 'lastsaved' means here. It used to denote the # last time a file was actually saved to disk for this channel # but I think it is more useful if it represents the maximum time for # which the channel has been checked. this will reduce lots # of unnecessary day checks if a channel is dormant for a while and then # suddenly has a message in it. This is only helpful if the # history export script is run on a periodic basis. room_state[channel_id]['lastsaved'] = end_time if not args.readonlystate: logger.debug('UPDATE state file') logger.debug('\n%s', pprint.pformat(room_state)) sf = open(state_file, 'wb') pickle.dump(room_state, sf) sf.close() else: logger.debug('Running in readonly state mode: SKIP updating state file') logger.info('END execution at %s\n------------------------\n\n', str(datetime.datetime.today()))
class RocketChatBot(object): def __init__(self, botname, passwd, server): self.botname = botname self.api = RocketChat(user=botname, password=passwd, server_url=server) self.commands = [([ 'echo', ], self.echo)] self.auto_answers = [] self.direct_answers = [] self.unknow_command = [ 'command not found', ] self.lastts = {} def echo(self, msg, user, channel_id): self.send_message('@' + user + ' : ' + msg, channel_id) def get_status(self, auser): return self.api.users_get_presence(username=auser) def send_message(self, msg, channel_id): self.api.chat_post_message(channel=channel_id, text=msg) def add_dm_handler(self, command, action): self.commands.append((command, action)) def add_auto_answer(self, triggers, answers): self.auto_answers.append((triggers, answers)) def add_direct_answer(self, triggers, answers): self.direct_answers.append((triggers, answers)) def handle_direct_message(self, message, channel_id): msg = message['msg'].lstrip('@' + self.botname).strip() if len(msg) > 0: command = msg.split()[0].lower() arguments = " ".join(msg.split()[1:]) user = message['u']['username'] for cmd_list in self.commands: if command.lower() in cmd_list[0]: cmd_list[1](arguments, user, channel_id) return if not self.handle_auto_answer(message, self.direct_answers, channel_id): self.send_message( '@' + user + ' :' + choice(self.unknow_command), channel_id) else: self.send_message('Here I am', channel_id) def handle_auto_answer(self, message, answers, channel_id): for kind in answers: for k in kind[0]: if k in message['msg'].lower(): self.send_message( choice(kind[1]) + ' @' + message['u']['username'], channel_id) return True return False def handle_messages(self, messages, channel_id): for message in messages['messages']: if message['u']['username'] != self.botname: pprint(message) if message['u']['username'] == 'rocket.cat': continue if message['msg'].startswith('@' + self.botname): Thread(target=self.handle_direct_message, args=(message, channel_id)).start() elif 'mentions' not in message or message.get( 'mentions') == []: Thread(target=self.handle_auto_answer, args=(message, self.auto_answers, channel_id)).start() def load_ts(self, channel_id, messages): if len(messages) > 0: self.lastts[channel_id] = messages[0]['ts'] else: self.lastts[channel_id] = '' def load_channel_ts(self, channel_id): self.load_ts(channel_id, self.api.channels_history(channel_id).json()['messages']) def load_group_ts(self, channel_id): self.load_ts(channel_id, self.api.groups_history(channel_id).json()['messages']) def load_im_ts(self, channel_id): response = self.api.im_history(channel_id).json() if response.get('success'): self.load_ts(channel_id, self.api.im_history(channel_id).json()['messages']) def process_messages(self, messages, channel_id): try: if "success" in messages: if messages['success'] == False: raise RuntimeError(messages['error']) if len(messages['messages']) > 0: self.lastts[channel_id] = messages['messages'][0]['ts'] self.handle_messages(messages, channel_id) except Exception as e: pprint(e) def process_channel(self, channel_id): if channel_id not in self.lastts: self.lastts[channel_id] = '' self.process_messages( self.api.channels_history(channel_id, oldest=self.lastts[channel_id]).json(), channel_id) def process_group(self, channel_id): if channel_id not in self.lastts: self.lastts[channel_id] = '' self.process_messages( self.api.groups_history(channel_id, oldest=self.lastts[channel_id]).json(), channel_id) def process_im(self, channel_id): if channel_id not in self.lastts: self.lastts[channel_id] = '' self.process_messages( self.api.im_history(channel_id, oldest=self.lastts[channel_id]).json(), channel_id) def run(self): for channel in self.api.channels_list_joined().json().get('channels'): self.load_channel_ts(channel.get('_id')) for group in self.api.groups_list().json().get('groups'): self.load_group_ts(group.get('_id')) for im in self.api.im_list().json().get('ims'): self.load_im_ts(im.get('_id')) while 1: for channel in self.api.channels_list_joined().json().get( 'channels'): Thread(target=self.process_channel, args=(channel.get('_id'), )).start() for group in self.api.groups_list().json().get('groups'): Thread(target=self.process_group, args=(group.get('_id'), )).start() for im in self.api.im_list().json().get('ims'): Thread(target=self.process_im, args=(im.get('_id'), )).start() sleep(1)