Update broadcast.py

plugins/broadcast.py

@@ -1,80 +1,23 @@
-import os
 import time
-import math
-import json
 import string
 import random
-import traceback
 import asyncio
 import datetime
-import aiofiles
-from io import BytesIO
-from pyrogram import Client, filters
-from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
-from pyrogram.errors import FloodWait, InputUserDeactivated, UserIsBlocked, PeerIdInvalid, UserNotParticipant, UserBannedInChannel
-from pyrogram.errors.exceptions.bad_request_400 import PeerIdInvalid
-
-
-class Database:
-    def __init__(self, url, database_name):
-        self._client = motor.motor_asyncio.AsyncIOMotorClient(url)
-        self.db = self._client[database_name]
-        self.col = self.db.users
-        self.cache = {}
-
-    def new_user(self, id):
-        return {"id": id, "language": "en"}
-
-    async def add_user(self, id):
-        user = self.new_user(id)
-        await self.col.insert_one(user)
-
-    async def get_user(self, id):
-        user = self.cache.get(id)
-        if user is not None:
-            return user
-
-        user = await self.col.find_one({"id": int(id)})
-        self.cache[id] = user
-        return user
-
-    async def is_user_exist(self, id):
-        user = await self.col.find_one({'id': int(id)})
-        return True if user else False
-
-    async def total_users_count(self):
-        count = await self.col.count_documents({})
-        return count
-
-    async def get_all_users(self):
-        all_users = self.col.find({})
-        return all_users
-
-    async def delete_user(self, user_id):
-        await self.col.delete_many({'id': int(user_id)})
-
-    async def get_lang(self, id):
-        user = await self.col.find_one({'id': int(id)})
-        if user:
-            return user.get("language", "en")
-        else:
-            return "en"
-
-    async def update_lang(self, id, language):
-        await self.col.update_one(
-            {"id": id}, {"$set": {"language": language}}
-        )
-
+import aiofiles.os
+import traceback
+from config import BROADCAST_AS_COPY
+from szbot.helpers.database.access_db import db
+from pyrogram.types import Message
+from pyrogram.errors import FloodWait, InputUserDeactivated, UserIsBlocked, PeerIdInvalid
 
-BOT_OWNER = int(os.environ.get("ADMIN"))
-DATABASE = os.environ.get("DB_URL")
-db = Database(DATABASE, "sztranslatorbot")
 broadcast_ids = {}
 
-
 async def send_msg(user_id, message):
     try:
-        await message.copy(chat_id=user_id)
+        if BROADCAST_AS_COPY is False:
+            await message.forward(chat_id=user_id)
+        elif BROADCAST_AS_COPY is True:
+            await message.copy(chat_id=user_id)
         return 200, None
     except FloodWait as e:
         await asyncio.sleep(e.x)
@@ -89,44 +32,66 @@ async def send_msg(user_id, message):
         return 500, f"{user_id} : {traceback.format_exc()}\n"
 
 
-@Client.on_message(filters.private & filters.command("broadcast") & filters.user(ADMIN) & filters.reply)
-async def broadcast(bot, update):
-	all_users = await db.get_all_users()
-	broadcast_msg = update.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 update.reply_text(text=f"Broadcast Started! You will be notified with log file when all the users are notified. @sztranslatorbot")
-	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))
-	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 update.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 update.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.")
-	os.remove('broadcast.txt')
+async def broadcast_handler(m: Message):
+    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 Started! You will be notified with log file when all the users are notified 😊\n\n@szimagebot 🤖"
+    )
+    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
+                    )
+                )
+    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.\n\n@szimagebot 🤖",
+            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.\n\n@szimagebot 🤖",
+            quote=True
+        )
+    await aiofiles.os.remove('broadcast.txt')