def __init__(self): if PENDING_MIGRATION: migration.do() update_id = "0" try: f = open(LAST_UPDATE_ID_FILE) update_id = f.read().split('\n')[0] f.close() except: pass self.mapper = PostgreSQLMapper(CHATS_COLLECTION_NAME) self.last_update_id = int(update_id) self.logger = logger.get_logger(__name__)
class EhBot: def __init__(self): if PENDING_MIGRATION: migration.do() update_id = "0" try: f = open(LAST_UPDATE_ID_FILE) update_id = f.read().split('\n')[0] f.close() except: pass self.mapper = PostgreSQLMapper(CHATS_COLLECTION_NAME) self.last_update_id = int(update_id) self.logger = logger.get_logger(__name__) def start(self): if ENVIRONMENT == "heroku": self.run_webhook() else: self.run_poll() def run_webhook(self): if not WEBHOOK or not BOTTLE_PORT: raise InvalidWebhookException("Webhook is empty") request = SetWebhookRequest(url=WEBHOOK) response, content = request.do() if response.status_code != 200 or not content["ok"]: raise InvalidWebhookException("Telegram response: %s" % content) self._app = bottle.Bottle() self.map_routes() self._app.run(host=BOTTLE_HOST, port=sys.argv[1]) def map_routes(self): self._app.route("/", method="POST", callback=self.handle_push_notification) self._app.route("/", method="GET", callback=self.handle_health) def run_poll(self): # disabling webhook request = SetWebhookRequest(url="") response, content = request.do() if response.status_code != 200 or not content["ok"]: raise InvalidWebhookException("Telegram response: %s" % content) while True: try: self.poll() except Exception as e: self.logger.error(str(traceback.format_exc())) self.logger.error(e) time.sleep(POLL_PERIOD) def poll(self): request = GetUpdatesRequest(offset = self.last_update_id) response, content = request.do() if response.status_code != 200 or not content["ok"]: raise GetUpdatesException("Failed to get updates") self.process_updates(content["result"]) def handle_health(self): return "I'm fine" def handle_push_notification(self): try: content = json.load(bottle.request.body) self.logger.info("Loaded from webhook: %s" % content) self.process_update(content) except Exception as e: self.logger.error(str(traceback.format_exc())) self.logger.error(e) def get_messages(self, results): return [message_from_json(m["message"]) for m in results if "message" in m] def get_last_update_id(self, results): return results[-1]["update_id"] if results else None def process_updates(self, updates): messages = self.get_messages(updates) last_update_id = self.get_last_update_id(updates) self.process_messages(messages) self.save_last_update_id(last_update_id) def process_update(self, update_json): message = message_from_json(update_json["message"]) self.process_messages([message]) def process_messages(self, messages): for message in messages: chat_id = message.chat_id text = message.text if not text: continue if text == "/help": self.process_help(chat_id) elif text.lower().find("/tldr") == 0: self.process_tldr_query(message) elif text.lower().find("/chatid") == 0: self.process_chat_id_query(chat_id) elif text.lower().find("/tag ") == 0: self.process_tag_command(message) elif text.lower().find("/deletetag ") == 0: self.process_delete_tag_command(message) elif BOT_TAG in text.lower(): self.process_tag_mention(message) def process_help(self, chat_id): request = SendMessageRequest(chat_id, HELP) response, content = request.do() if response.status_code != 200 or not content["ok"]: raise "Failed to send /help" def process_tag_command(self, message): chat_id = message.chat_id try: tag = self.get_tag_from_message(message, lambda text: text.split("/tag")[-1].strip()) self.add_tag(chat_id, tag) except UserTagLimitException as e: self.send_warning_to_user(message.user.id, "Stop spamming") self.logger.warn(e.message) def process_tag_mention(self, message): """ Takes a message that mentions EhBot and stores it as a tag """ chat_id = message.chat_id try: tag = self.get_tag_from_message(message, self.get_tag_text_from_mention) self.add_tag(chat_id, tag) except UserTagLimitException as e: self.send_warning_to_user(message.user.id, "Stop spamming") self.logger.warn(e.message) def process_chat_id_query(self, chat_id): text = "This chat's ID: {}\nUse it to call '/tldr {}'".format(chat_id, chat_id) request = SendMessageRequest(chat_id, text) response, content = request.do() if response.status_code != 200 or not content["ok"]: raise "Failed to send /help" def process_tldr_query(self, message): # try to get it from DB user = self.mapper.get_user_by_id(message.user.id) if not user: user = message.user query_chat = message.text.split("/tldr")[-1].strip() query_chat_id = None if query_chat: # chat id specified query_chat_id = query_chat user.last_tldr = query_chat_id elif message.chat_id == BOT_CHAT_ID and user.last_tldr: # sent directly to ehbot, try to provide last tldr query_chat_id = user.last_tldr else: # no query chat or tldr, provide current chat tldr query_chat_id = message.chat_id user.last_tldr = query_chat_id self.send_tags(user.id, query_chat_id) self.mapper.save_user(user) def process_delete_tag_command(self, message): chat_id = message.chat_id command = message.text.split(" ") tag_num = command[1] if len(command) >= 3: chat_id = command[2] if not tag_num.isdigit(): self.send_warning_to_user(user_id, "Tag number is not valid") tag_num = int(tag_num) - 1 user_id = message.user.id chat = self.mapper.get_chat_by_id(chat_id) if not chat: self.send_warning_to_user(user_id, "Chat doesn't have any tags") elif tag_num < 0 or tag_num >= len(chat.tags): self.send_warning_to_user(user_id, "Tag number is out of range") elif chat.tags[tag_num].user.id != user_id: # owns the tag? self.send_warning_to_user(user_id, "This tag is not yours") else: tag = chat.tags.pop(tag_num) self.mapper.save_chat(chat) self.send_warning_to_user(user_id, "Tag '%s' deleted" % tag.text) def send_tags(self, chat_id, query_chat_id): tags_text = "" chat = self.mapper.get_chat_by_id(query_chat_id) if not chat or len(chat.tags) == 0: tags_text = "No Tags found for this chat" else: tags = ["%s. %s" % (i+1, t.pretty_print()) for i, t in enumerate(chat.tags)] tags_text = "\n".join(tags) tags_text = "Tags for chat %s:\n%s" % (query_chat_id, tags_text) request = SendMessageRequest(chat_id, tags_text) response, content = request.do() if response.status_code != 200 or not content["ok"]: raise "Failed to send /tldr" def send_warning_to_user(self, user_id, warning_text): request = SendMessageRequest(user_id, warning_text) response, content = request.do() if response.status_code != 200 or not content["ok"]: raise "Failed to send warning to user" def add_tag(self, chat_id, tag): chat = self.mapper.get_chat_by_id(chat_id) if not chat: chat = Chat(chat_id) elif len(chat.tags) >= 5: chat.tags.pop(0) chat.tags.append(tag) self.mapper.save_chat(chat) def save_last_update_id(self, last_update_id): # if we found results, increment the last_update_id, else it stays the same if not last_update_id: return self.last_update_id = last_update_id + 1 f = open(LAST_UPDATE_ID_FILE, "w") f.write(str(self.last_update_id)) f.close() def get_tag_from_message(self, message, tag_func): date = datetime.fromtimestamp(message.date) chat_id = message.chat_id user = message.user chat = self.mapper.get_chat_by_id(chat_id) if chat: for t in chat.tags: d = datetime.fromtimestamp(t.date) if t.user.id == user.id and (message.date - t.date) <= 5 * 60: username = t.user.pretty_print() raise UserTagLimitException("User '%s' is submitting too rapidly" % username) text = tag_func(message.text) return Tag(text=text, user=user, date=message.date) def get_tag_text_from_mention(self, text): return text[text.lower().find(BOT_TAG) + len(BOT_TAG):].strip()