Add support for groups

zeke-cmd · Jun 28, 2020 · 108ddab · 108ddab
1 parent 393a5ba
commit 108ddab
Show file tree

Hide file tree

Showing 14 changed files with 214 additions and 40 deletions.
diff --git a/telegram_gcloner/handlers/add_group.py b/telegram_gcloner/handlers/add_group.py
@@ -8,7 +8,6 @@
 from telegram.utils.helpers import mention_html
 
 from utils.config_loader import config
-from utils.restricted import restricted
 
 logger = logging.getLogger(__name__)
 
@@ -18,7 +17,6 @@ def init(dispatcher: Dispatcher):
     dispatcher.add_handler(MessageHandler(Filters.status_update.new_chat_members, add_group))
 
 
-@restricted
 def add_group(update, context):
     message = 'joined: {} {}'.format(mention_html(update.message.new_chat_members[0].id,
                                                   html.escape(update.message.new_chat_members[0].full_name)),

diff --git a/telegram_gcloner/handlers/cancel.py b/telegram_gcloner/handlers/cancel.py
@@ -4,6 +4,8 @@
 
 from telegram.ext import Dispatcher, CallbackQueryHandler
 
+from utils.restricted import restricted
+
 logger = logging.getLogger(__name__)
 
 
@@ -12,7 +14,11 @@ def init(dispatcher: Dispatcher):
     dispatcher.add_handler(CallbackQueryHandler(cancel, pattern=r'^cancel$'))
 
 
+@restricted
 def cancel(update, context):
     query = update.callback_query
-    # query.message.edit_reply_markup(reply_markup=None)
-    query.message.delete()
+    if query.message.chat_id < 0 and \
+            (not query.message.reply_to_message or
+             query.from_user.id != query.message.reply_to_message.from_user.id):
+        # query.message.edit_reply_markup(reply_markup=None)
+        query.message.delete()
diff --git a/telegram_gcloner/handlers/choose_folder.py b/telegram_gcloner/handlers/choose_folder.py
@@ -35,7 +35,15 @@ def init(dispatcher: Dispatcher):
                                                 pattern=r'^(?:un)?set_folders(:?_page#\d+)?(?:\,[\dA-Za-z\-_]+)?$'))
 
 
+@restricted
 def chosen_folder(update, context):
+    query = update.callback_query
+    if query.message.chat_id < 0 and \
+            (not query.message.reply_to_message or
+             query.from_user.id != query.message.reply_to_message.from_user.id):
+        alert_users(context, update.effective_user, 'invalid caller', query.data)
+        query.answer(text='哟呵', show_alert=True)
+        return
     if update.effective_user.id in config.USER_IDS or update.effective_user.id in context.bot_data['vip']:
         max_folders = default_max_folders_vip
     else:
@@ -126,7 +134,9 @@ def choose_folder(update, context):
 
     if query:
         logger.debug('{}: {}'.format(update.effective_user.id, query.data))
-        if query.from_user.id != query.message.chat_id:
+        if query.message.chat_id < 0 and \
+                (not query.message.reply_to_message or
+                 query.from_user.id != query.message.reply_to_message.from_user.id):
             alert_users(context, update.effective_user, 'invalid caller', query.data)
             query.answer(text='哟呵', show_alert=True)
             return
@@ -229,6 +239,12 @@ def set_folders(update, context):
         rsp.done.wait(timeout=60)
         message_id = rsp.result().message_id
     else:
+        if query.message.chat_id < 0 and \
+                (not query.message.reply_to_message or
+                 query.from_user.id != query.message.reply_to_message.from_user.id):
+            alert_users(context, update.effective_user, 'invalid caller', query.data)
+            query.answer(text='哟呵', show_alert=True)
+            return
         message_id = query.message.message_id
     folder_ids = context.user_data.get(udkey_folders, None)
 

diff --git a/telegram_gcloner/handlers/contact.py b/telegram_gcloner/handlers/contact.py
@@ -7,8 +7,9 @@
 from telegram.ext import Dispatcher, CommandHandler
 from telegram.utils.helpers import mention_html
 
+from utils.callback import callback_delete_message
 from utils.config_loader import config
-from utils.restricted import restricted
+from utils.restricted import restricted_private
 
 logger = logging.getLogger(__name__)
 
@@ -18,7 +19,7 @@ def init(dispatcher: Dispatcher):
     dispatcher.add_handler(CommandHandler('4999baoyue', contact, pass_args=True))
 
 
-@restricted
+@restricted_private
 def contact(update, context):
     if update.message.text.strip('/4999baoyue'):
         context.bot.send_message(chat_id=config.USER_IDS[0],
@@ -29,8 +30,22 @@ def contact(update, context):
         context.bot.forward_message(chat_id=config.USER_IDS[0],
                                     from_chat_id=update.message.chat_id,
                                     message_id=update.message.message_id)
-        update.message.reply_text('收到。')
+        rsp = update.message.reply_text('收到。')
+        rsp.done.wait(timeout=60)
+        message_id = rsp.result().message_id
+        if update.message.chat_id < 0:
+            context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                       context=(update.message.chat_id, message_id))
+            context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                       context=(update.message.chat_id, update.message.message_id))
     else:
-        update.message.reply_text('这么害羞，不说点啥？\n' +
-                                  config.AD_STRING.format(context.bot.username),
-                                  ParseMode.HTML)
+        rsp = update.message.reply_text('这么害羞，不说点啥？\n' +
+                                        config.AD_STRING.format(context.bot.username),
+                                        ParseMode.HTML)
+        rsp.done.wait(timeout=60)
+        message_id = rsp.result().message_id
+        if update.message.chat_id < 0:
+            context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                       context=(update.message.chat_id, message_id))
+            context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                       context=(update.message.chat_id, update.message.message_id))
diff --git a/telegram_gcloner/handlers/get_help.py b/telegram_gcloner/handlers/get_help.py
@@ -21,12 +21,15 @@ def get_help(update, context):
     message = '发送google drive链接，或者转发带有google drive的信息即可手动转存。\n' \
               '需要使用 /sa 和 /folders 进行配置\n\n' \
               '以下是本BOT的命令：\n\n' \
-              '/start - 没啥用（仅限私聊）\n' \
               '/folders - 设置收藏文件夹\n' \
-              '/sa - 上传包含sa的ZIP文件夹，在标题写上/sa设置Service Account\n' \
-              '/4999baoyue - 商业洽谈，请附上留言\n' \
+              '/sa - 仅限私聊，上传包含sa的ZIP文件夹，在标题写上/sa设置Service Account\n' \
+              '/4999baoyue - 仅限私聊，商业洽谈，请附上留言\n' \
               '/help - 输出本帮助\n'
-    message_sent = update.message.reply_text(message)
+    rsp = update.message.reply_text(message)
+    rsp.done.wait(timeout=60)
+    message_id = rsp.result().message_id
     if update.message.chat_id < 0:
         context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
-                                   context=(update.message.chat_id, message_sent.message_id))
+                                   context=(update.message.chat_id, message_id))
+        context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                   context=(update.message.chat_id, update.message.message_id))
diff --git a/telegram_gcloner/handlers/get_id.py b/telegram_gcloner/handlers/get_id.py
@@ -4,6 +4,7 @@
 
 from telegram.ext import Dispatcher, CommandHandler
 
+from utils.callback import callback_delete_message
 from utils.config_loader import config
 from utils.restricted import restricted
 
@@ -17,7 +18,13 @@ def init(dispatcher: Dispatcher):
 
 @restricted
 def get_id(update, context):
-    if len(config.USER_IDS) and (update.message.chat.id not in config.USER_IDS):
-        return
-    update.message.reply_text(update.message.chat.id)
-    logger.info('telegram user {0} has requested its id.'.format(update.message.chat.id))
+    logger.info('telegram user {0} has requested its id.'.format(update.effective_user.id))
+    rsp = update.message.reply_text(update.effective_user.id)
+    rsp.done.wait(timeout=60)
+    message_id = rsp.result().message_id
+
+    if update.message.chat_id < 0:
+        context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                   context=(update.message.chat_id, message_id))
+        context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                   context=(update.message.chat_id, update.message.message_id))
diff --git a/telegram_gcloner/handlers/process_drive_links.py b/telegram_gcloner/handlers/process_drive_links.py
@@ -107,6 +107,12 @@ def save_to_folder_page(update, context):
     callback_query_prefix = 'save_to_folder'
 
     query = update.callback_query
+    if query.message.chat_id < 0 and \
+            (not query.message.reply_to_message or
+             query.from_user.id != query.message.reply_to_message.from_user.id):
+        alert_users(context, update.effective_user, 'invalid caller', query.data)
+        query.answer(text='哟呵', show_alert=True)
+        return
     match = re.search(r'^save_to_folder_page#(\d+)$', query.data)
     if not match:
         alert_users(context, update.effective_user, 'invalid query data', query.data)
@@ -133,7 +139,12 @@ def save_to_folder_page(update, context):
 
 def save_to_folder(update, context):
     query = update.callback_query
-
+    if query.message.chat_id < 0 and \
+            (not query.message.reply_to_message or
+             query.from_user.id != query.message.reply_to_message.from_user.id):
+        alert_users(context, update.effective_user, 'invalid caller', query.data)
+        query.answer(text='哟呵', show_alert=True)
+        return
     match = re.search(r'^save_to_folder(?:_page#[\d]+)?,\s*([\dA-Za-z\-_]+)$', query.data)
     fav_folders = context.user_data.get(udkey_folders, {})
     if not match or match.group(1) not in fav_folders:

diff --git a/telegram_gcloner/handlers/process_message.py b/telegram_gcloner/handlers/process_message.py
@@ -16,12 +16,20 @@
 def init(dispatcher: Dispatcher):
     """Provide handlers initialization."""
     dispatcher.add_handler(
-        MessageHandler(Filters.group & Filters.chat(config.GROUP_IDS), process_message))
+        MessageHandler(Filters.group & Filters.chat(config.GROUP_IDS) &
+                       (Filters.text | Filters.caption) &
+                       ~Filters.update.edited_message,
+                       process_message))
     dispatcher.add_handler(
-        MessageHandler(Filters.chat(config.USER_IDS[0]) & (Filters.text | Filters.caption),
+        MessageHandler(Filters.chat(config.USER_IDS[0]) &
+                       (Filters.text | Filters.caption) &
+                       ~Filters.update.edited_message,
                        process_message_from_authorised_user))
     dispatcher.add_handler(
-        MessageHandler((~Filters.group) & (Filters.text | Filters.photo | Filters.document), process_message))
+        MessageHandler((~Filters.group) &
+                       (Filters.text | Filters.caption) &
+                       ~Filters.update.edited_message,
+                       process_message))
 
     dispatcher.add_handler(CallbackQueryHandler(ignore_callback, pattern=r'^#$'))
     dispatcher.add_handler(CallbackQueryHandler(get_warning))
@@ -62,11 +70,10 @@ def process_message_from_authorised_user(update, context):
 
 @restricted
 def process_message(update, context):
-    if not update.message or update.message.chat_id < 0:
+    if not update.message:
         return
     if update.message.chat_id == config.USER_IDS[0]:
         pass
-        # process_message_for_gclone(update, context)
     else:
         logger.debug(update.message)
         if update.message.caption:

diff --git a/telegram_gcloner/handlers/sa.py b/telegram_gcloner/handlers/sa.py
@@ -12,20 +12,20 @@
 from telegram.ext import Dispatcher, CommandHandler, MessageHandler, Filters
 
 from utils.config_loader import config
-from utils.restricted import restricted
+from utils.restricted import restricted_private
 
 logger = logging.getLogger(__name__)
 
 
 def init(dispatcher: Dispatcher):
     """Provide handlers initialization."""
-    dispatcher.add_handler(CommandHandler('sa', get_sa))
-    dispatcher.add_handler(MessageHandler(Filters.document, get_sa))
+    dispatcher.add_handler(CommandHandler('sa', get_sa, filters=~Filters.update.edited_message))
+    dispatcher.add_handler(MessageHandler(Filters.private & Filters.document, get_sa))
 
 
-@restricted
+@restricted_private
 def get_sa(update, context):
-    instruction_text = '请上传含有SA文件的ZIP文件，并在telegram标题写上/sa。\n' \
+    instruction_text = '请私聊上传含有SA文件的ZIP文件，并在telegram标题写上/sa。\n' \
                        '手机用户请先上传ZIP文件，再以/sa回复该信息。'
     if update.message and update.message.caption and update.message.caption.startswith('/sa'):
         document = update.message.document

diff --git a/telegram_gcloner/handlers/start.py b/telegram_gcloner/handlers/start.py
@@ -4,6 +4,8 @@
 
 from telegram.ext import Dispatcher, CommandHandler
 
+from utils.callback import callback_delete_message
+from utils.config_loader import config
 from utils.restricted import restricted
 
 logger = logging.getLogger(__name__)
@@ -16,6 +18,13 @@ def init(dispatcher: Dispatcher):
 
 @restricted
 def start(update, context):
-    update.message.reply_text('先上传包含SA文件的ZIP压缩包，并在标题填写 /sa\n'
-                              '再 /folders 设置收藏的文件夹。\n'
-                              '随后转发或直接发送带有google drive链接的内容即可。')
+    rsp = update.message.reply_text('先私聊上传包含SA文件的ZIP压缩包，并在标题填写 /sa\n'
+                                    '再 /folders 设置收藏的文件夹。\n'
+                                    '随后转发或直接发送带有google drive链接的内容即可。')
+    rsp.done.wait(timeout=60)
+    message_id = rsp.result().message_id
+    if update.message.chat_id < 0:
+        context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                   context=(update.message.chat_id, message_id))
+        context.job_queue.run_once(callback_delete_message, config.TIMER_TO_DELETE_MESSAGE,
+                                   context=(update.message.chat_id, update.message.message_id))
diff --git a/telegram_gcloner/handlers/stop_task.py b/telegram_gcloner/handlers/stop_task.py
@@ -22,6 +22,12 @@ def init(dispatcher: Dispatcher):
 @restricted
 def stop_task(update, context):
     query = update.callback_query
+    if query.message.chat_id < 0 and \
+            (not query.message.reply_to_message or
+             query.from_user.id != query.message.reply_to_message.from_user.id):
+        alert_users(context, update.effective_user, 'invalid caller', query.data)
+        query.answer(text='哟呵', show_alert=True)
+        return
     if query.data:
         match = re.search(regex_stop_task, query.data)
         if match:

diff --git a/telegram_gcloner/utils/config_loader.py b/telegram_gcloner/utils/config_loader.py
@@ -20,6 +20,7 @@ def __init__(self):
         self._user_ids = ''
         self._group_ids = ''
         self._base_path = os.path.dirname(os.path.dirname(__file__))
+        self.TIMER_TO_DELETE_MESSAGE = 10
 
     def load_config(self):
         logger.debug('Loading config')

diff --git a/telegram_gcloner/utils/fire_save_files.py b/telegram_gcloner/utils/fire_save_files.py
@@ -32,15 +32,19 @@ def run(self):
         self.owner = update.effective_user.id
         thread_id = self.ident
         is_multiple_ids = len(folder_ids) > 1
-        chat_id = update.effective_user.id
-        gd = GoogleDrive(chat_id)
+        chat_id = update.effective_chat.id
+        user_id = update.effective_user.id
+        gd = GoogleDrive(user_id)
         message = '目标目录：{}\n\n'.format(dest_folder['path'])
         inline_keyboard = InlineKeyboardMarkup(
             [[InlineKeyboardButton(text=f'停止', callback_data=f'stop_task,{thread_id}')]])
 
+        reply_message_id = update.callback_query.message.reply_to_message.message_id \
+            if update.callback_query.message.reply_to_message else None
         rsp = context.bot.send_message(chat_id=chat_id, text=message,
                                        parse_mode=ParseMode.HTML,
                                        disable_web_page_preview=True,
+                                       reply_to_message_id=reply_message_id,
                                        reply_markup=inline_keyboard)
         rsp.done.wait(timeout=60)
         message_id = rsp.result().message_id
@@ -230,7 +234,7 @@ def run(self):
         update.callback_query.message.edit_reply_markup(reply_markup=InlineKeyboardMarkup(
             [[InlineKeyboardButton(text='已完成', callback_data='cancel')]]))
 
-        tasks = thread_pool.get(update.effective_user.id, None)
+        tasks = thread_pool.get(user_id, None)
         if tasks:
             for t in tasks:
                 if t.ident == thread_id: