From b1b50fca188fbc24f9e464c928114e7704dd0287 Mon Sep 17 00:00:00 2001 From: Egor3f Date: Wed, 3 Jun 2020 16:29:05 +0300 Subject: [PATCH] First commit --- .gitignore | 6 + setup.py | 14 ++ telecup_cli.py | 344 +++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 364 insertions(+) create mode 100644 .gitignore create mode 100644 setup.py create mode 100644 telecup_cli.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..b5de370 --- /dev/null +++ b/.gitignore @@ -0,0 +1,6 @@ +.idea +venv +*.egg-info +*.session +*.cfg +*.pyc diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..a5201ee --- /dev/null +++ b/setup.py @@ -0,0 +1,14 @@ +from setuptools import setup + +setup( + name='telecup', + version='0.1', + py_modules=['telecup'], + install_requires=[ + 'Click', + 'telethon' + ], + entry_points={ + 'console_scripts': 'telecup=telecup_cli:cli' + } +) diff --git a/telecup_cli.py b/telecup_cli.py new file mode 100644 index 0000000..c8e3f17 --- /dev/null +++ b/telecup_cli.py @@ -0,0 +1,344 @@ +import datetime +import hashlib +import math +import os +import sys +import textwrap +import time +import typing +from collections import namedtuple + +import click +from telethon.helpers import generate_random_long +from telethon.sync import TelegramClient, connection +from telethon.tl.custom.message import Message +from telethon.tl.functions.upload import SaveBigFilePartRequest +from telethon.tl.functions.upload import GetFileRequest +from telethon.tl.types import InputFileBig +from telethon.tl.types.upload import File +from telethon.utils import get_input_location + +MESSAGE_BLOCK_START = '_cup_start_' +MESSAGE_HEADER = f'TeleCup File Uploader https://lnurl.ru/telecup' +MESSAGE_BLOCK_END = '_cup_end_' + +UploadInfo = namedtuple('UploadInfo', ['file_list', 'real_size']) +DownloadInfo = namedtuple('DownloadInfo', ['message', 'part_info']) + + +class Config: + def __init__(self): + self.session_file = '.telegram-session' + self.dialog = 'me' + self.client: typing.Optional[TelegramClient] = None + + +pass_config = click.make_pass_decorator(Config, ensure=True) + + +def parse_file_size(humanString: str): + humanString = humanString.strip() + units = {"k": pow(2, 10), "m": pow(2, 20), "g": pow(2, 30)} + for unit in units.keys(): + if humanString.lower().endswith(unit): + return int(float(humanString[:-1]) * units[unit]) + return int(humanString) + + +def format_file_size(bytesCount: int): + for unit in ['', 'K', 'M', 'G']: + if bytesCount < pow(2, 10): + return '%3.2f%s' % (bytesCount, unit) + bytesCount /= pow(2, 10) + return '%.2f%s' % (bytesCount, 'T') + + +def parse_message(caption: str): + if MESSAGE_BLOCK_START not in caption or MESSAGE_BLOCK_END not in caption: + return None + blockStartPos = caption.index(MESSAGE_BLOCK_START) + blockEndPos = caption.index(MESSAGE_BLOCK_END) + block = caption[blockStartPos + len(MESSAGE_BLOCK_START):blockEndPos] + fileInfo = dict() + for line in block.split('\n'): + if len(line) > 0 and line[0] == '^': + name, value = line[1:].split('=') + fileInfo[name.strip()] = value.strip() + return fileInfo + + +def retrieve_app_hash(app_config): + app_id = None + app_hash = None + proxy = None + config_path = os.path.realpath(app_config) + if os.path.isfile(config_path): + with open(config_path, 'r') as f: + values = {line.split('=')[0].strip(): line.split('=')[1].strip() for line in f.readlines() if len(line.strip()) > 0} + if 'APP_ID' in values and 'APP_HASH' in values: + app_id = values['APP_ID'] + app_hash = values['APP_HASH'] + if 'PROXY_HOST' in values and 'PROXY_PORT' in values and 'PROXY_SECRET' in values: + proxy = (values['PROXY_HOST'], int(values['PROXY_PORT']), values['PROXY_SECRET']) + if not (app_id and app_hash): + click.echo('You have to create Telegram App configuration before using this tool') + click.echo('Please, visit https://my.telegram.org/apps page and create an app') + click.pause() + app_id = click.prompt('App id', type=click.INT) + app_hash = click.prompt('App hash').strip() + use_proxy = click.confirm('Use MTProto proxy? Required in some countries like Russia and Iran', default=False) + if use_proxy: + proxy_host = click.prompt('Proxy host') + proxy_port = click.prompt('Proxy port', type=click.INT) + proxy_secret = click.prompt('Proxy secret') + proxy = (proxy_host, proxy_port, proxy_secret) + with open(config_path, 'w') as f: + f.write(f'APP_ID={app_id}\nAPP_HASH={app_hash}\n') + if use_proxy: + f.write(f'PROXY_HOST={proxy_host}\nPROXY_PORT={proxy_port}\nPROXY_SECRET={proxy_secret}\n') + return app_hash, app_id, proxy + + +def check_logged_in(config: Config): + config.client.connect() + if not config.client.is_user_authorized(): + raise click.ClickException('You are not authorized. Please log in first') + + +def upload_file(client: TelegramClient, source: typing.BinaryIO, expectedStreamSize: int, nameHash: str, dialog: str, maxFileSize): + fileId = generate_random_long() + chunkSize = pow(2, 19) + maxChunksInFile = maxFileSize / chunkSize + fillByte = 0x00 + + inputFiles = [] + currentFileChunkPos = 0 + currentFileTotalChunks = math.ceil(expectedStreamSize / chunkSize) \ + if expectedStreamSize < maxFileSize \ + else math.ceil(maxFileSize / chunkSize) + currentFileBytesWritten = 0 + realBytesWritten = 0 + + while True: + buffer = source.read(chunkSize) + bufLen = len(buffer) + if not bufLen: + break + if realBytesWritten + bufLen > expectedStreamSize: + raise click.ClickException('Stream is larger than expected file size') + if bufLen < chunkSize: + buffer = buffer + bytearray([fillByte] * (chunkSize - bufLen)) + client(SaveBigFilePartRequest( + fileId, currentFileChunkPos, currentFileTotalChunks, buffer + )) + currentFileChunkPos += 1 + currentFileBytesWritten += bufLen + realBytesWritten += bufLen + if currentFileChunkPos % 10 == 0: + click.echo(f'{format_file_size(realBytesWritten)} bytes sent') + + if bufLen < chunkSize: + break + + if currentFileChunkPos == maxChunksInFile: + readyFile = InputFileBig( + fileId, currentFileTotalChunks, f'{int(time.time() * 1000)}.bin' + ) + client.send_file(dialog, readyFile, caption=textwrap.dedent(f''' + {MESSAGE_BLOCK_START} + {MESSAGE_HEADER} + #telecup_part #telecup_part_{nameHash} + ^name_hash={nameHash} + ^part={len(inputFiles) + 1} + ^real_size={currentFileBytesWritten} + {MESSAGE_BLOCK_END} + ''')) + inputFiles.append(readyFile) + + # Start a new file + fileId = generate_random_long() + currentFileTotalChunks = math.ceil((expectedStreamSize - realBytesWritten) / chunkSize) \ + if (expectedStreamSize - realBytesWritten) < maxFileSize \ + else math.ceil(maxFileSize / chunkSize) + currentFileChunkPos = 0 + currentFileBytesWritten = 0 + + if currentFileBytesWritten > 0: + if currentFileChunkPos < currentFileTotalChunks: + buffer = bytes(bytearray([fillByte] * chunkSize)) + for newChunkPos in range(currentFileChunkPos, currentFileTotalChunks): + client(SaveBigFilePartRequest( + fileId, newChunkPos, currentFileTotalChunks, buffer + )) + + readyFile = InputFileBig( + fileId, currentFileTotalChunks, f'{int(time.time() * 1000)}.bin' + ) + client.send_file(dialog, readyFile, caption=textwrap.dedent(f''' + {MESSAGE_BLOCK_START} + {MESSAGE_HEADER} + #telecup_part #telecup_part_{nameHash} + ^name_hash={nameHash} + ^part={len(inputFiles) + 1} + ^real_size={currentFileBytesWritten} + {MESSAGE_BLOCK_END} + ''')) + inputFiles.append(readyFile) + return UploadInfo(inputFiles, realBytesWritten) + + +def download_part(client: TelegramClient, dest: typing.BinaryIO, dInfo: DownloadInfo): + dcId, inputFileLocation = get_input_location(dInfo.message) + chunkSize = pow(2, 20) + realSize = int(dInfo.part_info['real_size']) + totalBytesDownloaded = 0 + while totalBytesDownloaded < realSize: + offset = totalBytesDownloaded + limit = chunkSize if offset + chunkSize < realSize else realSize - offset + if limit % pow(2, 12) != 0: + extraBytes = pow(2, 12) - limit % pow(2, 12) + limit += extraBytes + + downloadResult: File = client(GetFileRequest( + inputFileLocation, + offset, + limit, + precise=False, + cdn_supported=False + )) + buffer = downloadResult.bytes + bufLen = len(buffer) + totalBytesDownloaded += bufLen + + if totalBytesDownloaded > realSize: + extraBytes = totalBytesDownloaded - realSize + buffer = buffer[:-extraBytes] + + dest.write(buffer) + click.echo(f"Part {dInfo.part_info['part']}: {format_file_size(totalBytesDownloaded)} ready", err=True) + + +def download_file(client: TelegramClient, dest: typing.BinaryIO, fileInfo: dict, dialog: str): + partMessagesSearchResults = client.iter_messages(dialog, search=f"#telecup_part_{fileInfo['name_hash']}") + partMessages: typing.List[typing.Optional[DownloadInfo]] = [None] * int(fileInfo['total_parts']) + for msg in partMessagesSearchResults: + partInfo = parse_message(msg.message) + if not partInfo or partInfo['name_hash'] != fileInfo['name_hash']: + continue + partMessages[int(partInfo['part']) - 1] = DownloadInfo(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'])) + for part in partMessages: + download_part(client, dest, part) + click.echo(f"Part {part.part_info['part']} ready", err=True) + + +@click.group() +@click.option('--session-file', default='.telegram-session', show_default=True, + help='File where telegram session information will be stored (unique for each account)') +@click.option('--app-config', default='.telecup.cfg', show_default=True, help='App configuration file') +@click.option('--dialog', default='me', help='Name of the conversation, where files will be stored [default: Saved Messages]') +@pass_config +def cli(config: Config, session_file, app_config, dialog): + config.session_file = session_file + config.dialog = dialog + + APP_HASH, APP_ID, proxy = retrieve_app_hash(app_config) + + if not proxy: + client = TelegramClient(session_file, APP_ID, APP_HASH) + else: + client = TelegramClient(session_file, APP_ID, APP_HASH, + connection=connection.ConnectionTcpMTProxyAbridged, + proxy=proxy) + config.client = client + + +@cli.command() +@pass_config +def login(config: Config): + """ + Log into your telegram account interactively and save the login information to session file + """ + config.client.start() + click.echo(f'You are logged in as {config.client.get_me().first_name}') + click.echo('To switch user: either pass another session-file or remove existing') + + +@cli.command(name='list') +@pass_config +def list_files(config: Config): + """ + List all files, uploaded to your account by TeleCup + """ + check_logged_in(config) + fileMessages: typing.Iterator[Message] = config.client.iter_messages(config.dialog, search='#telecup_file') + for msg in fileMessages: + fileInfo = parse_message(msg.message) + if not fileInfo: + continue + timestamp = int(fileInfo['created_at']) + createdAtFormatted = datetime.datetime.utcfromtimestamp(timestamp).strftime('%d.%m.%Y %H:%M') + click.echo(f"{fileInfo['name']}\t{format_file_size(int(fileInfo['real_size']))}\t{createdAtFormatted}") + + +@cli.command() +@click.option('--part-size', default='1g', show_default=True, help='Maximum size of a single file') +@click.argument('filename') +@click.argument('stream_size') +@pass_config +def upload(config: Config, part_size, filename, stream_size): + """ + Upload a new file to your account + """ + source = sys.stdin.buffer + estimatedBytes = parse_file_size(stream_size) + partSize = parse_file_size(part_size) + if partSize > pow(2, 20) * 1536: + raise click.ClickException('Part size must be less than 1.5 Gib') + check_logged_in(config) + + nameHash = hashlib.sha256(filename.encode('utf-8')).hexdigest() + existingMessages = config.client.get_messages(config.dialog, search=f'#telecup_file_{nameHash}') + if len(existingMessages) > 0: + raise click.ClickException(f'File with name {filename} already exists in dialog {config.dialog}') + + uploadInfo = upload_file(config.client, source, estimatedBytes, nameHash, config.dialog, partSize) + if uploadInfo.real_size == 0: + raise click.ClickException('Input stream is empty') + config.client.send_message(config.dialog, textwrap.dedent(f''' + {MESSAGE_BLOCK_START} + {MESSAGE_HEADER} + #telecup_file #telecup_file_{nameHash} + ^name={filename} + ^name_hash={nameHash} + ^real_size={uploadInfo.real_size} + ^total_parts={len(uploadInfo.file_list)} + ^created_at={int(time.time())} + {MESSAGE_BLOCK_END} + ''')) + click.echo('OK') + config.client.disconnect() + + +@cli.command() +@click.argument('filename') +@pass_config +def download(config: Config, filename): + """ + Download an existing file from your account + """ + destination = sys.stdout.buffer + check_logged_in(config) + + fileMessages: typing.Iterator[Message] = config.client.iter_messages(config.dialog, search='#telecup_file') + for msg in fileMessages: + fileInfo = parse_message(msg.message) + if fileInfo and ( + fileInfo['name'].startswith(filename.strip()) or + fileInfo['name_hash'].startswith(filename.strip()) + ): + download_file(config.client, destination, fileInfo, config.dialog) + click.echo('OK', err=True) + config.client.disconnect()