Small fixes

This commit is contained in:
Egor Aristov 2020-06-25 09:03:08 +03:00
parent 0d299fe84f
commit 9fab87b045

View File

@ -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:
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.strip() or
fileInfo['name_hash'].startswith(filename.strip())
):
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()