lono/main.py

418 lines
15 KiB
Python
Raw Normal View History

2019-03-03 23:23:12 +00:00
#!/usr/bin/env python3
import logging
2019-09-06 12:08:11 +00:00
import os
2019-12-08 23:22:07 +00:00
import re
2019-03-15 23:27:42 +00:00
import traceback
2019-08-19 19:05:45 +00:00
from datetime import datetime, timedelta
2020-05-19 14:43:53 +00:00
from hashlib import sha1
2018-11-20 19:51:15 +00:00
from html import escape
2019-03-03 23:23:12 +00:00
from queue import Queue, Empty
2018-11-20 19:51:15 +00:00
from time import sleep
2020-05-19 14:43:53 +00:00
from threading import Thread, Event
2019-03-16 23:45:26 +00:00
from typing import Dict, List
2018-11-20 19:51:15 +00:00
2019-03-03 23:23:12 +00:00
import sentry_sdk
2020-05-19 14:43:53 +00:00
from redis import Redis
2018-11-20 19:51:15 +00:00
from telegram.error import Unauthorized, TelegramError
2020-05-19 14:43:53 +00:00
from telegram.ext import Updater, CommandHandler, MessageHandler, Filters, CallbackQueryHandler, CallbackContext
2019-03-16 23:45:26 +00:00
from telegram import Message, Update, Bot, InlineKeyboardMarkup, InlineKeyboardButton, User, InputMediaPhoto, \
InputMediaVideo, InputMediaAnimation, InputMediaAudio, InputMediaDocument
2018-11-20 19:51:15 +00:00
2019-03-15 23:27:42 +00:00
from config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT, DEBUG
2020-05-19 14:43:53 +00:00
from db import get_conn, Subscriber, commit
2019-03-03 23:23:12 +00:00
from send_users_list import send_users_list
2018-11-20 19:51:15 +00:00
2019-03-03 23:23:12 +00:00
logging.basicConfig(level=logging.WARNING)
queue = Queue()
sentry_sdk.init(dsn=SENTRY_DSN)
conn = get_conn()
2020-05-19 14:43:53 +00:00
redis = Redis()
2018-11-20 19:51:15 +00:00
2019-03-03 23:23:12 +00:00
MAX_MESSAGE_LENGTH = 4096
MAX_CAPTION_LENGTH = 1024
2018-11-20 19:51:15 +00:00
2020-05-19 14:43:53 +00:00
def _antispam(args):
if not args:
return True
args = '|'.join(map(str, args))
digest = sha1(args.encode()).digest()
key = 'lono-' + digest.hex()
if redis.get(key):
return False
redis.set(key, '1', ex=30)
return True
2019-03-03 23:23:12 +00:00
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)
2018-11-20 19:51:15 +00:00
2020-05-19 14:43:53 +00:00
def welcome(update: Update, ctx: CallbackContext):
2019-03-15 23:27:42 +00:00
if DEBUG:
2020-05-19 14:43:53 +00:00
_add_user(ctx.bot, update.effective_user.id)
2019-03-15 23:27:42 +00:00
update.message.reply_text('Добро пожаловать (debug)')
return
2019-03-03 23:23:12 +00:00
if update.effective_user.id in conn.root.subscribers:
update.message.reply_text('Вы уже являетесь участником ЛОНО')
else:
update.message.reply_text('Пожалуйста, обратитесь к @lono_contactbot')
2020-05-19 14:43:53 +00:00
_notify_access_request(ctx.bot, update.message.from_user)
2018-11-20 19:51:15 +00:00
2020-05-19 14:43:53 +00:00
def unsubscribe(update: Update, ctx: CallbackContext):
2019-03-15 23:27:42 +00:00
user = _remove_user(update.message.chat_id)
2019-03-03 23:23:12 +00:00
update.message.reply_text('Вы были отписаны от бота. '
'Обратитесь к @lono_contactbot если вы хотите подписаться снова.')
2020-05-19 14:43:53 +00:00
ctx.bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={user.id}">{escape(user.name)}</a> отписался')
2019-03-15 23:27:42 +00:00
def _add_user(bot, uid):
user = conn.root.subscribers[uid] = Subscriber.from_chat(bot.get_chat(uid))
commit()
return user
2019-03-03 23:23:12 +00:00
2020-05-19 14:43:53 +00:00
def add_user(update: Update, ctx: CallbackContext, groups=(), args=()):
2019-03-03 23:23:12 +00:00
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:
2020-05-19 14:43:53 +00:00
return ctx.bot.send_message(MANAGEMENT_CHAT, 'Укажите ID пользователя или ответьте на его сообщение')
2018-11-20 19:51:15 +00:00
try:
2019-03-03 23:23:12 +00:00
uid = int(uid)
except (ValueError, TypeError):
2018-11-20 19:51:15 +00:00
pass
2019-03-03 23:23:12 +00:00
try:
2020-05-19 14:43:53 +00:00
user = _add_user(ctx.bot, uid)
2019-03-03 23:23:12 +00:00
if update.callback_query:
update.callback_query.message.edit_reply_markup()
2020-05-19 14:43:53 +00:00
ctx.bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> был добавлен',
2019-03-03 23:23:12 +00:00
parse_mode='html')
2020-05-19 14:43:53 +00:00
ctx.bot.send_message(uid, 'Добро пожаловать. Снова.')
2019-03-03 23:23:12 +00:00
except TelegramError as e:
2020-05-19 14:43:53 +00:00
ctx.bot.send_message(MANAGEMENT_CHAT, str(e))
2019-03-03 23:23:12 +00:00
2019-03-15 23:27:42 +00:00
def _remove_user(uid):
user = conn.root.subscribers[uid]
del conn.root.subscribers[uid]
commit()
return user
2020-05-19 14:43:53 +00:00
def remove_user(update: Update, ctx: CallbackContext, groups=(), args=()):
2019-03-03 23:23:12 +00:00
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:
2020-05-19 14:43:53 +00:00
return ctx.bot.send_message(MANAGEMENT_CHAT, 'Укажите ID пользователя или ответьте на его сообщение')
2019-03-03 23:23:12 +00:00
try:
uid = int(uid)
except (ValueError, TypeError):
pass
try:
2019-03-15 23:27:42 +00:00
user = _remove_user(uid)
2020-05-19 14:43:53 +00:00
ctx.bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> был удален',
2019-03-03 23:23:12 +00:00
parse_mode='html')
if update.callback_query:
update.callback_query.message.edit_reply_markup()
except KeyError:
2020-05-19 14:43:53 +00:00
ctx.bot.send_message(MANAGEMENT_CHAT, f'Пользователь id={uid} не был найден')
2019-03-03 23:23:12 +00:00
2020-05-19 14:43:53 +00:00
def users(update: Update, ctx: CallbackContext):
2019-03-03 23:23:12 +00:00
send_users_list()
2020-05-19 14:43:53 +00:00
def msg(update: Update, ctx: CallbackContext):
2019-03-03 23:23:12 +00:00
queue.put(update.message)
def _sign_text(text, m: Message, limit):
if not text:
text = ''
2019-12-08 23:22:07 +00:00
text = re.sub(r'<a href="(tg://.*?)">.*?</a>', '<a href="\\1">\\1</a>', text)
2019-03-15 22:27:53 +00:00
sign = ''
2019-03-03 23:23:12 +00:00
if text.startswith('!sign') or text.startswith('/sign'):
2019-03-16 08:25:27 +00:00
text = text[5:]
2019-04-04 23:52:53 +00:00
sign = f'\n\n— <a href="tg://user?id={m.from_user.id}">{escape(m.from_user.full_name)}</a>'
2019-03-15 22:27:53 +00:00
return text[:limit - len(sign)] + sign
2019-03-03 23:23:12 +00:00
2019-03-16 23:45:26 +00:00
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'))
2019-03-17 12:07:38 +00:00
remove_uids = []
2019-03-16 23:45:26 +00:00
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:
2019-03-17 12:07:38 +00:00
remove_uids.append(uid)
2019-03-16 23:45:26 +00:00
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()
2019-03-17 12:07:38 +00:00
for uid in remove_uids:
_remove_user(uid)
2019-03-16 23:45:26 +00:00
2020-05-19 14:43:53 +00:00
def users_list(update: Update, ctx: CallbackContext):
2019-04-26 18:00:51 +00:00
current_chat = update.effective_chat.id
subs = conn.root.subscribers # type: Dict[int, Subscriber]
if current_chat not in subs:
return
messages = [f'Count: {len(subs)}\n']
for sub in subs.values(): # type: Subscriber
msg = f'<code>{sub.uid:>14}</code> '
if sub.uid < 0:
msg += str(sub.name)
else:
msg += f'<a href="tg://user?id={sub.uid}">{sub.name}</a>'
messages.append(msg)
for i in range(0, len(messages), 40):
update.effective_message.reply_text('\n'.join(messages[i:i+40]), parse_mode='html')
2019-03-03 23:23:12 +00:00
def _process_message(bot: Bot, m: Message):
2019-08-19 19:05:45 +00:00
if m.sticker or m.animation:
delta = datetime.now() - conn.root.last_media
if delta < timedelta(seconds=15):
bot.send_message(
m.from_user.id,
'Не вайпи, до следующей гифки/стикера осталось {} секунд'.format(15 - int(delta.total_seconds()))
)
return
conn.root.last_media = datetime.now()
commit()
2019-03-03 23:23:12 +00:00
current_chat = m.chat_id
2019-03-15 23:27:42 +00:00
users = conn.root.subscribers # type: Dict[int, Subscriber]
2019-03-03 23:23:12 +00:00
if current_chat not in users:
2019-03-15 23:27:42 +00:00
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')
2019-03-03 23:23:12 +00:00
text = _sign_text(m.text_html, m, MAX_MESSAGE_LENGTH)
caption = _sign_text(m.caption_html, m, MAX_CAPTION_LENGTH)
2019-03-15 23:27:42 +00:00
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]
2020-05-19 14:43:53 +00:00
func = None
args = []
kwargs = {}
if m.forward_date:
func = m.forward
elif hasattr(m, 'audio') and m.audio:
a = m.audio
func = bot.send_audio
args = [a.file_id, a.duration, a.performer, a.title, caption]
kwargs = dict(parse_mode='html')
elif hasattr(m, 'document') and m.document:
d = m.document
func = bot.send_document
args = [d.file_id, d.file_name, caption]
kwargs = dict(parse_mode='html')
elif hasattr(m, 'photo') and m.photo:
p = m.photo
func = bot.send_photo
args = [p[-1].file_id, caption]
kwargs = dict(parse_mode='html')
elif hasattr(m, 'sticker') and m.sticker:
s = m.sticker
func = bot.send_sticker
args = [s.file_id]
elif hasattr(m, 'video') and m.video:
v = m.video
func = bot.send_video
args = [v.file_id, v.duration, caption]
kwargs = dict(parse_mode='html')
elif hasattr(m, 'voice') and m.voice:
v = m.voice
func = bot.send_voice
args = [v.file_id, v.duration, caption]
kwargs = dict(parse_mode='html')
elif hasattr(m, 'video_note') and m.video_note:
vn = m.video_note
func = bot.send_video_note
args = [vn.file_id, vn.duration, vn.length]
elif hasattr(m, 'contact') and m.contact:
c = m.contact
func = bot.send_contact
args = [c.phone_number, c.first_name, c.last_name]
elif hasattr(m, 'location') and m.location:
l = m.location
func = bot.send_location
args = [l.latitude, l.longitude]
elif hasattr(m, 'venue') and m.venue:
v = m.venue
l = v.location
func = bot.send_venue
args = [l.latitude, l.longitude, v.title, v.address, v.foursquare_id]
elif hasattr(m, 'text') and m.text:
func = bot.send_message
args = [text, 'html']
if not _antispam(args):
return m.reply_text('Не вайпи', quote=True)
2019-03-17 12:07:38 +00:00
remove_uids = []
2020-05-19 14:43:53 +00:00
if func:
for uid, user in users.items():
sleep(.02)
2019-03-15 23:27:42 +00:00
2020-05-19 14:43:53 +00:00
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)
2019-03-15 23:27:42 +00:00
2020-05-19 14:43:53 +00:00
try:
r = func(*([uid] + args), **kwargs, 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()
2019-03-15 23:27:42 +00:00
conn.root.counter += 1
2019-03-03 23:23:12 +00:00
commit()
2019-03-17 12:07:38 +00:00
for uid in remove_uids:
_remove_user(uid)
2018-11-20 19:51:15 +00:00
2020-05-19 14:43:53 +00:00
def task_queue(u: Updater, stop_signal: Event):
2019-03-03 23:23:12 +00:00
while True:
2020-05-19 14:43:53 +00:00
if not u.running or stop_signal.is_set():
2019-03-03 23:23:12 +00:00
return
2018-11-20 19:51:15 +00:00
2019-03-03 23:23:12 +00:00
try:
m = queue.get(timeout=1) # type: Message
2019-03-16 23:45:26 +00:00
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)
2019-03-03 23:23:12 +00:00
except Empty:
pass
except:
2019-03-15 23:27:42 +00:00
traceback.print_exc()
2019-03-03 23:23:12 +00:00
sentry_sdk.capture_exception()
2018-11-20 19:51:15 +00:00
2020-05-19 14:43:53 +00:00
def main():
updater = Updater(BOT_TOKEN, workers=4, use_context=True)
2019-03-03 23:23:12 +00:00
updater.dispatcher.add_handler(CommandHandler('start', welcome, Filters.private))
updater.dispatcher.add_handler(CommandHandler('stop', unsubscribe, Filters.private))
2019-03-16 22:10:40 +00:00
updater.dispatcher.add_handler(CommandHandler('add', add_user, Filters.chat(MANAGEMENT_CHAT), pass_args=True))
2019-03-03 23:23:12 +00:00
updater.dispatcher.add_handler(CallbackQueryHandler(add_user, pattern=r'^add (\d+)$', pass_groups=True))
2019-03-16 22:10:40 +00:00
updater.dispatcher.add_handler(CommandHandler('remove', remove_user, Filters.chat(MANAGEMENT_CHAT), pass_args=True))
2019-03-03 23:23:12 +00:00
updater.dispatcher.add_handler(CallbackQueryHandler(remove_user, pattern=r'^remove (\d+)$', pass_groups=True))
2019-04-26 18:00:51 +00:00
updater.dispatcher.add_handler(CommandHandler('users', users_list, Filters.private))
2019-03-03 23:23:12 +00:00
updater.dispatcher.add_handler(MessageHandler(Filters.private, msg))
2018-11-20 19:51:15 +00:00
updater.start_polling()
2019-03-03 23:23:12 +00:00
2020-05-19 14:43:53 +00:00
stop_signal = Event()
tq = Thread(target=task_queue, args=(updater, stop_signal))
2019-03-03 23:23:12 +00:00
tq.start()
logging.warning('LONO has started')
2018-11-20 19:51:15 +00:00
updater.idle()
2020-05-19 14:43:53 +00:00
stop_signal.set()
2019-03-03 23:23:12 +00:00
logging.warning('LONO is stopping...')
commit()
conn.close()
2020-05-19 14:43:53 +00:00
if __name__ == '__main__':
main()