From 66c56dda5a1a53d9603c83260415a0ea1eca7de5 Mon Sep 17 00:00:00 2001 From: nocturn9x Date: Thu, 11 Jun 2020 08:58:50 +0000 Subject: [PATCH] Added the DELETE_MESSAGES option to the antiflood module, minor fixes --- BotBase/database/query.py | 24 ++++------- BotBase/modules/antiflood.py | 83 +++++++++++++++++++----------------- BotBase/modules/start.py | 16 +++---- bot.py | 4 +- 4 files changed, 62 insertions(+), 65 deletions(-) diff --git a/BotBase/database/query.py b/BotBase/database/query.py index 9c9e13b..5bf2902 100644 --- a/BotBase/database/query.py +++ b/BotBase/database/query.py @@ -17,11 +17,10 @@ def create_database(path: str, query: str): else: try: with database: - cursor = database.cursor() - cursor.executescript(query) - cursor.close() + database.executescript(query) + database.close() except sqlite3.Error as query_error: - logging.info(f"An error has occurred while executing query: {query_error}") + logging.info(f"An error has occurred while executing DB_CREATE: {query_error}") def get_user(tg_id: int): @@ -48,9 +47,8 @@ def update_name(tg_id: int, name: str): else: try: with database: - cursor = database.cursor() - query = cursor.execute(DB_UPDATE_NAME, (name, tg_id)) - return query.fetchone() + database.execute(DB_UPDATE_NAME, (name, tg_id)) + return True except sqlite3.Error as query_error: logging.error(f"An error has occurred while executing DB_UPDATE_NAME query: {query_error}") return query_error @@ -80,9 +78,7 @@ def set_user(tg_id: int, uname: str): else: try: with database: - cursor = database.cursor() - cursor.execute(DB_SET_USER, (None, tg_id, uname, time.strftime("%d/%m/%Y %T %p"), 0)) - cursor.close() + database.execute(DB_SET_USER, (None, tg_id, uname, time.strftime("%d/%m/%Y %T %p"), 0)) return True except sqlite3.Error as query_error: logging.error(f"An error has occurred while executing DB_GET_USERS query: {query_error}") @@ -97,9 +93,7 @@ def ban_user(tg_id: int): else: try: with database: - cursor = database.cursor() - cursor.execute(DB_BAN_USER, (tg_id, )) - cursor.close() + database.execute(DB_BAN_USER, (tg_id, )) return True except sqlite3.Error as query_error: logging.error(f"An error has occurred while executing DB_BAN_USER query: {query_error}") @@ -114,9 +108,7 @@ def unban_user(tg_id: int): else: try: with database: - cursor = database.cursor() - cursor.execute(DB_UNBAN_USER, (tg_id, )) - cursor.close() + database.execute(DB_UNBAN_USER, (tg_id, )) return True except sqlite3.Error as query_error: logging.error(f"An error has occurred while executing DB_UNBAN_USER query: {query_error}") diff --git a/BotBase/modules/antiflood.py b/BotBase/modules/antiflood.py index a9471ec..cd62c04 100644 --- a/BotBase/modules/antiflood.py +++ b/BotBase/modules/antiflood.py @@ -1,10 +1,11 @@ -from pyrogram import Client, Filters, Message +from pyrogram import Client, Filters, Message, CallbackQueryHandler from ..config import MAX_UPDATE_THRESHOLD, ANTIFLOOD_SENSIBILITY, BAN_TIME, ADMINS, BYPASS_FLOOD, FLOOD_NOTICE, \ - COUNT_CALLBACKS_SEPARATELY, FLOOD_PERCENTAGE, CACHE, PRIVATE_ONLY + COUNT_CALLBACKS_SEPARATELY, FLOOD_PERCENTAGE, CACHE, PRIVATE_ONLY, DELETE_MESSAGES from collections import defaultdict import logging import time from ..methods.safe_send import send_message +from ..methods.various import delete_messages # Some variables for runtime configuration @@ -22,49 +23,51 @@ def is_flood(updates: list): genexpr = [i <= ANTIFLOOD_SENSIBILITY for i in ((updates[i + 1] - timestamp) if i < (MAX_UPDATE_THRESHOLD - 1) else (timestamp - updates[i - 1]) for i, timestamp in enumerate(updates))] - limit = (len(genexpr) / 100) * FLOOD_PERCENTAGE - if genexpr.count(True) >= limit: - return True - else: - return False + return genexpr.count(True) >= int((len(genexpr) / 100) * FLOOD_PERCENTAGE) -@Client.on_callback_query(~BYPASS_USERS, group=-1) @Client.on_message(FILTER & ~BYPASS_USERS, group=-1) def anti_flood(client, update): """Anti flood module""" - VAR = MESSAGES if isinstance(update, Message) else QUERIES - if isinstance(VAR[update.from_user.id], tuple): - chat, date = VAR[update.from_user.id] - if time.time() - date >= BAN_TIME: - logging.warning(f"{update.from_user.id} has waited at least {BAN_TIME} seconds and can now text again") - BANNED_USERS.remove(update.from_user.id) - del VAR[update.from_user.id] - elif len(VAR[update.from_user.id]) >= MAX_UPDATE_THRESHOLD: - logging.info(f"MAX_MESS_THRESHOLD ({MAX_UPDATE_THRESHOLD}) Reached for {update.from_user.id}") - timestamps = VAR.pop(update.from_user.id) - if is_flood(timestamps): - logging.warning(f"Flood detected from {update.from_user.id} in chat {update.chat.id}") - if update.from_user.id in CACHE: - del CACHE[update.from_user.id] - BANNED_USERS.add(update.from_user.id) - if isinstance(update, Message): - chatid = update.chat.id - else: - chatid = update.from_user.id - VAR[update.from_user.id] = chatid, time.monotonic() - if FLOOD_NOTICE: - send_message(client, update.from_user.id, FLOOD_NOTICE) - else: - if update.from_user.id in VAR: - del VAR[update.from_user.id] + user_id = update.from_user.id + if isinstance(update, Message): + VAR = MESSAGES + chat = update.chat.id + date = update.date + message_id = update.message_id else: - if isinstance(update, Message): - date = update.date + VAR = QUERIES + message_id = None + chat = user_id + if update.message: + date = update.message.date else: - if update.message: - date = update.message.date - else: - date = time.monotonic() - VAR[update.from_user.id].append(date) + date = time.time() + if isinstance(VAR[user_id], tuple): + chat, date = VAR[user_id] + if time.time() - date >= BAN_TIME: + logging.warning(f"{user_id} has waited at least {BAN_TIME} seconds and can now text again") + BANNED_USERS.remove(user_id) + del VAR[user_id] + elif len(VAR[user_id]) >= MAX_UPDATE_THRESHOLD: + VAR[user_id].append({chat: (date, message_id)}) + logging.info(f"MAX_UPDATE_THRESHOLD ({MAX_UPDATE_THRESHOLD}) Reached for {user_id}") + user_data = VAR.pop(user_id) + timestamps = [list(*d.values())[0] for d in user_data] + updates = [list(*d.values())[1] for d in user_data] + if is_flood(timestamps): + logging.warning(f"Flood detected from {user_id} in chat {chat}") + if user_id in CACHE: + del CACHE[user_id] + BANNED_USERS.add(user_id) + VAR[user_id] = chat, time.time() + if FLOOD_NOTICE: + send_message(client, True, user_id, FLOOD_NOTICE) + if DELETE_MESSAGES and any(updates): + delete_messages(client, True, chat, filter(bool, updates)) + else: + if user_id in VAR: + del VAR[user_id] + else: + VAR[user_id].append({chat: (date, message_id)}) diff --git a/BotBase/modules/start.py b/BotBase/modules/start.py index 96b9ec8..ce52f78 100644 --- a/BotBase/modules/start.py +++ b/BotBase/modules/start.py @@ -14,6 +14,8 @@ def check_user_banned(tg_id: int): if isinstance(res, Exception): return False else: + if not res: + return False if res[-1]: return True else: @@ -50,13 +52,13 @@ def start_handler(client, message): ) -@Client.on_callback_query(Filters.callback_data("info")) +@Client.on_callback_query(Filters.callback_data("info") & ~BANNED_USERS) def bot_info(_, query): buttons = InlineKeyboardMarkup([[InlineKeyboardButton("🔙 Back", "back_start")]]) edit_message_text(query, True, CREDITS.format(), reply_markup=buttons) -@Client.on_callback_query(Filters.callback_data("back_start")) +@Client.on_callback_query(Filters.callback_data("back_start") & ~BANNED_USERS) def back_start(_, query): if query.from_user.first_name: name = query.from_user.first_name @@ -69,9 +71,7 @@ def back_start(_, query): if isinstance(data, list): for chatid, message_ids in data: delete_messages(_, True, chatid, message_ids) - start_handler(_, query) - else: - edit_message_text(query, True, - GREET.format(mention=f"[{name}](tg://user?id={query.from_user.id})", id=query.from_user.id, - username=query.from_user.username), - reply_markup=BUTTONS) + edit_message_text(query, True, + GREET.format(mention=f"[{name}](tg://user?id={query.from_user.id})", id=query.from_user.id, + username=query.from_user.username), + reply_markup=BUTTONS) diff --git a/bot.py b/bot.py index aaacc91..45030bd 100644 --- a/bot.py +++ b/bot.py @@ -1,5 +1,5 @@ import logging -from pyrogram import Client +from pyrogram import Client, CallbackQueryHandler from pyrogram.session import Session import importlib @@ -7,10 +7,12 @@ import importlib if __name__ == "__main__": MODULE_NAME = "BotBase" # Change this to match the FOLDER name that contains the config.py file conf = importlib.import_module(f"{MODULE_NAME}.config") + antiflood = importlib.import_module(f"{MODULE_NAME}.modules.antiflood") dbmodule = importlib.import_module(f"{MODULE_NAME}.database.query") logging.basicConfig(format=conf.LOGGING_FORMAT, datefmt=conf.DATE_FORMAT, level=conf.LOGGING_LEVEL) bot = Client(api_id=conf.API_ID, api_hash=conf.API_HASH, bot_token=conf.BOT_TOKEN, plugins=conf.PLUGINS_ROOT, session_name=conf.SESSION_NAME, workers=conf.WORKERS_NUM) + bot.add_handler(CallbackQueryHandler(antiflood.anti_flood, ~antiflood.BYPASS_USERS), group=-1) Session.notice_displayed = True try: logging.warning("Running create_database()")