room_state = pickle.load(sf) sf.close() logger.debug('\n' + 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) else: logger.debug('No state file at ' + state_file + ', so state will be created') room_state = {'_meta': {'schema_version': VERSION}} logger.debug('Initialize rocket.chat API connection') rocket = RocketChat(rc_user, rc_pass, server_url=rc_server) logger.debug('LOAD / UPDATE room state') assemble_state(room_state, rocket.channels_list_joined().json(), 'channels') assemble_state(room_state, rocket.im_list().json(), 'ims') assemble_state(room_state, rocket.groups_list().json(), 'groups') 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: ' + str(start_time)) logger.debug('Global end time: ' + str(end_time)) logger.debug('Room start ts: ' + str(channel_data['begintime'])) logger.debug('Last message: ' + str(channel_data['lastmessage'])) logger.debug('Last saved: ' + str(channel_data['lastsaved']))
class TestChannels(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.channels_add_owner('GENERAL', 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') def tearDown(self): self.rocket.users_delete(self.testuser_id) def test_channels_list(self): channels_list = self.rocket.channels_list().json() self.assertTrue(channels_list.get('success')) self.assertIn('channels', channels_list) def test_channels_list_joined(self): channels_list_joined = self.rocket.channels_list_joined().json() self.assertTrue(channels_list_joined.get('success')) self.assertIn('channels', channels_list_joined) def test_channels_info(self): channels_info = self.rocket.channels_info(room_id='GENERAL').json() self.assertTrue(channels_info.get('success')) self.assertIn('channel', channels_info) self.assertEqual(channels_info.get('channel').get('_id'), 'GENERAL') channel_name = channels_info.get('channel').get('name') channels_info = self.rocket.channels_info(channel=channel_name).json() self.assertTrue(channels_info.get('success')) self.assertIn('channel', channels_info) self.assertEqual(channels_info.get('channel').get('_id'), 'GENERAL') self.assertEqual( channels_info.get('channel').get('name'), channel_name) with self.assertRaises(RocketMissingParamException): self.rocket.channels_info() def test_channels_history(self): channels_history = self.rocket.channels_history( room_id='GENERAL').json() self.assertTrue(channels_history.get('success')) self.assertIn('messages', channels_history) def test_channels_add_all(self): channels_add_all = self.rocket.channels_add_all('GENERAL').json() self.assertTrue(channels_add_all.get('success')) def test_channels_add_and_remove_moderator(self): me = self.rocket.me().json() channels_add_moderator = self.rocket.channels_add_moderator( 'GENERAL', me.get('_id')).json() self.assertTrue(channels_add_moderator.get('success')) channels_remove_moderator = self.rocket.channels_remove_moderator( 'GENERAL', me.get('_id')).json() self.assertTrue(channels_remove_moderator.get('success')) def test_channels_add_and_remove_owner(self): channels_add_owner = self.rocket.channels_add_owner( 'GENERAL', user_id=self.testuser_id).json() self.assertTrue(channels_add_owner.get('success'), channels_add_owner.get('error')) channels_remove_owner = self.rocket.channels_remove_owner( 'GENERAL', user_id=self.testuser_id).json() self.assertTrue(channels_remove_owner.get('success'), channels_remove_owner.get('error')) with self.assertRaises(RocketMissingParamException): self.rocket.channels_add_owner(room_id='GENERAL') def test_channels_archive_unarchive(self): channels_archive = self.rocket.channels_archive('GENERAL').json() self.assertTrue(channels_archive.get('success')) channels_unarchive = self.rocket.channels_unarchive('GENERAL').json() self.assertTrue(channels_unarchive.get('success')) def test_channels_close_open(self): channels_close = self.rocket.channels_close('GENERAL').json() self.assertTrue(channels_close.get('success')) channels_open = self.rocket.channels_open('GENERAL').json() self.assertTrue(channels_open.get('success')) def test_channels_create_delete(self): name = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() self.assertTrue(channels_create.get('success')) self.assertEqual(name, channels_create.get('channel').get('name')) channels_delete = self.rocket.channels_delete(channel=name).json() self.assertTrue(channels_delete.get('success')) channels_create = self.rocket.channels_create(name).json() self.assertTrue(channels_create.get('success')) room_id = channels_create.get('channel').get('_id') channels_delete = self.rocket.channels_delete(room_id=room_id).json() self.assertTrue(channels_delete.get('success')) with self.assertRaises(RocketMissingParamException): self.rocket.channels_delete() def test_channels_get_integrations(self): channels_get_integrations = self.rocket.channels_get_integrations( room_id='GENERAL').json() self.assertTrue(channels_get_integrations.get('success')) def test_channels_invite(self): channels_invite = self.rocket.channels_invite('GENERAL', self.testuser_id).json() self.assertTrue(channels_invite.get('success')) def test_channels_kick(self): channels_kick = self.rocket.channels_kick('GENERAL', self.testuser_id).json() self.assertTrue(channels_kick.get('success')) def test_channels_leave(self): channels_leave = self.rocket.channels_leave('GENERAL').json() self.assertFalse(channels_leave.get('success')) self.assertEqual(channels_leave.get('errorType'), 'error-you-are-last-owner') name = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() self.rocket.channels_invite( room_id=channels_create.get('channel').get('_id'), user_id=self.testuser_id) self.rocket.channels_add_owner( channels_create.get('channel').get('_id'), user_id=self.testuser_id).json() channels_leave = self.rocket.channels_leave( channels_create.get('channel').get('_id')).json() self.assertTrue(channels_leave.get('success')) def test_channels_rename(self): name = str(uuid.uuid1()) name2 = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() channels_rename = self.rocket.channels_rename( room_id=channels_create.get('channel').get('_id'), name=name2).json() self.assertTrue(channels_rename.get('success')) self.assertEqual(channels_rename.get('channel').get('name'), name2) def test_channels_set_description(self): description = str(uuid.uuid1()) channels_set_description = self.rocket.channels_set_description( 'GENERAL', description).json() self.assertTrue(channels_set_description.get('success')) self.assertEqual(channels_set_description.get('description'), description, 'Description does not match') def test_channels_set_join_code(self): join_code = str(uuid.uuid1()) channels_set_join_code = self.rocket.channels_set_join_code( 'GENERAL', join_code).json() self.assertTrue(channels_set_join_code.get('success')) def test_channels_set_read_only(self): channels_set_read_only = self.rocket.channels_set_read_only( 'GENERAL', True).json() self.assertTrue(channels_set_read_only.get('success')) channels_set_read_only = self.rocket.channels_set_read_only( 'GENERAL', False).json() self.assertTrue(channels_set_read_only.get('success')) def test_channels_set_topic(self): topic = str(uuid.uuid1()) channels_set_topic = self.rocket.channels_set_topic('GENERAL', topic).json() self.assertTrue(channels_set_topic.get('success')) self.assertEqual(channels_set_topic.get('topic'), topic, 'Topic does not match') def test_channels_set_type(self): name = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() self.assertTrue(channels_create.get('success')) channels_set_type = self.rocket.channels_set_type( channels_create.get('channel').get('_id'), 'p').json() self.assertTrue(channels_set_type.get('success')) self.assertTrue(channels_set_type.get('channel').get('t'), 'p') channels_set_type = self.rocket.channels_set_type( channels_create.get('channel').get('_id'), 'c').json() # should fail because this is no more a channel self.assertFalse(channels_set_type.get('success')) def test_channels_set_announcement(self): announcement = str(uuid.uuid1()) channels_set_announcement = self.rocket.channels_set_announcement( 'GENERAL', announcement).json() self.assertTrue(channels_set_announcement.get('success')) self.assertEqual(channels_set_announcement.get('announcement'), announcement, 'Topic does not match') def test_channels_set_custom_fields(self): cf = {'key': 'value'} channels_set_custom_fields = self.rocket.channels_set_custom_fields( 'GENERAL', cf).json() self.assertTrue(channels_set_custom_fields.get('success')) self.assertEqual(cf, channels_set_custom_fields['channel']['customFields']) def test_channels_members(self): channels_members = self.rocket.channels_members( room_id='GENERAL').json() self.assertTrue(channels_members.get('success')) channels_members = self.rocket.channels_members( channel='general').json() self.assertTrue(channels_members.get('success')) with self.assertRaises(RocketMissingParamException): self.rocket.channels_members() def test_channels_roles(self): channels_roles = self.rocket.channels_roles(room_id='GENERAL').json() self.assertTrue(channels_roles.get('success')) self.assertIsNotNone(channels_roles.get('roles')) channels_roles = self.rocket.channels_roles(room_name='general').json() self.assertTrue(channels_roles.get('success')) self.assertIsNotNone(channels_roles.get('roles')) with self.assertRaises(RocketMissingParamException): self.rocket.channels_roles() def test_channels_files(self): channels_files = self.rocket.channels_files(room_id='GENERAL').json() self.assertTrue(channels_files.get('success')) channels_files = self.rocket.channels_files(room_name='general').json() self.assertTrue(channels_files.get('success')) with self.assertRaises(RocketMissingParamException): self.rocket.channels_files() def test_channels_get_all_user_mentions_by_channel(self): channels_get_all_user_mentions_by_channel = self.rocket.channels_get_all_user_mentions_by_channel( room_id='GENERAL').json() self.assertTrue( channels_get_all_user_mentions_by_channel.get('success'))
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 TestChannels(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.channels_add_owner('GENERAL', username=self.user) self.rocket = RocketChat(self.user, self.password) self.testuser = self.rocket.users_create('*****@*****.**', 'testuser1', 'password', 'testuser1').json() if not self.testuser.get('success'): user_id = self.rocket.users_info( username='******').json().get('user').get('_id') self.rocket.users_delete(user_id) self.testuser = self.rocket.users_create('*****@*****.**', 'testuser1', 'password', 'testuser1').json() def tearDown(self): self.rocket.users_delete(self.testuser.get('user').get('_id')) def test_channels_list(self): channels_list = self.rocket.channels_list().json() self.assertTrue(channels_list.get('success')) self.assertIn('channels', channels_list) def test_channels_list_joined(self): channels_list_joined = self.rocket.channels_list_joined().json() self.assertTrue(channels_list_joined.get('success')) self.assertIn('channels', channels_list_joined) def test_channels_info(self): channels_info = self.rocket.channels_info(room_id='GENERAL').json() self.assertTrue(channels_info.get('success')) self.assertIn('channel', channels_info) self.assertEqual(channels_info.get('channel').get('_id'), 'GENERAL') def test_channels_history(self): channels_history = self.rocket.channels_history( room_id='GENERAL').json() self.assertTrue(channels_history.get('success')) self.assertIn('messages', channels_history) def test_channels_add_all(self): channels_add_all = self.rocket.channels_add_all('GENERAL').json() self.assertTrue(channels_add_all.get('success')) def test_channels_add_and_remove_moderator(self): me = self.rocket.me().json() channels_add_moderator = self.rocket.channels_add_moderator( 'GENERAL', me.get('_id')).json() self.assertTrue(channels_add_moderator.get('success')) channels_remove_moderator = self.rocket.channels_remove_moderator( 'GENERAL', me.get('_id')).json() self.assertTrue(channels_remove_moderator.get('success')) def test_channels_add_and_remove_owner(self): channels_add_owner = self.rocket.channels_add_owner( 'GENERAL', user_id=self.testuser.get('user').get('_id')).json() self.assertTrue(channels_add_owner.get('success'), channels_add_owner.get('error')) channels_remove_owner = self.rocket.channels_remove_owner( 'GENERAL', user_id=self.testuser.get('user').get('_id')).json() self.assertTrue(channels_remove_owner.get('success'), channels_remove_owner.get('error')) def test_channels_archive_unarchive(self): channels_archive = self.rocket.channels_archive('GENERAL').json() self.assertTrue(channels_archive.get('success')) channels_unarchive = self.rocket.channels_unarchive('GENERAL').json() self.assertTrue(channels_unarchive.get('success')) def test_channels_close_open(self): channels_close = self.rocket.channels_close('GENERAL').json() self.assertTrue(channels_close.get('success')) channels_open = self.rocket.channels_open('GENERAL').json() self.assertTrue(channels_open.get('success')) def test_channels_create(self): name = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() self.assertTrue(channels_create.get('success')) self.assertEqual(name, channels_create.get('channel').get('name')) def test_channels_get_integrations(self): channels_get_integrations = self.rocket.channels_get_integrations( room_id='GENERAL').json() self.assertTrue(channels_get_integrations.get('success')) def test_channels_invite(self): channels_invite = self.rocket.channels_invite( 'GENERAL', self.testuser.get('user').get('_id')).json() self.assertTrue(channels_invite.get('success')) self.assertIn( self.testuser.get('user').get('username'), channels_invite.get('channel').get('usernames')) def test_channels_kick(self): channels_kick = self.rocket.channels_kick( 'GENERAL', self.testuser.get('user').get('_id')).json() self.assertTrue(channels_kick.get('success')) self.assertNotIn( self.testuser.get('user').get('username'), channels_kick.get('channel').get('usernames')) def test_channels_leave(self): channels_leave = self.rocket.channels_leave('GENERAL').json() self.assertFalse(channels_leave.get('success')) self.assertEqual(channels_leave.get('errorType'), 'error-you-are-last-owner') name = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() self.rocket.channels_invite( room_id=channels_create.get('channel').get('_id'), user_id=self.testuser.get('user').get('_id')) self.rocket.channels_add_owner( channels_create.get('channel').get('_id'), user_id=self.testuser.get('user').get('_id')).json() channels_leave = self.rocket.channels_leave( channels_create.get('channel').get('_id')).json() self.assertTrue(channels_leave.get('success')) def test_channels_rename(self): name = str(uuid.uuid1()) name2 = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() channels_rename = self.rocket.channels_rename( room_id=channels_create.get('channel').get('_id'), name=name2).json() self.assertTrue(channels_rename.get('success')) self.assertEqual(channels_rename.get('channel').get('name'), name2) def test_channels_set_description(self): description = str(uuid.uuid1()) channels_set_description = self.rocket.channels_set_description( 'GENERAL', description).json() self.assertTrue(channels_set_description.get('success')) self.assertEqual(channels_set_description.get('description'), description, 'Description does not match') def test_channels_set_join_code(self): join_code = str(uuid.uuid1()) channels_set_join_code = self.rocket.channels_set_join_code( 'GENERAL', join_code).json() self.assertTrue(channels_set_join_code.get('success')) def test_channels_set_read_only(self): channels_set_read_only = self.rocket.channels_set_read_only( 'GENERAL', True).json() self.assertTrue(channels_set_read_only.get('success')) channels_set_read_only = self.rocket.channels_set_read_only( 'GENERAL', False).json() self.assertTrue(channels_set_read_only.get('success')) def test_channels_set_topic(self): topic = str(uuid.uuid1()) channels_set_topic = self.rocket.channels_set_topic('GENERAL', topic).json() self.assertTrue(channels_set_topic.get('success')) self.assertEqual(channels_set_topic.get('topic'), topic, 'Topic does not match') def test_channels_set_type(self): name = str(uuid.uuid1()) channels_create = self.rocket.channels_create(name).json() self.assertTrue(channels_create.get('success')) channels_set_type = self.rocket.channels_set_type( channels_create.get('channel').get('_id'), 'p').json() self.assertTrue(channels_set_type.get('success')) self.assertTrue(channels_set_type.get('channel').get('t'), 'p') channels_set_type = self.rocket.channels_set_type( channels_create.get('channel').get('_id'), 'c').json() self.assertFalse(channels_set_type.get( 'success')) # should fail because this is no more a channel
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)
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)