FileStream/bot/plugins/Admin/admin.py

import os
import time
import string
import random
import asyncio
import aiofiles
import datetime

from FileStream.utils.FileProcessors.broadcast_helper import send_msg
from FileStream.Database import Database
from FileStream.bot import FileStream
from FileStream.Exceptions import FileNotFound
from FileStream.config import Telegram, Server
from pyrogram import filters, Client
from pyrogram.types import Message
from pyrogram.enums.parse_mode import ParseMode

db = Database(Telegram.DATABASE_URL, Telegram.SESSION_NAME)
broadcast_ids = {}


#Authorize User to Use the Services
@FileStream.on_message(filters.command("auth") & filters.private & filters.user(Telegram.OWNER_ID))
async def sts(c: Client, m: Message):
  await m.reply_text(
      text=f"""**Total Users in DB:** `{await db.total_users_count()}`
**Banned Users in DB:** `{await db.total_banned_users_count()}`
**Total Links Generated: ** `{await db.total_files()}`""",
      parse_mode=ParseMode.MARKDOWN,
      quote=True)

@FileStream.on_message(filters.command("add_user") & filters.private )
async def add_user(c: Client, m: Message):
  if await db.is_admin(m.from_user.id):
    if len(m.command) == 1:
      return await m.reply_text("**Usage:**\n /add_user <user_id>")
    try:
      user_id = int(m.command[1])
    except ValueError or UnboundLocalError:
      return await m.reply_text(f"**Usage:**\n <code> /add_admin <{user_id}> </code> ")
    await db.add_user(user_id)
    await m.reply_text(f"**User[{m.from_user.first_name} {m.from_user.last_name}] \n  User ID : {user_id} Added Successfully**")
  else:
    await m.reply_text(f"** Sorry Sir {user_id} You are not Admin **")
  
@FileStream.on_message(filters.command("add_admin") &  filters.user(Telegram.OWNER_ID) )
async def add_user(c: Client, m: Message):
  if await db.is_admin(m.from_user.id):
    if len(m.command) == 1:
      return await m.reply_text(f"**Usage:**\n <code> /add_admin <user_id> </code>")
    try:
      user_id = int(m.command[1])
    except ValueError or UnboundLocalError:
      return await m.reply_text(f"**Usage:**\n <code> /add_admin <{user_id}> </code>")
    await db.add_admin(user_id)
    await m.reply_text(f"**Admin [{m.from_user.first_name} {m.from_user.last_name}]\n  {user_id} Added Successfully**")
  else:
    await m.reply_text(f"** Sorry Sir [{m.from_user.first_name} {m.from_user.last_name}] {user_id} You are not Admin **")
  
  
@FileStream.on_message(filters.command("status") & filters.private & filters.user(Telegram.OWNER_ID))
async def sts(c: Client, m: Message):
  await m.reply_text(
      text=f"""**Total Users in DB:** `{await db.total_users_count()}`
**Banned Users in DB:** `{await db.total_banned_users_count()}`
**Total Links Generated: ** `{await db.total_files()}`""",
      parse_mode=ParseMode.MARKDOWN,
      quote=True)


@FileStream.on_message(filters.command("ban") & filters.private & filters.user(Telegram.OWNER_ID))
async def sts(b, m: Message):
  id = m.text.split("/ban ")[-1]
  if not await db.is_user_banned(int(id)):
    try:
      await db.ban_user(int(id))
      await db.delete_user(int(id))
      await m.reply_text(text=f"`{id}`** is Banned** ",
                         parse_mode=ParseMode.MARKDOWN,
                         quote=True)
      if not str(id).startswith('-100'):
        await b.send_message(chat_id=id,
                             text="**Your Banned to Use The Bot**",
                             parse_mode=ParseMode.MARKDOWN,
                             disable_web_page_preview=True)
    except Exception as e:
      await m.reply_text(text=f"**something went wrong: {e}** ",
                         parse_mode=ParseMode.MARKDOWN,
                         quote=True)
  else:
    await m.reply_text(text=f"`{id}`** is Already Banned** ",
                       parse_mode=ParseMode.MARKDOWN,
                       quote=True)


@FileStream.on_message(filters.command("unban") & filters.private & filters.user(Telegram.OWNER_ID))
async def sts(b, m: Message):
  id = m.text.split("/unban ")[-1]
  if await db.is_user_banned(int(id)):
    try:
      await db.unban_user(int(id))
      await m.reply_text(text=f"`{id}`** is Unbanned** ",
                         parse_mode=ParseMode.MARKDOWN,
                         quote=True)
      if not str(id).startswith('-100'):
        await b.send_message(chat_id=id,
                             text="**Your Unbanned now Use can use The Bot**",
                             parse_mode=ParseMode.MARKDOWN,
                             disable_web_page_preview=True)
    except Exception as e:
      await m.reply_text(text=f"** something went wrong: {e}**",
                         parse_mode=ParseMode.MARKDOWN,
                         quote=True)
  else:
    await m.reply_text(text=f"`{id}`** is not Banned** ",
                       parse_mode=ParseMode.MARKDOWN,
                       quote=True)


@FileStream.on_message(
    filters.command("broadcast") & filters.private
    & filters.user(Telegram.OWNER_ID) & filters.reply)
async def broadcast_(c, m):
  all_users = await db.get_all_users()
  broadcast_msg = m.reply_to_message
  while True:
    broadcast_id = ''.join(
        [random.choice(string.ascii_letters) for i in range(3)])
    if not broadcast_ids.get(broadcast_id):
      break
  out = await m.reply_text(
      text=
      f"Broadcast initiated! You will be notified with log file when all the users are notified."
  )
  start_time = time.time()
  total_users = await db.total_users_count()
  done = 0
  failed = 0
  success = 0
  broadcast_ids[broadcast_id] = dict(total=total_users,
                                     current=done,
                                     failed=failed,
                                     success=success)
  async with aiofiles.open('broadcast.txt', 'w') as broadcast_log_file:
    async for user in all_users:
      sts, msg = await send_msg(user_id=int(user['id']), message=broadcast_msg)
      if msg is not None:
        await broadcast_log_file.write(msg)
      if sts == 200:
        success += 1
      else:
        failed += 1
      if sts == 400:
        await db.delete_user(user['id'])
      done += 1
      if broadcast_ids.get(broadcast_id) is None:
        break
      else:
        broadcast_ids[broadcast_id].update(
            dict(current=done, failed=failed, success=success))
        try:
          await out.edit_text(
              f"Broadcast Status\n\ncurrent: {done}\nfailed:{failed}\nsuccess: {success}"
          )
        except:
          pass
  if broadcast_ids.get(broadcast_id):
    broadcast_ids.pop(broadcast_id)
  completed_in = datetime.timedelta(seconds=int(time.time() - start_time))
  await asyncio.sleep(3)
  await out.delete()
  if failed == 0:
    await m.reply_text(
        text=
        f"broadcast completed in `{completed_in}`\n\nTotal users {total_users}.\nTotal done {done}, {success} success and {failed} failed.",
        quote=True)
  else:
    await m.reply_document(
        document='broadcast.txt',
        caption=
        f"broadcast completed in `{completed_in}`\n\nTotal users {total_users}.\nTotal done {done}, {success} success and {failed} failed.",
        quote=True)
  os.remove('broadcast.txt')


@FileStream.on_message(
    filters.command("del") & filters.private & filters.user(Telegram.OWNER_ID))
async def sts(c: Client, m: Message):
  file_id = m.text.split(" ")[-1]
  try:
    file_info = await db.get_file(file_id)
  except FIleNotFound:
    await m.reply_text(text=f"**File Already Deleted**", quote=True)
    return
  await db.delete_one_file(file_info['_id'])
  await db.count_links(file_info['user_id'], "-")
  await m.reply_text(text=f"**Fɪʟᴇ Dᴇʟᴇᴛᴇᴅ Sᴜᴄᴄᴇssғᴜʟʟʏ !** ", quote=True)