First commit

.gitignore

@@ -0,0 +1,6 @@
+.idea
+venv
+*.egg-info
+*.session
+*.cfg
+*.pyc

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'
+    }
+)

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()