diff --git a/bot/__init__.py b/bot/__init__.py index aa3c9c0..59ef952 100644 --- a/bot/__init__.py +++ b/bot/__init__.py @@ -1,19 +1,15 @@ -from pyrogram import Client +from telethon import TelegramClient from logging import getLogger from logging.config import dictConfig from .config import Telegram, LOGGER_CONFIG_JSON dictConfig(LOGGER_CONFIG_JSON) -version = 1.4 +version = 1.5 logger = getLogger('bot') -TelegramBot = Client( - name="bot", - api_id = Telegram.API_ID, - api_hash = Telegram.API_HASH, - bot_token = Telegram.BOT_TOKEN, - plugins={"root": "bot/plugins"}, - workers = Telegram.BOT_WORKERS, - max_concurrent_transmissions=1000 +TelegramBot = TelegramClient( + session='bot', + api_id=Telegram.API_ID, + api_hash=Telegram.API_HASH ) diff --git a/bot/__main__.py b/bot/__main__.py index da32186..7c8d32e 100644 --- a/bot/__main__.py +++ b/bot/__main__.py @@ -1,7 +1,22 @@ +from importlib import import_module +from pathlib import Path from bot import TelegramBot, logger +from bot.config import Telegram from bot.server import server +def load_plugins(): + count = 0 + for path in Path('bot/plugins').rglob('*.py'): + import_module(f'bot.plugins.{path.stem}') + count += 1 + logger.info(f'Loaded {count} {"plugins" if count > 1 else "plugin"}.') + if __name__ == '__main__': - logger.info('Initializing...') + logger.info('initializing...') TelegramBot.loop.create_task(server.serve()) - TelegramBot.run() + TelegramBot.start(bot_token=Telegram.BOT_TOKEN) + logger.info('Telegram client is now started.') + logger.info('Loading bot plugins...') + load_plugins() + logger.info('Bot is now ready!') + TelegramBot.run_until_disconnected() diff --git a/bot/config.py b/bot/config.py index d2ce2dc..49ffd6e 100644 --- a/bot/config.py +++ b/bot/config.py @@ -8,7 +8,6 @@ class Telegram: BOT_USERNAME = env.get("TELEGRAM_BOT_USERNAME", "BotFather") BOT_TOKEN = env.get("TELEGRAM_BOT_TOKEN", "1234:abcd") CHANNEL_ID = int(env.get("TELEGRAM_CHANNEL_ID", -1001234567890)) - BOT_WORKERS = int(env.get("BOT_WORKERS", 10)) SECRET_CODE_LENGTH = int(env.get("SECRET_CODE_LENGTH", 12)) class Server: diff --git a/bot/modules/decorators.py b/bot/modules/decorators.py index 3c24daa..cd8d7b9 100644 --- a/bot/modules/decorators.py +++ b/bot/modules/decorators.py @@ -1,16 +1,22 @@ -from pyrogram import Client -from pyrogram.types import Message, CallbackQuery -from typing import Union, Callable +from telethon.events import NewMessage, CallbackQuery +from telethon.tl.custom import Message +from typing import Callable from functools import wraps from bot.config import Telegram +from bot.modules.static import * -def verify_user(func: Callable): +def verify_user(private: bool = False): - @wraps(func) - async def decorator(client: Client, update: Union[Message, CallbackQuery]): - chat_id = str(update.from_user.id if update.from_user else update.chat.id) + def decorator(func: Callable): + @wraps(func) + async def wrapper(update: NewMessage.Event | CallbackQuery.Event): + if private and not update.is_private: + return - if not Telegram.ALLOWED_USER_IDS or chat_id in Telegram.ALLOWED_USER_IDS: - return await func(client, update) - + chat_id = str(update.chat_id) + + if not Telegram.ALLOWED_USER_IDS or chat_id in Telegram.ALLOWED_USER_IDS: + return await func(update) + + return wrapper return decorator diff --git a/bot/modules/static.py b/bot/modules/static.py index f3347e6..dfd9865 100644 --- a/bot/modules/static.py +++ b/bot/modules/static.py @@ -9,6 +9,21 @@ Add me to your channel to instantly generate links for any downloadable media. O - /log to get bot logs. (admin only!) """ +UserInfoText = \ +""" +**First Name:** +`{sender.first_name}` + +**Last Name:** +`{sender.last_name}` + +**User ID:** +`{sender.id}` + +**Username:** +`@{sender.username}` +""" + FileLinksText = \ """ **Download Link:** @@ -45,4 +60,9 @@ The link has been revoked. It may take some time for the changes to take effect. InvalidPayloadText = \ """ Invalid payload. -""" \ No newline at end of file +""" + +MediaTypeNotSupportedText = \ +""" +Sorry, this media type is not supported. +""" diff --git a/bot/modules/telegram.py b/bot/modules/telegram.py index 7b500ec..452fb6e 100644 --- a/bot/modules/telegram.py +++ b/bot/modules/telegram.py @@ -1,56 +1,68 @@ -from pyrogram.types import Message +from telethon.events import NewMessage +from telethon.tl.custom import Message from datetime import datetime from mimetypes import guess_type from bot import TelegramBot from bot.config import Telegram from bot.server.error import abort -async def get_message(message_id: int): +async def get_message(message_id: int) -> Message | None: message = None try: - message = await TelegramBot.get_messages( - chat_id=Telegram.CHANNEL_ID, - message_ids=message_id - ) - if message.empty: message = None + message = await TelegramBot.get_messages(Telegram.CHANNEL_ID, ids=message_id) except Exception: pass return message -async def get_file_properties(msg: Message): - attributes = ( - 'document', - 'video', - 'audio', - 'voice', - 'photo', - 'video_note' +async def send_message(message:Message, send_to:int = Telegram.CHANNEL_ID) -> Message: + message.forward + return await TelegramBot.send_message( + entity=send_to, + message=message ) - for attribute in attributes: - media = getattr(msg, attribute, None) - if media: - file_type = attribute - break - - if not media: abort(400, 'Unknown file type.') +def filter_files(update: NewMessage.Event | Message): + return bool( + ( + update.document + or update.photo + or update.video + or update.video_note + or update.audio + or update.gif + ) + and not update.sticker + ) - file_name = getattr(media, 'file_name', None) - file_size = getattr(media, 'file_size', 0) +def get_file_properties(message: Message): + file_name = message.file.name + file_size = message.file.size or 0 + mime_type = message.file.mime_type if not file_name: - file_format = { + attributes = { 'video': 'mp4', 'audio': 'mp3', 'voice': 'ogg', 'photo': 'jpg', 'video_note': 'mp4' - }.get(attribute) + } + + for attribute in attributes: + media = getattr(message, attribute, None) + if media: + file_type, file_format = attribute, attributes[attribute] + break + + if not media: + abort(400, 'Invalid media type.') + date = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") file_name = f'{file_type}-{date}.{file_format}' - - mime_type = guess_type(file_name)[0] or 'application/octet-stream' - + + if not mime_type: + mime_type = guess_type(file_name)[0] or 'application/octet-stream' + return file_name, file_size, mime_type diff --git a/bot/plugins/callback.py b/bot/plugins/callback.py index 007ec9b..3ebf08c 100644 --- a/bot/plugins/callback.py +++ b/bot/plugins/callback.py @@ -1,27 +1,22 @@ -from pyrogram.types import CallbackQuery +from telethon.events import CallbackQuery from bot import TelegramBot -from bot.modules.static import * from bot.modules.decorators import verify_user +from bot.modules.static import * from bot.modules.telegram import get_message -@TelegramBot.on_callback_query() -@verify_user -async def manage_callback(bot, q: CallbackQuery): - query = q.data - if query.startswith('rm_'): - sq = query.split('_') +@TelegramBot.on(CallbackQuery(pattern=r'^rm_')) +@verify_user(private=True) +async def delete_file(event: CallbackQuery.Event): + query_data = event.query.data.decode().split('_') - if len(sq) != 3: - return await q.answer(InvalidQueryText, show_alert=True) - - message = await get_message(int(sq[1])) - - if not message: - return await q.answer(MessageNotExist, show_alert=True) - if sq[2] != message.caption: - return await q.answer(InvalidQueryText, show_alert=True) + if len(query_data) != 3: + return await event.answer(InvalidQueryText, alert=True) - await message.delete() - await q.answer(LinkRevokedText, show_alert=True) - else: - await q.answer(InvalidQueryText, show_alert=True) \ No newline at end of file + message = await get_message(int(query_data[1])) + + if not message: + return await event.answer(MessageNotExist, alert=True) + + await message.delete() + + return await event.answer(LinkRevokedText, alert=True) diff --git a/bot/plugins/commands.py b/bot/plugins/commands.py index 0eaf037..d00e648 100644 --- a/bot/plugins/commands.py +++ b/bot/plugins/commands.py @@ -1,45 +1,28 @@ -from pyrogram import filters -from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton -from aiofiles import open as async_open -from aiofiles.os import remove as async_rm -from bot import TelegramBot, logger +from telethon import Button +from telethon.events import NewMessage +from telethon.tl.custom.message import Message +from bot import TelegramBot from bot.config import Telegram from bot.modules.static import * -from .deeplinks import deeplinks from bot.modules.decorators import verify_user -@TelegramBot.on_message(filters.command('start') & filters.private) -@verify_user -async def start(_, msg: Message): - if len(msg.command) != 1: - return await deeplinks(msg, msg.command[1]) - - await msg.reply( - text=WelcomeText % {'first_name': msg.from_user.first_name}, - quote=True, - reply_markup=InlineKeyboardMarkup( +@TelegramBot.on(NewMessage(incoming=True, pattern=r'^/start$')) +@verify_user(private=True) +async def welcome(event: NewMessage.Event | Message): + await event.reply( + message=WelcomeText % {'first_name': event.sender.first_name}, + buttons=[ [ - [ - InlineKeyboardButton('Add to Channel', url=f'https://t.me/{Telegram.BOT_USERNAME}?startchannel&admin=post_messages+edit_messages+delete_messages') - ] + Button.url('Add to Channel', f'https://t.me/{Telegram.BOT_USERNAME}?startchannel&admin=post_messages+edit_messages+delete_messages') ] - ) + ] ) -@TelegramBot.on_message(filters.command('info') & filters.private) -@verify_user -async def user_info(_, msg: Message): - await msg.reply(text=f'`{msg.from_user}`', quote=True) +@TelegramBot.on(NewMessage(incoming=True, pattern=r'^/info$')) +@verify_user(private=True) +async def user_info(event: Message): + await event.reply(UserInfoText.format(sender=event.sender)) - filename = f'{msg.from_user.id}.json' - async with async_open(filename, "w") as file: - await file.write(f'{msg.from_user}') - - await msg.reply_document(filename) - await async_rm(filename) - -@TelegramBot.on_message(filters.private & filters.command('log') & filters.user(Telegram.OWNER_ID)) -async def send_log(_, msg: Message): - await msg.reply_document('event-log.txt', quote=True) - -logger.info('Bot is now started!') \ No newline at end of file +@TelegramBot.on(NewMessage(chats=Telegram.OWNER_ID, incoming=True, pattern=r'^/log$')) +async def send_log(event: NewMessage.Event | Message): + await event.reply(file='event-log.txt') diff --git a/bot/plugins/deeplinks.py b/bot/plugins/deeplinks.py index bdfd3e7..0e22b14 100644 --- a/bot/plugins/deeplinks.py +++ b/bot/plugins/deeplinks.py @@ -1,21 +1,22 @@ -from pyrogram.types import Message +from telethon.events import NewMessage +from telethon.tl.custom import Message +from bot import TelegramBot +from bot.modules.decorators import verify_user +from bot.modules.telegram import get_message, send_message from bot.modules.static import * -from bot.modules.telegram import get_message -async def deeplinks(msg: Message, payload: str): - if payload.startswith('file_'): - sp = payload.split('_') +@TelegramBot.on(NewMessage(incoming=True, pattern=r'^/start file_')) +@verify_user(private=True) +async def send_file(event: NewMessage.Event | Message): + payload = event.raw_text.split()[-1].split('_') - if len(sp) != 3: - return await msg.reply(InvalidPayloadText, quote=True) - - message = await get_message(int(sp[1])) + if len(payload) != 3: + return await event.reply(InvalidPayloadText) + + message = await get_message(int(payload[1])) - if not message: - return await msg.reply(MessageNotExist) - if sp[2] != message.caption: - return await msg.reply(InvalidPayloadText, quote=True) - - await message.copy(chat_id=msg.from_user.id, caption="") - else: - await msg.reply(InvalidPayloadText, quote=True) \ No newline at end of file + if not message: + return await event.reply(MessageNotExist) + + message.raw_text = '' + await send_message(message, send_to=event.chat_id) diff --git a/bot/plugins/files.py b/bot/plugins/files.py index 6cb8aa4..afa5e96 100644 --- a/bot/plugins/files.py +++ b/bot/plugins/files.py @@ -1,126 +1,92 @@ -from pyrogram import filters, errors -from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton +from telethon import Button +from telethon.events import NewMessage +from telethon.errors import MessageAuthorRequiredError, MessageNotModifiedError, MessageIdInvalidError +from telethon.tl.custom import Message from secrets import token_hex from bot import TelegramBot from bot.config import Telegram, Server from bot.modules.decorators import verify_user +from bot.modules.telegram import send_message, filter_files from bot.modules.static import * -@TelegramBot.on_message( - filters.private - & ( - filters.document - | filters.video - | filters.video_note - | filters.audio - | filters.voice - | filters.photo - ) -) -@verify_user -async def handle_user_file(_, msg: Message): +@TelegramBot.on(NewMessage(incoming=True, func=filter_files)) +@verify_user(private=True) +async def user_file_handler(event: NewMessage.Event | Message): secret_code = token_hex(Telegram.SECRET_CODE_LENGTH) - file = await msg.copy( - chat_id=Telegram.CHANNEL_ID, - caption=f'`{secret_code}`' - ) - file_id = file.id + event.message.text = f'`{secret_code}`' + message = await send_message(event.message) + message_id = message.id - dl_link = f'{Server.BASE_URL}/dl/{file_id}?code={secret_code}' - tg_link = f'{Server.BASE_URL}/file/{file_id}?code={secret_code}' - deep_link = f'https://t.me/{Telegram.BOT_USERNAME}?start=file_{file_id}_{secret_code}' + dl_link = f'{Server.BASE_URL}/dl/{message_id}?code={secret_code}' + tg_link = f'{Server.BASE_URL}/file/{message_id}?code={secret_code}' + deep_link = f'https://t.me/{Telegram.BOT_USERNAME}?start=file_{message_id}_{secret_code}' - if (msg.document and 'video' in msg.document.mime_type) or msg.video: - stream_link = f'{Server.BASE_URL}/stream/{file_id}?code={secret_code}' - await msg.reply( - text=MediaLinksText % {'dl_link': dl_link, 'tg_link': tg_link, 'stream_link': stream_link}, - quote=True, - reply_markup=InlineKeyboardMarkup( + if (event.document and 'video' in event.document.mime_type) or event.video: + stream_link = f'{Server.BASE_URL}/stream/{message_id}?code={secret_code}' + await event.reply( + message= MediaLinksText % {'dl_link': dl_link, 'tg_link': tg_link, 'tg_link': tg_link, 'stream_link': stream_link}, + buttons=[ [ - [ - InlineKeyboardButton('Download', url=dl_link), - InlineKeyboardButton('Stream', url=stream_link) - ], - [ - InlineKeyboardButton('Get File', url=deep_link), - InlineKeyboardButton('Revoke', callback_data=f'rm_{file_id}_{secret_code}') - ] + Button.url('Download', dl_link), + Button.url('Stream', stream_link) + ], + [ + Button.url('Get File', deep_link), + Button.inline('Revoke', f'rm_{message_id}_{secret_code}') ] - ) + ] ) else: - await msg.reply( - text=FileLinksText % {'dl_link': dl_link, 'tg_link': tg_link}, - quote=True, - reply_markup=InlineKeyboardMarkup( + await event.reply( + message=FileLinksText % {'dl_link': dl_link, 'tg_link': tg_link}, + buttons=[ [ - [ - InlineKeyboardButton('Download', url=dl_link), - InlineKeyboardButton('Get File', url=deep_link) - ], - [ - InlineKeyboardButton('Revoke', callback_data=f'rm_{file_id}_{secret_code}') - ] + Button.url('Download', dl_link), + Button.url('Get File', deep_link) + ], + [ + Button.inline('Revoke', f'rm_{message_id}_{secret_code}') ] - ) + ] ) -@TelegramBot.on_message( - filters.channel - & ~filters.forwarded - & ~filters.media_group - & ( - filters.document - | filters.video - | filters.video_note - | filters.audio - | filters.voice - | filters.photo - ) -) -@verify_user -async def handle_channel_file(_, msg: Message): - if msg.caption and '#pass' in msg.caption: - return - +@TelegramBot.on(NewMessage(incoming=True, func=filter_files, forwards=False)) +@verify_user() +async def channel_file_handler(event: NewMessage.Event | Message): secret_code = token_hex(Telegram.SECRET_CODE_LENGTH) + event.message.text = f"`{secret_code}`" + message = await send_message(event.message) + message_id = message.id - try: - file = await msg.copy( - chat_id=Telegram.CHANNEL_ID, - caption=f'`{secret_code}`' - ) - except (errors.ChatForwardsRestricted, errors.MessageIdInvalid, errors.ChannelPrivate): - return + dl_link = f"{Server.BASE_URL}/dl/{message_id}?code={secret_code}" + tg_link = f"{Server.BASE_URL}/file/{message_id}?code={secret_code}" - file_id = file.id + if (event.document and "video" in event.document.mime_type) or event.video: + stream_link = f"{Server.BASE_URL}/stream/{message_id}?code={secret_code}" - dl_link = f'{Server.BASE_URL}/dl/{file_id}?code={secret_code}' - tg_link = f'{Server.BASE_URL}/file/{file_id}?code={secret_code}' - - if (msg.document and 'video' in msg.document.mime_type) or msg.video: - stream_link = f'{Server.BASE_URL}/stream/{file_id}?code={secret_code}' - await msg.edit_reply_markup( - InlineKeyboardMarkup( - [ - [ - InlineKeyboardButton('Download', url=dl_link), - InlineKeyboardButton('Stream', url=stream_link) - ], - [ - InlineKeyboardButton('Get File', url=tg_link) - ] + try: + await event.edit( + buttons=[ + [Button.url("Download", dl_link), Button.url("Stream", stream_link)], + [Button.url("Get File", tg_link)], ] ) - ) + except ( + MessageAuthorRequiredError, + MessageIdInvalidError, + MessageNotModifiedError, + ): + pass else: - await msg.edit_reply_markup( - InlineKeyboardMarkup( - [ - [ - InlineKeyboardButton('Download', url=dl_link), - InlineKeyboardButton('Get File', url=tg_link) - ] + try: + await event.edit( + buttons=[ + [Button.url("Download", dl_link), Button.url("Get File", tg_link)] ] ) - ) \ No newline at end of file + except ( + MessageAuthorRequiredError, + MessageIdInvalidError, + MessageNotModifiedError, + ): + pass diff --git a/bot/server/error.py b/bot/server/error.py index 1462b6a..b18ff72 100644 --- a/bot/server/error.py +++ b/bot/server/error.py @@ -28,4 +28,4 @@ async def http_error(error: HTTPError): return error.description or error_message, error.status_code def abort(status_code: int = 500, description: str = None): - raise HTTPError(status_code, description) \ No newline at end of file + raise HTTPError(status_code, description) diff --git a/bot/server/main.py b/bot/server/main.py index 558a33a..b736778 100644 --- a/bot/server/main.py +++ b/bot/server/main.py @@ -1,9 +1,9 @@ from quart import Blueprint, Response, request, render_template, redirect -from math import ceil, floor +from .error import abort from bot import TelegramBot from bot.config import Telegram, Server +from math import ceil, floor from bot.modules.telegram import get_message, get_file_properties -from .error import abort bp = Blueprint('main', __name__) @@ -13,22 +13,22 @@ async def home(): @bp.route('/dl/') async def transmit_file(file_id): - file = await get_message(int(file_id)) or abort(404) + file = await get_message(message_id=int(file_id)) or abort(404) code = request.args.get('code') or abort(401) range_header = request.headers.get('Range', 0) - if code != file.caption: + if code != file.raw_text: abort(403) - file_name, file_size, mime_type = await get_file_properties(file) - + file_name, file_size, mime_type = get_file_properties(file) + if range_header: from_bytes, until_bytes = range_header.replace("bytes=", "").split("-") from_bytes = int(from_bytes) until_bytes = int(until_bytes) if until_bytes else file_size - 1 else: - from_bytes = 0 - until_bytes = file_size -1 + from_bytes = 0 + until_bytes = file_size - 1 if (until_bytes > file_size) or (from_bytes < 0) or (until_bytes < from_bytes): abort(416, 'Invalid range.') @@ -36,26 +36,24 @@ async def transmit_file(file_id): chunk_size = 1024 * 1024 until_bytes = min(until_bytes, file_size - 1) - offset = from_bytes // chunk_size - first_part_cut = from_bytes - (from_bytes - (from_bytes % chunk_size)) + offset = from_bytes - (from_bytes % chunk_size) + first_part_cut = from_bytes - offset last_part_cut = until_bytes % chunk_size + 1 req_length = until_bytes - from_bytes + 1 - part_count = ceil(until_bytes / chunk_size) - floor((from_bytes - (from_bytes % chunk_size)) / chunk_size) + part_count = ceil(until_bytes / chunk_size) - floor(offset / chunk_size) + + headers = { + "Content-Type": f"{mime_type}", + "Content-Range": f"bytes {from_bytes}-{until_bytes}/{file_size}", + "Content-Length": str(req_length), + "Content-Disposition": f'attachment; filename="{file_name}"', + "Accept-Ranges": "bytes", + } - disposition = 'inline' if 'video' in mime_type or 'audio' in mime_type or 'html' in mime_type else 'attachment' - headers={ - 'Content-Type': f'{mime_type}', - 'Content-Range': f'bytes {from_bytes}-{until_bytes}/{file_size}', - 'Content-Length': str(req_length), - 'Content-Disposition': f'{disposition}; filename="{file_name}"', - 'Accept-Ranges': 'bytes', - } - - async def file_streamer(): + async def file_generator(): current_part = 1 - async for chunk in TelegramBot.stream_media(file, offset = offset): - + async for chunk in TelegramBot.iter_download(file, offset=offset, chunk_size=chunk_size, stride=chunk_size, file_size=file_size): if not chunk: break elif part_count == 1: @@ -66,12 +64,13 @@ async def transmit_file(file_id): yield chunk[:last_part_cut] else: yield chunk + current_part += 1 if current_part > part_count: break - return Response(file_streamer(), headers=headers, status=206 if range_header else 200) + return Response(file_generator(), headers=headers, status=206 if range_header else 200) @bp.route('/stream/') async def stream_file(file_id): diff --git a/bot/server/templates/player.html b/bot/server/templates/player.html index 93bd1e3..9371a26 100644 --- a/bot/server/templates/player.html +++ b/bot/server/templates/player.html @@ -1,4 +1,3 @@ - diff --git a/requirements.txt b/requirements.txt index 7cd3141..960ce6d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,4 @@ -pyrogram -tgcrypto +telethon +cryptg quart uvicorn -aiofiles