From 9fab87b0451e2dd6b45971c16ff3b22b1200dab7 Mon Sep 17 00:00:00 2001
From: Egor3f <ef@efprojects.com>
Date: Thu, 25 Jun 2020 09:03:08 +0300
Subject: [PATCH] Small fixes

---
 telecup_cli.py | 27 +++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/telecup_cli.py b/telecup_cli.py
index 977ed27..7546a6f 100644
--- a/telecup_cli.py
+++ b/telecup_cli.py
@@ -28,7 +28,7 @@ MESSAGE_HEADER = f'TeleCup File Uploader https://lnurl.ru/telecup'
 MESSAGE_BLOCK_END = '_cup_end_'
 
 UploadInfo = namedtuple('UploadInfo', ['file_list', 'real_size'])
-DownloadInfo = typing.NamedTuple('DownloadInfo', [('message', Message), ('part_info', dict)])
+DownloadInfo = typing.NamedTuple('DownloadInfo', [('dialog', str), ('message', Message), ('part_info', dict)])
 
 
 class Config:
@@ -239,7 +239,7 @@ async def download_part(client: TelegramClient, dest: typing.BinaryIO, dInfo: Do
         except BadRequestError as e:
             if 'expire' in str(e.message).lower():
                 click.echo('Reloading message...', err=True)
-                newMessage = await client.get_messages(ids=dInfo.message.id)
+                newMessage = await client.get_messages(dInfo.dialog, ids=dInfo.message.id)
                 dcId, inputFileLocation = get_input_location(newMessage)
                 continue
         buffer = downloadResult.bytes
@@ -265,7 +265,7 @@ async def download_file(client: TelegramClient, dest: typing.BinaryIO, fileInfo:
         partInfo = parse_message(msg.message)
         if not partInfo or partInfo['part_id'] != fileInfo['part_id']:
             continue
-        partMessages[int(partInfo['part']) - 1] = DownloadInfo(msg, partInfo)
+        partMessages[int(partInfo['part']) - 1] = DownloadInfo(dialog, msg, partInfo)
     if any(part is None for part in partMessages):
         raise click.ClickException('Missing some parts')
     partMessages.sort(key=lambda dInfo: int(dInfo.part_info['part']))
@@ -377,13 +377,16 @@ async def download(config: Config, filename):
     destination = sys.stdout.buffer
     await check_logged_in(config)
 
-    fileMessages = config.client.iter_messages(config.dialog, search='#telecup_file')
-    async for msg in fileMessages:
-        fileInfo = parse_message(msg.message)
-        if fileInfo and (
-                fileInfo['name'] == filename.strip() or
-                fileInfo['name_hash'].startswith(filename.strip())
-        ):
-            await download_file(config.client, destination, fileInfo, config.dialog)
-            click.echo('OK', err=True)
+    filename = filename.strip()
+    nameHash = hashlib.sha256(filename.encode('utf-8')).hexdigest()
+    fileMessages = await config.client.get_messages(config.dialog, search=f'#telecup_file_{nameHash}')
+    if len(fileMessages) == 0:
+        click.echo(f'`{filename}` not found', err=True)
+    msg = fileMessages[0]
+    fileInfo = parse_message(msg.message)
+    if fileInfo and fileInfo['name'] == filename:
+        await download_file(config.client, destination, fileInfo, config.dialog)
+        click.echo('OK', err=True)
+    else:
+        click.echo('Message FileInfo corrupt', err=True)
     await config.client.disconnect()