lono/main.py

#!/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, List

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, InputMediaPhoto, \
    InputMediaVideo, InputMediaAnimation, InputMediaAudio, InputMediaDocument

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


logging.basicConfig(level=logging.WARNING)
queue = Queue()
sentry_sdk.init(dsn=SENTRY_DSN)
conn = get_conn()

MAX_MESSAGE_LENGTH = 4096
MAX_CAPTION_LENGTH = 1024


def _notify_access_request(bot: Bot, user: User):
    markup = InlineKeyboardMarkup([[InlineKeyboardButton('Добавить', callback_data=f'add {user.id}')]])
    bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={user.id}">{escape(user.full_name)}</a> запросил доступ',
                     parse_mode='html', reply_markup=markup)


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:
        update.message.reply_text('Пожалуйста, обратитесь к @lono_contactbot')
        _notify_access_request(bot, update.message.from_user)


def unsubscribe(bot: Bot, update: Update):
    user = _remove_user(update.message.chat_id)
    update.message.reply_text('Вы были отписаны от бота. '
                              'Обратитесь к @lono_contactbot если вы хотите подписаться снова.')
    bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={user.id}">{escape(user.name)}</a> отписался')


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=()):
    if update.callback_query:
        update.callback_query.answer()

    if groups:
        if update.callback_query.message.chat.id != MANAGEMENT_CHAT:
            return
        uid = groups[0]
    elif args:
        uid = args[0]
    elif update.message and update.message.reply_to_message and update.message.reply_to_message.forward_from:
        uid = update.message.reply_to_message.forward_from.id
    else:
        return bot.send_message(MANAGEMENT_CHAT, 'Укажите ID пользователя или ответьте на его сообщение')

    try:
        uid = int(uid)
    except (ValueError, TypeError):
        pass

    try:
        user = _add_user(bot, uid)
        if update.callback_query:
            update.callback_query.message.edit_reply_markup()
        bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> был добавлен',
                         parse_mode='html')
        bot.send_message(uid, 'Добро пожаловать. Снова.')
    except TelegramError as e:
        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()

    if groups:
        if update.callback_query.message.chat.id != MANAGEMENT_CHAT:
            return
        uid = groups[0]
    elif args:
        uid = args[0]
    elif update.message and update.message.reply_to_message and update.message.reply_to_message.forward_from:
        uid = update.message.reply_to_message.forward_from.id
    else:
        return bot.send_message(MANAGEMENT_CHAT, 'Укажите ID пользователя или ответьте на его сообщение')

    try:
        uid = int(uid)
    except (ValueError, TypeError):
        pass

    try:
        user = _remove_user(uid)
        bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> был удален',
                         parse_mode='html')
        if update.callback_query:
            update.callback_query.message.edit_reply_markup()
    except KeyError:
        bot.send_message(MANAGEMENT_CHAT, f'Пользователь id={uid} не был найден')


def users(bot: Bot, update: Update):
    send_users_list()


def msg(bot: Bot, update: Update):
    queue.put(update.message)


def _sign_text(text, m: Message, limit):
    if not text:
        text = ''

    sign = ''
    if text.startswith('!sign') or text.startswith('/sign'):
        text = text[5:]
        sign = f'\n\n____________\n' \
               f'by <a href="tg://user?id={m.from_user.id}">{escape(m.from_user.full_name)}</a>'

    return text[:limit - len(sign)] + sign


def _process_media_group(bot: Bot, messages: List[Message]):
    if not messages:
        return

    m = messages[0]

    current_chat = m.chat_id
    users = conn.root.subscribers  # type: Dict[int, Subscriber]
    if current_chat not in users:
        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')

    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]

    media_group = []
    for message in messages:
        caption = _sign_text(message.caption_html, message, MAX_CAPTION_LENGTH)

        if hasattr(message, 'photo') and message.photo:
            media_group.append(InputMediaPhoto(message.photo[-1].file_id, caption=caption, parse_mode='html'))
        elif hasattr(message, 'video') and message.video:
            media_group.append(InputMediaVideo(message.video.file_id, caption=caption, parse_mode='html'))
        elif hasattr(message, 'animation') and message.animation:
            media_group.append(InputMediaAnimation(message.animation.file_id, caption=caption, parse_mode='html'))
        elif hasattr(message, 'document') and message.document:
            media_group.append(InputMediaDocument(message.document.file_id, caption=caption, parse_mode='html'))
        elif hasattr(message, 'audio') and message.audio:
            media_group.append(InputMediaAudio(message.audio.file_id, caption=caption, parse_mode='html'))

    remove_uids = []

    for uid, user in users.items():
        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:
            sent_messages = bot.send_media_group(uid, media_group, reply_to_message_id=reply_to_message_id)
            if sent_messages:
                user.update_from_message(sent_messages[0])
                for r in sent_messages:
                    user.messages_forward[r.message_id] = conn.root.counter
                    user.messages_reverse[conn.root.counter] = r.message_id
        except Unauthorized:
            remove_uids.append(uid)
            bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{user.name}</a> был удален '
                                              f'из-за блокировки бота', parse_mode='html')
        except Exception:
            traceback.print_exc()
            sentry_sdk.capture_exception()
    conn.root.counter += len(messages)
    commit()
    for uid in remove_uids:
        _remove_user(uid)


def _process_message(bot: Bot, m: Message):
    current_chat = m.chat_id
    users = conn.root.subscribers  # type: Dict[int, Subscriber]
    if current_chat not in users:
        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]

    remove_uids = []

    for uid, user in users.items():
        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,
                                   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, 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, 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, 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, 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, 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,
                                        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,
                                     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, 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,
                                   reply_to_message_id=reply_to_message_id)
            elif hasattr(m, 'text') and m.text:
                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:
            remove_uids.append(uid)
            bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{user.name}</a> был удален '
                                              f'из-за блокировки бота', parse_mode='html')
        except Exception:
            traceback.print_exc()
            sentry_sdk.capture_exception()
    conn.root.counter += 1
    commit()
    for uid in remove_uids:
        _remove_user(uid)


def task_queue(u: Updater):
    while True:
        if not u.running:
            return

        try:
            m = queue.get(timeout=1)  # type: Message
            if m.media_group_id:
                media_group_id = m.media_group_id
                group_messages = [m]
                while queue.qsize() and queue.queue[0].media_group_id == media_group_id:
                    group_messages.append(queue.get(block=False))
                _process_media_group(u.bot, group_messages)
            else:
                _process_message(u.bot, m)
        except Empty:
            pass
        except:
            traceback.print_exc()
            sentry_sdk.capture_exception()


if __name__ == '__main__':
    updater = Updater(BOT_TOKEN, workers=4)

    updater.dispatcher.add_handler(CommandHandler('start', welcome, Filters.private))
    updater.dispatcher.add_handler(CommandHandler('stop', unsubscribe, Filters.private))
    updater.dispatcher.add_handler(CommandHandler('add', add_user, Filters.chat(MANAGEMENT_CHAT), pass_args=True))
    updater.dispatcher.add_handler(CallbackQueryHandler(add_user, pattern=r'^add (\d+)$', pass_groups=True))
    updater.dispatcher.add_handler(CommandHandler('remove', remove_user, Filters.chat(MANAGEMENT_CHAT), pass_args=True))
    updater.dispatcher.add_handler(CallbackQueryHandler(remove_user, pattern=r'^remove (\d+)$', pass_groups=True))
    updater.dispatcher.add_handler(CommandHandler('users', users, Filters.chat(MANAGEMENT_CHAT)))
    updater.dispatcher.add_handler(MessageHandler(Filters.private, msg))

    updater.start_polling()

    tq = Thread(target=task_queue, args=(updater,))
    tq.start()

    logging.warning('LONO has started')
    updater.idle()
    logging.warning('LONO is stopping...')
    commit()
    conn.close()