Added the DELETE_MESSAGES option to the antiflood module, minor fixes

This commit is contained in:
nocturn9x 2020-06-11 08:58:50 +00:00
parent 9e6e158ac8
commit 66c56dda5a
4 changed files with 62 additions and 65 deletions

View File

@ -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}")

View File

@ -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)})

View File

@ -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)

4
bot.py
View File

@ -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()")