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 Porter: DEFAULTS_PATH = "default_settings/" def __init__(self, user="", password="", host="http://rocketchat:3000", destination="settings.json"): self.username = os.environ.get('API_USER', user) self.password = os.environ.get('API_PASS', password) self.host = os.environ.get('API_HOST', host) self.destination = os.environ.get('SETTINGS_PATH', destination) self.session = Session() self._wait_for_rocketchat() def _wait_for_rocketchat(self, rocket=None): self.rocket = rocket timeout = 60 logger.info("starting to wait for rocketchat...") while timeout: time.sleep(1) try: if self.rocket is None: self.rocket = RocketChat(self.username, self.password, server_url=self.host, session=self.session) ret = self.rocket.info().json() if ret.get('success'): logger.info("rocketchat available!") return except ConnectionError: pass except: pass timeout =- 1 logger.error("Timeout while waiting for rocketchat!") def import_settings(self): with open(self.destination, 'r') as f: settings_to_import = json.load(f) logger.info(f"Importing {len(settings_to_import)} settings...") self._import(settings_to_import) logger.info("Finished importing settings!") def _import(self, settings): import_oauth = {} for _id, value in settings.items(): if _id.startswith("Accounts_OAuth_Custom-"): custom_name = _id.split('-') if len(custom_name) == 2: import_oauth[custom_name[1]] = True continue self.rocket.settings_update(_id, value) if import_oauth: # create all neccessary OAuth services for i, oauth_service_name in enumerate(import_oauth.keys()): logger.debug(f'create {oauth_service_name}') json_ = json.dumps({'msg': 'method', 'method': 'addOAuthService', 'params': [oauth_service_name], 'id': 1337+i}) self.rocket.call_api_post("method.call/addOAuthService", message=json_) for _id, value in settings.items(): if _id.startswith("Accounts_OAuth_Custom-"): self.rocket.settings_update(_id, value) def _export(self): data = self.rocket.settings().json() settings = {sett_.get('_id'): sett_.get('value') for sett_ in data.get('settings', [])} if not settings: logger.error('Nothing returned from server, just "{}"!'.format(data)) return {} api_default_count = settings.get('API_Default_Count', 50) # Get remaining for call in range(1, int(data.get('total', 1) / api_default_count)+1): data = self.rocket.settings(count=api_default_count, offset=call*api_default_count).json() for sett_ in data.get('settings', []): settings[sett_.get('_id')] = sett_.get('value') return settings def export_changed(self): logger.info("Exporting all non-default settings...") version = self.rocket.info().json().get('info', {}).get('version', 0) if not version: logger.error('Could not get server version!') return try: with open(os.path.join(self.DEFAULTS_PATH, '{}.json'.format(version)), 'r') as f: settings_defaults = json.load(f) except FileNotFoundError: logger.error('No defaults for version {} found in the directory {}!'.format(version, self.DEFAULTS_PATH)) return settings = self._export() settings_updated = {} for setting_id, setting_value in settings.items(): if setting_id not in settings_defaults or setting_value != settings_defaults.get(setting_id): settings_updated[setting_id] = setting_value with open(self.destination, 'w') as f: json.dump(settings_updated, f) logger.info(f"Finished exporting {len(settings_updated)} updated settings") def export_all(self): logger.info("Exporting all settings...") settings = self._export() with open(self.destination, 'w') as f: json.dump(settings, f) logger.info(f"Finished exporting {len(settings)} settings") def close(self): self.session.close()