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_' MESSAGE_BLOCK_END = '_cup_end_'
UploadInfo = namedtuple('UploadInfo', ['file_list', 'real_size']) 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: class Config:
@ -239,7 +239,7 @@ async def download_part(client: TelegramClient, dest: typing.BinaryIO, dInfo: Do
except BadRequestError as e: except BadRequestError as e:
if 'expire' in str(e.message).lower(): if 'expire' in str(e.message).lower():
click.echo('Reloading message...', err=True) 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) dcId, inputFileLocation = get_input_location(newMessage)
continue continue
buffer = downloadResult.bytes buffer = downloadResult.bytes
@ -265,7 +265,7 @@ async def download_file(client: TelegramClient, dest: typing.BinaryIO, fileInfo:
partInfo = parse_message(msg.message) partInfo = parse_message(msg.message)
if not partInfo or partInfo['part_id'] != fileInfo['part_id']: if not partInfo or partInfo['part_id'] != fileInfo['part_id']:
continue 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): if any(part is None for part in partMessages):
raise click.ClickException('Missing some parts') raise click.ClickException('Missing some parts')
partMessages.sort(key=lambda dInfo: int(dInfo.part_info['part'])) partMessages.sort(key=lambda dInfo: int(dInfo.part_info['part']))
@ -377,13 +377,16 @@ async def download(config: Config, filename):
destination = sys.stdout.buffer destination = sys.stdout.buffer
await check_logged_in(config) await check_logged_in(config)
fileMessages = config.client.iter_messages(config.dialog, search='#telecup_file') filename = filename.strip()
async for msg in fileMessages: nameHash = hashlib.sha256(filename.encode('utf-8')).hexdigest()
fileInfo = parse_message(msg.message) fileMessages = await config.client.get_messages(config.dialog, search=f'#telecup_file_{nameHash}')
if fileInfo and ( if len(fileMessages) == 0:
fileInfo['name'] == filename.strip() or click.echo(f'`{filename}` not found', err=True)
fileInfo['name_hash'].startswith(filename.strip()) msg = fileMessages[0]
): fileInfo = parse_message(msg.message)
await download_file(config.client, destination, fileInfo, config.dialog) if fileInfo and fileInfo['name'] == filename:
click.echo('OK', err=True) 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() await config.client.disconnect()