lono/main.py
2021-04-01 11:25:31 +03:00

432 lines
16 KiB
Python
Executable File
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
import json
import logging
import os
import re
import traceback
from datetime import datetime, timedelta
from hashlib import sha1
from html import escape
from queue import Queue, Empty
from time import sleep
from threading import Thread, Event
from typing import Dict, List
import sentry_sdk
from redis import Redis
from telegram.error import Unauthorized, TelegramError
from telegram.ext import Updater, CommandHandler, MessageHandler, Filters, CallbackQueryHandler, CallbackContext
from telegram import Message, Update, Bot, InlineKeyboardMarkup, InlineKeyboardButton, User, InputMediaPhoto, \
InputMediaVideo, InputMediaAnimation, InputMediaAudio, InputMediaDocument, PhotoSize
from config import BOT_TOKEN, SENTRY_DSN, MANAGEMENT_CHAT, DEBUG
from db import get_conn, Subscriber, 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()
redis = Redis()
MAX_MESSAGE_LENGTH = 4096
MAX_CAPTION_LENGTH = 1024
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
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(update: Update, ctx: CallbackContext):
if DEBUG:
_add_user(ctx.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(ctx.bot, update.message.from_user)
def unsubscribe(update: Update, ctx: CallbackContext):
user = _remove_user(update.message.chat_id)
update.message.reply_text('Вы были отписаны от бота. '
'Обратитесь к @lono_contactbot если вы хотите подписаться снова.')
ctx.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(update: Update, ctx: CallbackContext):
if update.callback_query:
update.callback_query.answer()
if ctx.match:
if update.callback_query.message.chat.id != MANAGEMENT_CHAT:
return
uid = ctx.match.group(1)
elif ctx.args:
uid = ctx.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 ctx.bot.send_message(MANAGEMENT_CHAT, 'Укажите ID пользователя или ответьте на его сообщение')
try:
uid = int(uid)
except (ValueError, TypeError):
pass
try:
user = _add_user(ctx.bot, uid)
if update.callback_query:
update.callback_query.message.edit_reply_markup()
ctx.bot.send_message(MANAGEMENT_CHAT, f'<a href="tg://user?id={uid}">{escape(user.name)}</a> ({uid}) был добавлен',
parse_mode='html')
ctx.bot.send_message(uid, 'Добро пожаловать. Снова.')
except TelegramError as e:
ctx.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(update: Update, ctx: CallbackContext):
if update.callback_query:
update.callback_query.answer()
if ctx.match:
if update.callback_query.message.chat.id != MANAGEMENT_CHAT:
return
uid = ctx.match.groups(1)
elif ctx.args:
uid = ctx.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 ctx.bot.send_message(MANAGEMENT_CHAT, 'Укажите ID пользователя или ответьте на его сообщение')
try:
uid = int(uid)
except (ValueError, TypeError):
pass
try:
user = _remove_user(uid)
ctx.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:
ctx.bot.send_message(MANAGEMENT_CHAT, f'Пользователь id={uid} не был найден')
def users(update: Update, ctx: CallbackContext):
send_users_list()
def msg(update: Update, ctx: CallbackContext):
queue.put(update.message)
def _sign_text(text, m: Message, limit):
if not text:
text = ''
text = re.sub(r'<a href="(tg://.*?)">.*?</a>', '<a href="\\1">\\1</a>', text)
sign = ''
if text.startswith('!sign') or text.startswith('/sign'):
text = text[5:]
sign = f'\n\n— <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 users_list(update: Update, ctx: CallbackContext):
current_chat = update.effective_chat.id
subs = conn.root.subscribers # type: Dict[int, Subscriber]
if current_chat not in subs:
return
try:
with open('fake_users.json') as f:
subs = {data[0]: Subscriber(int(data[0]), data[1]) for data in json.load(f).items()}
except:
subs = conn.root.subscribers
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')
def _process_message(bot: Bot, m: Message):
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()
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('Пожалуйста, GTFO')
if m.text and len(m.text) > 140 or m.caption and len(m.caption) > 140:
return m.reply_text('Сообщение не может содержать более 140 символов')
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]
func = None
args = []
hash_args = None
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]
hash_args = [p[-1].file_unique_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(hash_args or args):
return m.reply_text('Не вайпи', quote=True)
remove_uids = []
if func:
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:
if func != m.forward:
kwargs['reply_to_message_id'] = reply_to_message_id
r = func(*([uid] + args), **kwargs)
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, stop_signal: Event):
while True:
if not u.running or stop_signal.is_set():
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()
def main():
updater = Updater(BOT_TOKEN, workers=4, use_context=True)
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+)$'))
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+)$'))
updater.dispatcher.add_handler(CommandHandler('users', users_list, Filters.private))
updater.dispatcher.add_handler(MessageHandler(Filters.private, msg))
updater.start_polling()
stop_signal = Event()
tq = Thread(target=task_queue, args=(updater, stop_signal))
tq.start()
logging.warning('LONO has started')
updater.idle()
stop_signal.set()
logging.warning('LONO is stopping...')
commit()
conn.close()
if __name__ == '__main__':
main()