shift 支持转发媒体组

shift/main.py

@@ -1,22 +1,32 @@
 """ PagerMaid module for channel help. """
 
 import contextlib
+import datetime
 from asyncio import sleep
 from random import uniform
-from typing import Any, List, Union, Set
+from typing import Any, List, Literal, Optional
 
-from pyrogram.enums.chat_type import ChatType
-from pyrogram.errors.exceptions.flood_420 import FloodWait
+import pytz
+from pyrogram.enums import ChatType
+from pyrogram.errors import FloodWait
 from pyrogram.types import Chat
 
-from pagermaid import log, logs
+from pagermaid import log, Config
 from pagermaid.enums import Client, Message
+from pagermaid.enums.command import CommandHandler
 from pagermaid.listener import listener
-from pagermaid.single_utils import sqlite
+from pagermaid.services import bot, scheduler, sqlite, logs
 from pagermaid.utils import lang
 
 WHITELIST = [-1001441461877]
+AVAILABLE_OPTIONS_TYPE = Literal["silent", "text", "all", "photo", "document", "video"]
 AVAILABLE_OPTIONS = {"silent", "text", "all", "photo", "document", "video"}
+HELP_TEXT = """set [from channel] [to channel] (silent) 自动转发频道新消息（可以使用频道用户名或者 id）
+del [from channel] 删除转发
+backup [from channel] [to channel] (silent) 备份频道（可以使用频道用户名或者 id）
+list 顯示目前轉發的頻道
+选项说明：
+silent: 禁用通知, text: 文字, all: 全部訊息都傳, photo: 圖片, document: 檔案, video: 影片"""
 
 
 def try_cast_or_fallback(val: Any, t: type) -> Any:
@@ -28,142 +38,144 @@ def try_cast_or_fallback(val: Any, t: type) -> Any:
 
 def check_chat_available(chat: Chat):
     assert (
-            chat.type in [ChatType.CHANNEL, ChatType.GROUP, ChatType.SUPERGROUP, ChatType.BOT, ChatType.PRIVATE]
-            and not chat.has_protected_content
+        chat.type
+        in [
+            ChatType.CHANNEL,
+            ChatType.GROUP,
+            ChatType.SUPERGROUP,
+            ChatType.BOT,
+            ChatType.PRIVATE,
+        ]
+        and not chat.has_protected_content
     )
 
 
 @listener(
     command="shift",
     description="开启转发频道新消息功能",
-    parameters="set [from channel] [to channel] (silent) 自动转发频道新消息（可以使用频道用户名或者 id）\n"
-               "del [from channel] 删除转发\n"
-               "backup [from channel] [to channel] (silent) 备份频道（可以使用频道用户名或者 id）\n"
-               "list 顯示目前轉發的頻道\n\n"
-               "选项说明：\n"
-               "silent: 禁用通知, text: 文字, all: 全部訊息都傳, photo: 圖片, document: 檔案, video: 影片",
+    parameters=HELP_TEXT,
 )
-async def shift_set(client: Client, message: Message):
-    if not message.parameter:
-        await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
-        return
-    if message.parameter[0] == "set":
-        if len(message.parameter) < 3:
-            return await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
-        options = set(message.parameter[3:] if len(message.parameter) > 3 else ())
-        if set(options).difference(AVAILABLE_OPTIONS):
-            return await message.edit("出错了呜呜呜 ~ 无法识别的选项。")
-        # 检查来源频道
-        try:
-            source = await client.get_chat(
-                try_cast_or_fallback(message.parameter[1], int)
-            )
-            assert isinstance(source, Chat)
-            check_chat_available(source)
-        except Exception:
-            return await message.edit("出错了呜呜呜 ~ 无法识别的来源对话。")
-        if source.id in WHITELIST:
-            return await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
-        # 检查目标频道
-        try:
-            target = await client.get_chat(
-                try_cast_or_fallback(message.parameter[2], int)
-            )
-            assert isinstance(target, Chat)
-        except Exception:
-            return await message.edit("出错了呜呜呜 ~ 无法识别的目标对话。")
-        if target.id in WHITELIST:
-            await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
-            return
-        sqlite[f"shift.{source.id}"] = target.id
-        sqlite[f"shift.{source.id}.options"] = (
-            message.parameter[3:] if len(message.parameter) > 3 else ["all"]
-        )
-        await message.edit(f"已成功配置将对话 {source.id} 的新消息转发到 {target.id} 。")
-        await log(f"已成功配置将对话 {source.id} 的新消息转发到 {target.id} 。")
-    elif message.parameter[0] == "del":
-        if len(message.parameter) != 2:
-            return await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
-        # 检查来源频道
-        try:
-            source = await client.get_chat(
-                try_cast_or_fallback(message.parameter[1], int)
-            )
-            assert isinstance(source, Chat)
-        except Exception:
-            return await message.edit("出错了呜呜呜 ~ 无法识别的来源对话。")
-        try:
-            del sqlite[f"shift.{source.id}"]
-            with contextlib.suppress(Exception):
-                del sqlite[f"shift.{source.id}.options"]
-        except Exception:
-            return await message.edit("emm...当前对话不存在于自动转发列表中。")
-        await message.edit(f"已成功关闭对话 {str(source.id)} 的自动转发功能。")
-        await log(f"已成功关闭对话 {str(source.id)} 的自动转发功能。")
-    elif message.parameter[0] == "backup":
-        if len(message.parameter) < 3:
-            return await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
-        options = set(message.parameter[3:] if len(message.parameter) > 3 else ())
-        if set(options).difference(AVAILABLE_OPTIONS):
-            return await message.edit("出错了呜呜呜 ~ 无法识别的选项。")
-        # 检查来源频道
-        try:
-            source = await client.get_chat(
-                try_cast_or_fallback(message.parameter[1], int)
-            )
-            assert isinstance(source, Chat)
-            check_chat_available(source)
-        except Exception:
-            return await message.edit("出错了呜呜呜 ~ 无法识别的来源对话。")
-        if source.id in WHITELIST:
-            return await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
-        # 检查目标频道
-        try:
-            target = await client.get_chat(
-                try_cast_or_fallback(message.parameter[2], int)
-            )
-            assert isinstance(target, Chat)
-        except Exception:
-            return await message.edit("出错了呜呜呜 ~ 无法识别的目标对话。")
-        if target.id in WHITELIST:
-            return await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
-        # 开始遍历消息
-        await message.edit(f"开始备份频道 {source.id} 到 {target.id} 。")
+async def shift_func(message: Message):
+    await message.edit(HELP_TEXT)
 
-        # 如果有把get_chat_history方法merge進去就可以實現從舊訊息到新訊息,https://github.com/pyrogram/pyrogram/pull/1046
-        # async for msg in client.get_chat_history(source.id,reverse=True):
 
-        async for msg in client.search_messages(source.id):  # type: ignore
-            await sleep(uniform(0.5, 1.0))
-            await loosely_forward(
-                message,
-                msg,
-                target.id,
-                options,
-                disable_notification="silent" in options,
-            )
-        await message.edit(f"备份频道 {source.id} 到 {target.id} 已完成。")
-    # 列出要轉存的頻道
-    elif message.parameter[0] == "list":
-        from_ids = list(
-            filter(
-                lambda x: (x.startswith("shift.") and (not x.endswith("options"))),
-                list(sqlite.keys()),
-            )
-        )
-        if not from_ids:
-            return await message.edit("沒有要轉存的頻道")
-        output = "總共有 %d 個頻道要轉存\n\n" % len(from_ids)
-        for from_id in from_ids:
-            to_id = sqlite[from_id]
-            output += "%s -> %s\n" % (
-                format_channel_id(from_id[6:]),
-                format_channel_id(to_id),
-            )
-        await message.edit(output)
-    else:
-        await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
+shift_func: "CommandHandler"
+
+
+@shift_func.sub_command(command="set")
+async def shift_func_set(client: Client, message: Message):
+    if len(message.parameter) < 3:
+        return await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
+    options = set(message.parameter[3:] if len(message.parameter) > 3 else ())
+    if set(options).difference(AVAILABLE_OPTIONS):
+        return await message.edit("出错了呜呜呜 ~ 无法识别的选项。")
+    # 检查来源频道
+    try:
+        source = await client.get_chat(try_cast_or_fallback(message.parameter[1], int))
+        assert isinstance(source, Chat)
+        check_chat_available(source)
+    except Exception:
+        return await message.edit("出错了呜呜呜 ~ 无法识别的来源对话。")
+    if source.id in WHITELIST:
+        return await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
+    # 检查目标频道
+    try:
+        target = await client.get_chat(try_cast_or_fallback(message.parameter[2], int))
+        assert isinstance(target, Chat)
+    except Exception:
+        return await message.edit("出错了呜呜呜 ~ 无法识别的目标对话。")
+    if target.id in WHITELIST:
+        await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
         return
+    sqlite[f"shift.{source.id}"] = target.id
+    sqlite[f"shift.{source.id}.options"] = (
+        message.parameter[3:] if len(message.parameter) > 3 else ["all"]
+    )
+    await message.edit(f"已成功配置将对话 {source.id} 的新消息转发到 {target.id} 。")
+    await log(f"已成功配置将对话 {source.id} 的新消息转发到 {target.id} 。")
+
+
+@shift_func.sub_command(command="del")
+async def shift_func_del(message: Message):
+    if len(message.parameter) != 2:
+        return await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
+    # 检查来源频道
+    try:
+        source = try_cast_or_fallback(message.parameter[1], int)
+        assert isinstance(source, int)
+    except Exception:
+        return await message.edit("出错了呜呜呜 ~ 无法识别的来源对话。")
+    try:
+        del sqlite[f"shift.{source}"]
+        with contextlib.suppress(Exception):
+            del sqlite[f"shift.{source}.options"]
+    except Exception:
+        return await message.edit("emm...当前对话不存在于自动转发列表中。")
+    await message.edit(f"已成功关闭对话 {str(source)} 的自动转发功能。")
+    await log(f"已成功关闭对话 {str(source)} 的自动转发功能。")
+
+
+@shift_func.sub_command(command="backup")
+async def shift_func_backup(client: Client, message: Message):
+    if len(message.parameter) < 3:
+        return await message.edit(f"{lang('error_prefix')}{lang('arg_error')}")
+    options = set(message.parameter[3:] if len(message.parameter) > 3 else ())
+    if set(options).difference(AVAILABLE_OPTIONS):
+        return await message.edit("出错了呜呜呜 ~ 无法识别的选项。")
+    # 检查来源频道
+    try:
+        source = await client.get_chat(try_cast_or_fallback(message.parameter[1], int))
+        assert isinstance(source, Chat)
+        check_chat_available(source)
+    except Exception:
+        return await message.edit("出错了呜呜呜 ~ 无法识别的来源对话。")
+    if source.id in WHITELIST:
+        return await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
+    # 检查目标频道
+    try:
+        target = await client.get_chat(try_cast_or_fallback(message.parameter[2], int))
+        assert isinstance(target, Chat)
+    except Exception:
+        return await message.edit("出错了呜呜呜 ~ 无法识别的目标对话。")
+    if target.id in WHITELIST:
+        return await message.edit("出错了呜呜呜 ~ 此对话位于白名单中。")
+    # 开始遍历消息
+    await message.edit(f"开始备份频道 {source.id} 到 {target.id} 。")
+
+    # 如果有把get_chat_history方法merge進去就可以實現從舊訊息到新訊息,https://github.com/pyrogram/pyrogram/pull/1046
+    # async for msg in client.get_chat_history(source.id,reverse=True):
+
+    async for msg in client.search_messages(source.id):  # type: ignore
+        await sleep(uniform(0.5, 1.0))
+        await loosely_forward(
+            message,
+            msg,
+            target.id,
+            list(options),
+            disable_notification="silent" in options,
+        )
+    await message.edit(f"备份频道 {source.id} 到 {target.id} 已完成。")
+
+
+# 列出要轉存的頻道
+@shift_func.sub_command(command="list")
+async def shift_func_list(message: Message):
+    from_ids = list(
+        filter(
+            lambda x: (x.startswith("shift.") and (not x.endswith("options"))),
+            list(sqlite.keys()),
+        )
+    )
+    if not from_ids:
+        return await message.edit("沒有要轉存的頻道")
+    output = "總共有 %d 個頻道要轉存\n\n" % len(from_ids)
+    for from_id in from_ids:
+        to_id = sqlite[from_id]
+        output += "%s -> %s\n" % (
+            format_channel_id(from_id[6:]),
+            format_channel_id(to_id),
+        )
+    await message.edit(output)
 
 
 def format_channel_id(channel_id: str):
@@ -174,62 +186,123 @@ def format_channel_id(channel_id: str):
 @listener(is_plugin=True, incoming=True, ignore_edited=True)
 async def shift_channel_message(message: Message):
     """Event handler to auto forward channel messages."""
-    d = dict(sqlite)
     source = message.chat.id
-
-    # 找訊息類型video、document...
-    media_type = message.media.value if message.media else "text"
-    target = d.get(f"shift.{source}")
+    target = sqlite.get(f"shift.{source}")
     if not target:
         return
     if message.chat.has_protected_content:
         del sqlite[f"shift.{source}"]
         return
-    options = d.get(f"shift.{source}.options") or []
+    options = sqlite.get(f"shift.{source}.options") or []
 
     with contextlib.suppress(Exception):
-        if (not options) or "all" in options:
-            await message.forward(
+        if message.media_group_id:
+            add_or_replace_forward_group_media(
                 target,
+                source,
+                message.id,
+                message.media_group_id,
+                options,
                 disable_notification="silent" in options,
             )
-        elif media_type in options:
-            await message.forward(
-                target,
-                disable_notification="silent" in options,
-            )
-        else:
-            logs.debug("skip message type: %s", media_type)
+            return
+        await loosely_forward(
+            None,
+            message,
+            target,
+            options,
+            disable_notification="silent" in options,
+        )
 
 
 async def loosely_forward(
-        notifier: Message,
-        message: Message,
-        chat_id: int,
-        options: Union[List[str], Set[str]],
-        disable_notification: bool = False,
+    notifier: Optional[Message],
+    message: Message,
+    chat_id: int,
+    options: List[AVAILABLE_OPTIONS_TYPE],
+    disable_notification: bool = False,
 ):
     # 找訊息類型video、document...
     media_type = message.media.value if message.media else "text"
+    if (not options) or "all" in options:
+        await forward_messages(
+            chat_id, message.chat.id, [message.id], disable_notification, notifier
+        )
+    elif media_type in options:
+        await forward_messages(
+            chat_id, message.chat.id, [message.id], disable_notification, notifier
+        )
+    else:
+        logs.debug("skip message type: %s", media_type)
+
+
+async def forward_messages(
+    cid: int,
+    from_id: int,
+    message_ids: List[int],
+    disable_notification: bool,
+    notifier: Optional["Message"],
+):
     try:
-        if (not options) or "all" in options:
-            await message.forward(
-                chat_id,
-                disable_notification=disable_notification,
-            )
-        elif media_type in options:
-            await message.forward(
-                chat_id,
-                disable_notification=disable_notification,
-            )
-        else:
-            logs.debug("skip message type: %s", media_type)
-        # await message.forward(chat_id, disable_notification=disable_notification)
+        await bot.forward_messages(
+            cid, from_id, message_ids, disable_notification=disable_notification
+        )
     except FloodWait as ex:
-        min: int = ex.value  # type: ignore
-        delay = min + uniform(0.5, 1.0)
-        await notifier.edit(f"触发 Flood ，暂停 {delay} 秒。")
+        min_time: int = ex.value  # type: ignore
+        delay = min_time + uniform(0.5, 1.0)
+        if notifier:
+            await notifier.edit(f"触发 Flood ，暂停 {delay} 秒。")
         await sleep(delay)
-        await loosely_forward(notifier, message, chat_id, options, disable_notification)
+        await forward_messages(
+            cid, from_id, message_ids, disable_notification, notifier
+        )
     except Exception:
         pass  # drop other errors
+
+
+async def forward_group_media(
+    cid: int,
+    from_id: int,
+    group_id: int,
+    options: List[AVAILABLE_OPTIONS_TYPE],
+    disable_notification: bool,
+):
+    try:
+        msgs = await bot.get_media_group(from_id, group_id)
+    except Exception:
+        logs.debug("get_media_group failed for %d %d", from_id, group_id)
+        return
+    real_msgs = []
+    for message in msgs:
+        media_type = message.media.value if message.media else "text"
+        if (not options) or "all" in options:
+            real_msgs.append(message)
+        elif media_type in options:
+            real_msgs.append(message)
+        else:
+            logs.debug("skip message type: %s", media_type)
+    if not real_msgs:
+        return
+    real_msgs_ids = [msg.id for msg in real_msgs]
+    await forward_messages(cid, from_id, real_msgs_ids, disable_notification, None)
+
+
+def add_or_replace_forward_group_media(
+    cid: int,
+    from_id: int,
+    message_id: int,
+    group_id: int,
+    options: List[AVAILABLE_OPTIONS_TYPE],
+    disable_notification: bool,
+):
+    key = f"shift.forward_group_media.{group_id}"
+    scheduler.add_job(
+        forward_group_media,
+        trigger="date",
+        args=(cid, from_id, message_id, options, disable_notification),
+        id=key,
+        name=key,
+        replace_existing=True,
+        run_date=datetime.datetime.now(pytz.timezone(Config.TIME_ZONE))
+        + datetime.timedelta(seconds=4),
+    )