diff --git a/db.py b/db.py index e74d2e0..f4420a7 100644 --- a/db.py +++ b/db.py @@ -1,3 +1,4 @@ +from BTrees.IIBTree import IIBTree from ZODB import DB from ZODB.Connection import Connection from ZODB.FileStorage import FileStorage @@ -13,6 +14,15 @@ def get_conn(read_only=False) -> Connection: conn = db.open() if not hasattr(conn.root, 'subscribers'): conn.root.subscribers = PersistentMapping() + # migration 1 + if not hasattr(conn.root, 'counter'): + conn.root.counter = 0 + for user in conn.root.subscribers.values(): # type: Subscriber + if not hasattr(user, 'messages_forward') or not isinstance(user.messages_forward, IIBTree): + user.messages_forward = IIBTree() + user.messages_reverse = IIBTree() + # end migrations + commit() return conn @@ -20,6 +30,8 @@ class Subscriber(Persistent): def __init__(self, user_id, name): self.uid = user_id self.name = name + self.messages_forward = IIBTree() + self.messages_reverse = IIBTree() def update_from_message(self, m: Message): self.name = Subscriber.get_name(m.chat) diff --git a/main.py b/main.py index 62aa3b7..4adaca5 100755 --- a/main.py +++ b/main.py @@ -1,17 +1,19 @@ #!/usr/bin/env python3 import logging +import traceback from html import escape from queue import Queue, Empty from time import sleep from threading import Thread +from typing import Dict import sentry_sdk from telegram.error import Unauthorized, TelegramError from telegram.ext import Updater, CommandHandler, MessageHandler, Filters, CallbackQueryHandler from telegram import Message, Update, Bot, InlineKeyboardMarkup, InlineKeyboardButton, User -from config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT +from config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT, DEBUG from db import get_conn, Subscriber, PersistentMapping, commit from send_users_list import send_users_list @@ -32,6 +34,11 @@ def _notify_access_request(bot: Bot, user: User): def welcome(bot: Bot, update: Update): + if DEBUG: + _add_user(bot, update.effective_user.id) + update.message.reply_text('Добро пожаловать (debug)') + return + if update.effective_user.id in conn.root.subscribers: update.message.reply_text('Вы уже являетесь участником ЛОНО') else: @@ -40,12 +47,16 @@ def welcome(bot: Bot, update: Update): def unsubscribe(bot: Bot, update: Update): - del conn.root.subscribers[update.message.chat_id] - commit() + user = _remove_user(update.message.chat_id) update.message.reply_text('Вы были отписаны от бота. ' 'Обратитесь к @lono_contactbot если вы хотите подписаться снова.') - user = update.message.from_user - bot.send_message(MANAGEMENT_CHAT, f'{escape(user.full_name)} отписался') + bot.send_message(MANAGEMENT_CHAT, f'{escape(user.name)} отписался') + + +def _add_user(bot, uid): + user = conn.root.subscribers[uid] = Subscriber.from_chat(bot.get_chat(uid)) + commit() + return user def add_user(bot: Bot, update: Update, groups=(), args=()): @@ -69,8 +80,7 @@ def add_user(bot: Bot, update: Update, groups=(), args=()): pass try: - user = conn.root.subscribers[uid] = Subscriber.from_chat(bot.get_chat(uid)) - commit() + user = _add_user(bot, uid) if update.callback_query: update.callback_query.message.edit_reply_markup() bot.send_message(MANAGEMENT_CHAT, f'{escape(user.name)} был добавлен', @@ -80,6 +90,13 @@ def add_user(bot: Bot, update: Update, groups=(), args=()): bot.send_message(MANAGEMENT_CHAT, str(e)) +def _remove_user(uid): + user = conn.root.subscribers[uid] + del conn.root.subscribers[uid] + commit() + return user + + def remove_user(bot: Bot, update: Update, groups=(), args=()): if update.callback_query: update.callback_query.answer() @@ -101,10 +118,8 @@ def remove_user(bot: Bot, update: Update, groups=(), args=()): pass try: - name = conn.root.subscribers[uid].name - del conn.root.subscribers[uid] - commit() - bot.send_message(MANAGEMENT_CHAT, f'{escape(name)} был удален', + user = _remove_user(uid) + bot.send_message(MANAGEMENT_CHAT, f'{escape(user.name)} был удален', parse_mode='html') if update.callback_query: update.callback_query.message.edit_reply_markup() @@ -134,66 +149,87 @@ def _sign_text(text, m: Message, limit): def _process_message(bot: Bot, m: Message): current_chat = m.chat_id - users = conn.root.subscribers # type: PersistentMapping + users = conn.root.subscribers # type: Dict[int, Subscriber] if current_chat not in users: - _notify_access_request(bot, m.from_user) - return m.reply_text('Пожалуйста, обратитесь к @lono_contactbot') + if DEBUG: + _add_user(bot, current_chat) + m.reply_text('Добро пожаловать (debug)') + else: + _notify_access_request(bot, m.from_user) + return m.reply_text('Пожалуйста, обратитесь к @lono_contactbot') text = _sign_text(m.text_html, m, MAX_MESSAGE_LENGTH) caption = _sign_text(m.caption_html, m, MAX_CAPTION_LENGTH) + reply_to_message_internal_id = None + if m.reply_to_message and m.reply_to_message.message_id in users[current_chat].messages_forward: + reply_to_message_internal_id = users[current_chat].messages_forward[m.reply_to_message.message_id] + for uid, user in users.items(): - if uid == current_chat: - continue sleep(.02) + + reply_to_message_id = None + if reply_to_message_internal_id: + reply_to_message_id = user.messages_reverse.get(reply_to_message_internal_id, None) + try: r = None if m.forward_date: r = m.forward(uid) elif hasattr(m, 'audio') and m.audio: a = m.audio - r = bot.send_audio(uid, a.file_id, a.duration, a.performer, a.title, caption, parse_mode='html') + r = bot.send_audio(uid, a.file_id, a.duration, a.performer, a.title, caption, + reply_to_message_id=reply_to_message_id, parse_mode='html') elif hasattr(m, 'document') and m.document: d = m.document - r = bot.send_document(uid, d.file_id, d.file_name, caption, parse_mode='html') + r = bot.send_document(uid, d.file_id, d.file_name, caption, reply_to_message_id=reply_to_message_id, + parse_mode='html') elif hasattr(m, 'photo') and m.photo: p = m.photo - r = bot.send_photo(uid, p[-1].file_id, caption, parse_mode='html') + r = bot.send_photo(uid, p[-1].file_id, caption, reply_to_message_id=reply_to_message_id, + parse_mode='html') elif hasattr(m, 'sticker') and m.sticker: s = m.sticker - r = bot.send_sticker(uid, s.file_id) + r = bot.send_sticker(uid, s.file_id, reply_to_message_id=reply_to_message_id) elif hasattr(m, 'video') and m.video: v = m.video - r = bot.send_video(uid, v.file_id, v.duration, caption, parse_mode='html') + r = bot.send_video(uid, v.file_id, v.duration, caption, reply_to_message_id=reply_to_message_id, + parse_mode='html') elif hasattr(m, 'voice') and m.voice: v = m.voice - r = bot.send_voice(uid, v.file_id, v.duration, caption, parse_mode='html') + r = bot.send_voice(uid, v.file_id, v.duration, caption, reply_to_message_id=reply_to_message_id, + parse_mode='html') elif hasattr(m, 'video_note') and m.video_note: vn = m.video_note - r = bot.send_video_note(uid, vn.file_id, vn.duration, vn.length) + r = bot.send_video_note(uid, vn.file_id, vn.duration, vn.length, + reply_to_message_id=reply_to_message_id) elif hasattr(m, 'contact') and m.contact: c = m.contact - r = bot.send_contact(uid, c.phone_number, c.first_name, c.last_name) + r = bot.send_contact(uid, c.phone_number, c.first_name, c.last_name, + reply_to_message_id=reply_to_message_id) elif hasattr(m, 'location') and m.location: l = m.location - r = bot.send_location(uid, l.latitude, l.longitude) + r = bot.send_location(uid, l.latitude, l.longitude, reply_to_message_id=reply_to_message_id) elif hasattr(m, 'venue') and m.venue: v = m.venue l = v.location - r = bot.send_venue(uid, l.latitude, l.longitude, v.title, v.address, v.foursquare_id) + r = bot.send_venue(uid, l.latitude, l.longitude, v.title, v.address, v.foursquare_id, + reply_to_message_id=reply_to_message_id) elif hasattr(m, 'text') and m.text: - r = bot.send_message(uid, text, 'html') + r = bot.send_message(uid, text, 'html', reply_to_message_id=reply_to_message_id) if r: user.update_from_message(r) + user.messages_forward[r.message_id] = conn.root.counter + user.messages_reverse[conn.root.counter] = r.message_id except Unauthorized: - name = conn.root.subscribers[uid].name - del conn.root.subscribers[uid] + user = _remove_user(uid) commit() - bot.send_message(MANAGEMENT_CHAT, f'{name} был удален ' + bot.send_message(MANAGEMENT_CHAT, f'{user.name} был удален ' f'из-за блокировки бота', parse_mode='html') - except TelegramError: + except Exception: + traceback.print_exc() sentry_sdk.capture_exception() - + conn.root.counter += 1 commit() @@ -208,6 +244,7 @@ def task_queue(u: Updater): except Empty: pass except: + traceback.print_exc() sentry_sdk.capture_exception()