FileStream/bot/plugins/Admin/admin.py

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

from pyrogram.types import Message
from pyrogram import filters, Client
from FileStream.bot import FileStream
from FileStream.Database import Database
from FileStream.Exceptions import FileNotFound
from FileStream.config import Telegram, Server

from pyrogram.enums.parse_mode import ParseMode
from FileStream.utils.FileProcessors.translation import LANG, BUTTON
from FileStream.utils.FileProcessors.broadcast_helper import send_msg

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

@FileStream.on_message(filters.command("admin") & filters.private & filters.user(Telegram.OWNER_ID))
async def iamadmin(c: Client, m: Message):
  if await db.is_admin(m.from_user.id):
    await m.reply_text(
        text=LANG.ADMIN_GUIDELINES_TEXT_Y.format("/add_user <UserID>","/add_admin <UserID>"),
        parse_mode=ParseMode.HTML,
    )
  

#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 & 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("**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> ")
    details= await FileStream.get_users(user_id)
    #print("Details :",details)
    await db.add_user(user_id, details)
    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>")
    details= await FileStream.get_users(user_id)
    #print("Details :",details)
    """Details : {
    "_": "User",
    "id": ,
    "is_self": false,
    "is_contact": false,
    "is_mutual_contact": false,
    "is_deleted": false,
    "is_bot": false,
    "is_verified": false,
    "is_restricted": false,
    "is_scam": false,
    "is_fake": false,
    "is_support": false,
    "is_premium": false,
    "is_contacts_only": false,
    "first_name": "",
    "last_name": "",
    "status": "UserStatus.RECENTLY",
    "username": "",
    "language_code": "en"
    }"""
    await db.add_admin(user_id, details )
    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)