FileStream/bot/plugins/FileHandlers/stream.py

import os
import sys
import random
import urllib3
import logging
import asyncio
import traceback
import mimetypes
import time, math
from datetime import datetime
from typing import Dict, Union
from typing import Union, BinaryIO, List, Optional, Callable

from pyrogram import raw
from pyrogram import types
from pyrogram import filters, Client
from pyrogram import utils as pgutils
from pyrogram import StopTransmission, enums
from pyrogram import Client, utils, raw

from pyrogram.file_id import FileType
from pyrogram.session import Session, Auth
from pyrogram.enums import ParseMode, ChatType
from pyrogram.enums.parse_mode import ParseMode
from pyrogram.file_id import FileId, FileType, ThumbnailSource
from pyrogram.errors import FilePartMissing, AuthBytesInvalid, FloodWait
from pyrogram.types import Message, InlineKeyboardMarkup, InlineKeyboardButton

#--------------------------Local Imports-------------#

from FileStream.Database import Database
from FileStream.config import Telegram, Server
from FileStream.Tools.tool import TimeFormatter
from FileStream.Tools.progress import progress
from FileStream.bot import FileStream, req_client
from FileStream.bot import MULTI_CLIENTS, WORK_LOADS
from FileStream.utils.FileProcessors.bot_utils import is_user_banned, is_user_exist, is_user_joined, gen_link, is_channel_banned, is_channel_exist, is_user_authorized, upload_type_func
from FileStream.utils.FileProcessors.custom_dl import ByteStreamer
from FileStream.utils.FileProcessors.custom_ul import  TeleUploader
from FileStream.utils.FileProcessors.custom_mix import TGFileController
from FileStream.utils.FileProcessors.translation import LANG, BUTTON
from FileStream.utils.FileProcessors.human_readable import humanbytes
from FileStream.utils.FileProcessors.file_properties import get_file_ids, get_file_info, get_private_file_ids

db = Database(Telegram.DATABASE_URL, Telegram.SESSION_NAME)


@FileStream.on_message(filters.private & (filters.document  | filters.video   | filters.video_note   | filters.audio   | filters.voice   | filters.animation   | filters.photo),    group=4,)
async def private_receive_handler(bot: Client, message: Message):
    try:
        """
        //All Values Generated Here is Decided in get_file_info () so No Need of this Block
        message_id = message.id
        if message.chat.type == ChatType.PRIVATE:
        user_idx = message.from_user.id
        else:
        user_idx = message.chat.id
        """
        #For First Set All Files PUBLIC
        instruction = {
            "privacy_type":"PUBLIC",
            "user_id":message.from_user.id if (message.chat.type == ChatType.PRIVATE) else message.chat.id,
            "user_type": "TELEGRAM",
            "IMDB_id": "Unknown",
            "poster": "Unknown",   
            "title": "Unknown",
            "release_date":"Unknown",
            "keywords":"Unknown",
            "type": "Unknown",
            "description":"Unknown",
            "genre":"Unknown",
            "quality":"Unknown",
        }

        reply = await message.reply_text(LANG.PROCESSING_TEXT)
        file_info=get_file_info(message, instruction)
        #This Will give a option to Public Private and temporary Upload
        response = await upload_type_func(get_file_info(message, instruction),reply)
        print("Response at Upload_type_func:",response)
        if response["type"]=="ExistingFile":
            await reply.reply_photo(photo=response["poster"],
                              caption=response["stream_text"],
                              parse_mode=ParseMode.HTML,
                              reply_markup= response["reply_markup"],
            )
        else:
            await reply.edit_text(text=response["stream_text"],
                                parse_mode=ParseMode.HTML,
                                disable_web_page_preview=True,
                                reply_markup= response["reply_markup"],
                                )
    except FloodWait as e:
        print(f"Sleeping for {str(e.value)}s")
        
        await asyncio.sleep(e.value)
        await bot.send_message(
            chat_id=Telegram.ULOG_CHANNEL,
            text=
            f"Gᴏᴛ FʟᴏᴏᴅWᴀɪᴛ ᴏғ {str(e.value)}s ғʀᴏᴍ [{message.from_user.first_name}](tg://user?id={message.from_user.id})\n\n**ᴜsᴇʀ ɪᴅ :** `{str(message.from_user.id)}`",
            disable_web_page_preview=True,
            parse_mode=ParseMode.MARKDOWN)


@FileStream.on_message(filters.channel & ~filters.forwarded & ~filters.media_group & (filters.document  | filters.video  | filters.video_note  | filters.audio  | filters.voice  | filters.photo))
async def channel_receive_handler(bot: Client, message: Message):
    if await is_channel_banned(bot, message):
        return
    await is_channel_exist(bot, message)

    try:
        inserted_id = await db.add_file(file_info=get_file_info(message),  db_type="TEMPORARY")
        await get_file_ids(False, inserted_id, MULTI_CLIENTS, message)
        response = await gen_link(_id=inserted_id)
        await bot.edit_message_reply_markup(
            chat_id=message.chat.id,
            message_id=message.id,
            reply_markup=InlineKeyboardMarkup([[
                InlineKeyboardButton("Dᴏᴡɴʟᴏᴀᴅ ʟɪɴᴋ 📥",url=f"https://t.me/{FileStream.username}?start=stream_{str(inserted_id)}")
            ]]))

    except FloodWait as w:
        print(f"Sleeping for {str(w.x)}s")
        await asyncio.sleep(w.x)
        await bot.send_message(
            chat_id=Telegram.ULOG_GROUP,
            text=
            f"ɢᴏᴛ ғʟᴏᴏᴅᴡᴀɪᴛ ᴏғ {str(w.x)}s ғʀᴏᴍ {message.chat.title}\n\n**ᴄʜᴀɴɴᴇʟ ɪᴅ :** `{str(message.chat.id)}`",
            disable_web_page_preview=True)
    except Exception as e:
        exc_type, exc_obj, exc_tb = sys.exc_info()
        await bot.send_message(chat_id=Telegram.ULOG_GROUP,
                               text=f"**#EʀʀᴏʀTʀᴀᴄᴋᴇʙᴀᴄᴋ:** `{e} {exc_type, fname, exc_tb.tb_lineno}`",
                               disable_web_page_preview=True)
        print(
            f"Cᴀɴ'ᴛ Eᴅɪᴛ Bʀᴏᴀᴅᴄᴀsᴛ Mᴇssᴀɢᴇ!\nEʀʀᴏʀ:  **Gɪᴠᴇ ᴍᴇ ᴇᴅɪᴛ ᴘᴇʀᴍɪssɪᴏɴ ɪɴ ᴜᴘᴅᴀᴛᴇs ᴀɴᴅ ʙɪɴ Cʜᴀɴɴᴇʟ!{e}**"
        )


@FileStream.on_message(filters.reply)
async def reply_handler(
    bot: Client,
    message: Message,
    thumb: Union[str, BinaryIO] = None,
):
    replied_message = message.reply_to_message
    resp = ["n", "no", "not","N","NOT","NO"]
    if not replied_message:
        return
    if not replied_message.media:
        return
    text = message.text
    try:
        main_message = replied_message.reply_to_message
        instruction = {
            "privacy_type":"PRIVATE",
            "user_id":message.from_user.id if (message.chat.type == ChatType.PRIVATE) else message.chat.id,
            "user_type":"TELEGRAM"
        }
        #file_infos = get_file_info(main_message, instruction)
        #All the Time Get_file_ids should be called before update privacy or else tagged_users will be {}
        file_info = await db.update_privacy(get_file_info(main_message, instruction))

        #file_info = await db.get_file_by_fileuniqueid_only(
        #    file_infos['file']['file_unique_id'])
        mime_type = file_info['file']['mime_type']
        file_id = file_info['file']['file_id']
        file_size = file_info['file']['file_size']
        message_id = replied_message.id
        user_idx = message.from_user.id if (message.chat.type == ChatType.PRIVATE) else message.chat.id

        if text.lower() not in resp:
            ext = mimetypes.guess_extension(mime_type)
            file_info['file']['file_name'] = f"{message.text}{ext}"

        #index = min(WORK_LWORK_LOADS, key=WORK_LWORK_LOADS.get)
        #faster_client = MULTI_CLIENTS[index]
        client=await req_client()
        print(f"Using {client['index']} for {file_info['file']['file_name']}")
        logging.info(f"Client {client['index']} is now serving {request.headers.get('X-FORWARDED-FOR',request.remote)}")


        tg_connect = TGFileController(client['client'])
        req_file_id = await tg_connect.get_file_properties(file_info['_id'], MULTI_CLIENTS)

        start_time = time.time()
        file = await tg_connect.upload_file(
            index=client['index'],
            file_id=req_file_id,
            file_name=file_info['file']['file_name'],
            file_size=file_info['file']['file_size'],
            progress=progress,
            progress_args=(replied_message, file_info['file']['file_name'],file_info['file']['file_size'], start_time))

        success_text = LANG.SUCCESS_PRIV_FILE.format(file_info['file']['file_name'], humanbytes(file_size))
        await replied_message.edit_text(
            text=success_text,
            parse_mode=ParseMode.HTML,
            disable_web_page_preview=True,
        )

        media = raw.types.InputMediaUploadedDocument(
            mime_type=mime_type,
            file=file,
            force_file=False,
            thumb=thumb,
            attributes=[
                raw.types.DocumentAttributeFilename(
                    file_name=file_info['file']['file_name'])
            ])
        message_txt = f"* File Name: { file_info['file']['file_name']} \n* File Size :{humanbytes(file_info['file']['file_size'])} \n\n🔔 ᴛʜɪꜱ ꜰɪʟᴇ ɪꜱ ᴜᴘʟᴏᴀᴅᴇᴅ ᴛᴏ ᴘʀɪᴠᴀᴛᴇ ꜱᴘᴀᴄᴇ ʙʏ @ꜱᴘᴀʀᴋᴅʀɪᴠᴇ_ʙᴏᴛ"

        send_file = await tg_connect.send(media=media,
                                          caption=message_txt,
                                          reply_to_msg_id=None,
                                          chat_id=Telegram.PFLOG_CHANNEL)

        response = await gen_link(_id=file_info['_id'])
        await replied_message.edit_text(
            text=f"{message_txt}\n{response['stream_text']}",
            parse_mode=ParseMode.HTML,
            disable_web_page_preview=True,
            reply_markup=response['reply_markup'],
        )
        instruction = {
            "privacy_type":"PRIVATE",
            "user_id":message.from_user.id if(message.chat.type == ChatType.PRIVATE) else message.chat.id,
            "user_type":"TELEGRAM",
            "REUPLOADED":"True"
        }
        file_info = get_file_info(send_file, instruction)
        file_info['reupload']="True"
        # Here we are Adding the File Into the database First
        await db.add_private_file(file_info)
        #await get_private_file_ids(False, inserted_id, message)
        #All the Time Get_file_ids should be called before update privacy or else tagged_users will be {}
        #await db.update_privacy(file_info)

    except FloodWait as w:
        print(f"Sleeping for {str(w.wait_time)}s")
        await asyncio.sleep(w.wait_time)
        await bot.send_message(
            chat_id=Telegram.ULOG_GROUP,
            text=
            f"ɢᴏᴛ ғʟᴏᴏᴅᴡᴀɪᴛ ᴏғ {str(w.x)}s ғʀᴏᴍ {message.chat.title}\n\n**ᴄʜᴀɴɴᴇʟ ɪᴅ :** `{str(message.chat.id)}`",
            disable_web_page_preview=True)

    except Exception as e:
        exc_type, exc_obj, exc_tb = sys.exc_info()
        await bot.send_message(chat_id=Telegram.ULOG_GROUP,
                               text=f"**#EʀʀᴏʀTʀᴀᴄᴋᴇʙᴀᴄᴋ:** `{e} {exc_type, fname, exc_tb.tb_lineno}`",
                               disable_web_page_preview=True)
        print(f"Cᴀɴ'ᴛ Eᴅɪᴛ Bʀᴏᴀᴅᴄᴀsᴛ Mᴇssᴀɢᴇ!\nEʀʀᴏʀ:  **Gɪᴠᴇ ᴍᴇ ᴇᴅɪᴛ ᴘᴇʀᴍɪssɪᴏɴ ɪɴ ᴜᴘᴅᴀᴛᴇs ᴀɴᴅ ʙɪɴ Cʜᴀɴɴᴇʟ!{traceback.print_exc()}  (upload_file)**")
        #await db.delete_one_privfile(inserted_id)