sukuinote/sukuinote/__init__.py

import os
import html
import time
import logging
import asyncio
import traceback
import functools
import yaml
import aiohttp
from datetime import timedelta
from pyrogram import Client, StopPropagation, ContinuePropagation
from pyrogram.types import Chat, User
from pyrogram.parser import parser
from pyrogram.errors.exceptions.bad_request_400 import PeerIdInvalid, ChannelInvalid

logging.basicConfig(level=logging.INFO)
with open('config.yaml') as config:
    config = yaml.safe_load(config)
loop = asyncio.get_event_loop()
help_dict = dict()

apps = []
app_user_ids = dict()
# this code here exists because i can't be fucked
class Parser(parser.Parser):
    async def parse(self, text, mode):
        if mode == 'through':
            return text
        return await super().parse(text, mode)
for session_name in config['config']['sessions']:
    app = Client(session_name, api_id=config['telegram']['api_id'], api_hash=config['telegram']['api_hash'], plugins={'root': os.path.join(__package__, 'plugins')}, parse_mode='html', workdir='sessions')
    app.parser = Parser(app)
    apps.append(app)
slave = Client('sukuinote-slave', api_id=config['telegram']['api_id'], api_hash=config['telegram']['api_hash'], plugins={'root': os.path.join(__package__, 'slave-plugins')}, parse_mode='html', bot_token=config['telegram']['slave_bot_token'], workdir='sessions')
slave.parser = Parser(slave)
session = aiohttp.ClientSession()

async def get_entity(client, entity):
    entity_client = client
    if not isinstance(entity, Chat):
        try:
            entity = int(entity)
        except ValueError:
            pass
        except TypeError:
            entity = entity.id
        try:
            entity = await client.get_chat(entity)
        except (PeerIdInvalid, ChannelInvalid):
            for app in apps:
                if app != client:
                    try:
                        entity = await app.get_chat(entity)
                    except (PeerIdInvalid, ChannelInvalid):
                        pass
                    else:
                        entity_client = app
                        break
            else:
                entity = await slave.get_chat(entity)
                entity_client = slave
    return entity, entity_client

async def get_user(client, entity):
    entity_client = client
    if not isinstance(entity, User):
        try:
            entity = int(entity)
        except ValueError:
            pass
        except TypeError:
            entity = entity.id
        try:
            entity = await client.get_users(entity)
        except PeerIdInvalid:
            for app in apps:
                if app != client:
                    try:
                        entity = await app.get_users(entity)
                    except PeerIdInvalid:
                        pass
                    else:
                        entity_client = app
                        break
            else:
                entity = await slave.get_users(entity)
                entity_client = slave
    return entity, entity_client

def log_errors(func):
    @functools.wraps(func)
    async def wrapper(client, *args):
        try:
            await func(client, *args)
        except (StopPropagation, ContinuePropagation):
            raise
        except Exception:
            tb = traceback.format_exc()
            try:
                await slave.send_message(config['config']['log_chat'], f'Exception occured in {func.__name__}\n\n{tb}', parse_mode=None)
            except Exception:
                logging.exception('Failed to log exception for %s as slave', func.__name__)
                tb = traceback.format_exc()
                for app in apps:
                    try:
                        await app.send_message(config['config']['log_chat'], f'Exception occured in {func.__name__}\n\n{tb}', parse_mode=None)
                    except Exception:
                        logging.exception('Failed to log exception for %s as app', func.__name__)
                        tb = traceback.format_exc()
                    else:
                        break
                raise
            raise
    return wrapper

def public_log_errors(func):
    @functools.wraps(func)
    async def wrapper(client, message):
        try:
            await func(client, message)
        except (StopPropagation, ContinuePropagation):
            raise
        except Exception:
            await message.reply_text(traceback.format_exc(), parse_mode=None)
            raise
    return wrapper

# https://stackoverflow.com/a/49361727
def format_bytes(size):
    size = int(size)
    # 2**10 = 1024
    power = 1024
    n = 0
    power_labels = {0 : '', 1: 'K', 2: 'M', 3: 'G', 4: 'T'}
    while size > power:
        size /= power
        n += 1
    return f"{size:.2f} {power_labels[n]+'B'}"

# https://stackoverflow.com/a/34325723
def return_progress_string(current, total):
    filled_length = int(30 * current // total)
    return '[' + '=' * filled_length + ' ' * (30 - filled_length) + ']'

# https://stackoverflow.com/a/852718
# https://stackoverflow.com/a/775095
def calculate_eta(current, total, start_time):
    if not current:
        return '00:00:00'
    end_time = time.time()
    elapsed_time = end_time - start_time
    seconds = (elapsed_time * (total / current)) - elapsed_time
    thing = ''.join(str(timedelta(seconds=seconds)).split('.')[:-1]).split(', ')
    thing[-1] = thing[-1].rjust(8, '0')
    return ', '.join(thing)

progress_callback_data = dict()
async def progress_callback(current, total, reply, text, upload):
    message_identifier = (reply.chat.id, reply.message_id)
    last_edit_time, prevtext, start_time = progress_callback_data.get(message_identifier, (0, None, time.time()))
    if current == total:
        try:
            progress_callback_data.pop(message_identifier)
        except KeyError:
            pass
    elif (time.time() - last_edit_time) > 1:
        handle = 'Upload' if upload else 'Download'
        if last_edit_time:
            speed = format_bytes((total - current) / (time.time() - start_time))
        else:
            speed = '0 B'
        text = f'''{text}
<code>{return_progress_string(current, total)}</code>

<b>Total Size:</b> {format_bytes(total)}
<b>{handle}ed Size:</b> {format_bytes(current)}
<b>{handle} Speed:</b> {speed}/s
<b>ETA:</b> {calculate_eta(current, total, start_time)}'''
        if prevtext != text:
            await reply.edit_text(text)
            prevtext = text
            last_edit_time = time.time()
            progress_callback_data[message_identifier] = last_edit_time, prevtext, start_time