Bug修复

修复文件不在一个消息内时无法接受 多个文件，修复视频文件接受问题
#38 #36

bot/config.py

@@ -46,7 +46,7 @@ def change_password():
     }
 
     change_result = requests.put(url=change_url, json=zh_data, headers=headers)
-    
+
     if change_result.status_code==403:
         headers['X-Auth']=html.text
         change_result = requests.put(url=change_url, json=zh_data, headers=headers)

bot/modules/moretg.py

@@ -13,45 +13,145 @@
 nest_asyncio.apply()
 os.system("df -lh")
 
+import asyncio
+import functools
+import pyrogram
+
+def patch(obj):
+    def is_patchable(item):
+        return getattr(item[1], 'patchable', False)
+
+    def wrapper(container):
+        for name,func in filter(is_patchable, container.__dict__.items()):
+            old = getattr(obj, name, None)
+            setattr(obj, 'old'+name, old)
+            setattr(obj, name, func)
+        return container
+    return wrapper
+
+def patchable(func):
+    func.patchable = True
+    return func
+
+loop = asyncio.get_event_loop()
+
+
+class ListenerCanceled(Exception):
+    pass
+
+
+pyrogram.errors.ListenerCanceled = ListenerCanceled
+
+@patch(pyrogram.client.Client)
+class Client():
+    @patchable
+    def __init__(self, *args, **kwargs):
+        self.listening = {}
+        self.using_mod = True
+
+        self.old__init__(*args, **kwargs)
+
+    @patchable
+    async def listen(self, chat_id, filters=None, timeout=None):
+        if type(chat_id) != int:
+            chat = await self.get_chat(chat_id)
+            chat_id = chat.id
+
+        future = loop.create_future()
+        future.add_done_callback(
+            functools.partial(self.clear_listener, chat_id)
+        )
+        self.listening.update({
+            chat_id: {"future": future, "filters": filters}
+        })
+        return await asyncio.wait_for(future, timeout)
+
+    @patchable
+    async def ask(self, chat_id, text, filters=None, timeout=None, *args, **kwargs):
+        request = await self.send_message(chat_id, text, *args, **kwargs)
+        response = await self.listen(chat_id, filters, timeout)
+        response.request = request
+        return response
+
+    @patchable
+    async def none(self, chat_id, filters=None, timeout=None, *args, **kwargs):
+
+        response = await self.listen(chat_id, filters, timeout)
+
+        return response
+
+    @patchable
+    def clear_listener(self, chat_id, future):
+        if future == self.listening[chat_id]:
+            self.listening.pop(chat_id, None)
+
+    @patchable
+    def cancel_listener(self, chat_id):
+        listener = self.listening.get(chat_id)
+        if not listener or listener['future'].done():
+            return
+
+        listener['future'].set_exception(ListenerCanceled())
+        self.clear_listener(chat_id, listener['future'])
+
 
 async def start_down_telegram_file(client, message,file_list):
     try:
         if len(file_list)==0:
             answer = await client.ask(chat_id=message.chat.id, text='请发送TG文件,或输入 /cancel 取消')
+            #answer = await client.none(chat_id=message.chat.id)
         else:
-            answer = await client.ask(chat_id=message.chat.id, text=f'已接收{len(file_list)}个文件，请继续发送TG文件，输入 /finish 取消,或输入 /cancel 取消')
+            answer = await client.ask(chat_id=message.chat.id, text=f'已接收{len(file_list)}个文件，请继续发送TG文件，输入 /finish 结束,或输入 /cancel 取消')
+
+
         info = answer
-        print(info)
+
+
+
         if info.media_group_id !=None:
             media=await client.get_media_group(chat_id=info.chat.id,message_id=info.message_id)
             print(media)
             for a in media:
-                if a.document == None and a.video == None:
+                if not a.media:
                     await client.send_message(text="发送的不是文件", chat_id=message.chat.id, parse_mode='markdown')
                     await start_down_telegram_file(client, message, file_list)
                     return file_list
                 else:
                     file_list.append(a)
-            await start_down_telegram_file(client, message, file_list)
+            file_list = await start_down_telegram_file(client, message, file_list)
             return file_list
 
 
 
         elif info.text == "/cancel":
             await client.send_message(text="取消发送", chat_id=message.chat.id, parse_mode='markdown')
+
             return []
         elif info.text == "/finish":
             await client.send_message(text=f"接收文件完成,共有{len(file_list)}个文件", chat_id=message.chat.id, parse_mode='markdown')
             return file_list
-        elif info.document == None and info.video == None:
+        elif not info.media:
             await client.send_message(text="发送的不是文件", chat_id=message.chat.id, parse_mode='markdown')
-            await start_down_telegram_file(client, message,file_list)
+            file_list = await start_down_telegram_file(client, message,file_list)
             return file_list
 
         else:
             try:
                 file_list.append(info)
-                temp_file=await start_down_telegram_file(client, message,file_list)
+
+                temp_num = int(info.message_id)+1
+                while True:
+                    try:
+                        temp_info = await client.get_messages(chat_id=message.chat.id,message_ids=temp_num)
+                    except:
+                        break
+                    if not temp_info.media:
+                        break
+                    else:
+                        file_list.append(temp_info)
+                        temp_num = temp_num + 1
+
+                temp_file= await start_down_telegram_file(client, message,file_list)
                 file_list=temp_file
                 return file_list
 
@@ -74,28 +174,32 @@ def progress(current, total, client, message, name):
         None
 
 
-def tgfile_download(client, message, new_message):
-    if new_message.document != None:
-        file_name = new_message.document.file_name
-    elif new_message.video != None:
+async def tgfile_download(client, message, file_list):
 
-        file_name = new_message.video.file_name
-    info = client.send_message(text="开始下载", chat_id=message.chat.id, parse_mode='markdown')
+    for new_message in file_list:
+        if new_message.document != None:
+            file_name = new_message.document.file_name
+        elif new_message.video != None:
 
-    file = client.download_media(message=new_message, progress=progress, progress_args=(client, info, file_name,))
+            file_name = new_message.video.file_name
+        info = await client.send_message(text="开始下载", chat_id=message.chat.id, parse_mode='markdown')
 
-    try:
-        print("开始上传")
-        file_dir = file
-        files_num = 1
-        run_rclone(file_dir, file_name, info=info, file_num=files_num, client=client, message=message,gid=0)
-        os.remove(path=file)
-        return
+        file = await client.download_media(message=new_message, progress=progress, progress_args=(client, info, file_name,))
 
-    except Exception as e:
-        print(e)
-        print("Upload Issue!")
-        return
+        try:
+            print("开始上传")
+            file_dir = file
+            files_num = 1
+            t1 = threading.Thread(target=run_rclone, args=(file_dir, file_name, info, files_num, client, message,0),)
+            t1.start()
+
+            #os.remove(path=file)
+            continue
+
+        except Exception as e:
+            print(e)
+            print("Upload Issue!")
+            return
 
 
 
@@ -111,18 +215,12 @@ async def get_telegram_file(client, message):
         print(temp)
         print(type(temp))
         sys.stdout.flush()
-        
+
         if len(temp)==0:
             return
-        elif len(temp)==1:
-            t1 = threading.Thread(target=tgfile_download, args=(client, message, temp[0]))
-            t1.start()
-            return
-
         else:
-            for a in temp:
-                t1 = threading.Thread(target=tgfile_download, args=(client, message, a))
-                t1.start()
+            await tgfile_download(client, message, temp)
+
             return
     except Exception as e:
         print(f"start_down_telegram_file {e}")