Merge pull request #13 from kotolex/master

FEATURE: mute and warnings

src/commands.py

@@ -7,9 +7,9 @@ from .admin_commands import bot
 from .constants import (LOG_COMM, FAQ, LIB, RULES, RUS, RUS_ENG_TABLE, ENG_RUS_TABLE, PYTHONCHATRU, ZEN, LUTZ_ID,
                         BDMTSS_ID)
 from .filters import in_spam_list, in_caption_spam_list, in_delete_list, is_hello_text, is_invalid_name
-from .helpers import represent_as_get, detect_args, is_admin, fetch_rule
+from .helpers import represent_as_get, detect_args, is_admin, fetch_rule, has_warnings, warn_user, warnings_count
 from .report import update_stats, increment
-from .stan import act, speak
+from .stan import act, speak, mute_for_one_day
 
 
 async def send_or_reply(message: types.Message, answer, **kwargs):
@@ -358,7 +358,15 @@ async def handle_invalid_name(message: types.Message):
     rule_num = 6
     markup = types.InlineKeyboardMarkup([[RULES]], 1)
     logging.info(LOG_COMM % (message.chat.title, message.from_user.id, message.from_user.first_name, message.text))
-    await send_or_reply(message, f'<b><a href="tg:user?id={message.from_user.id}">{message.from_user.full_name}</a>, правило {rule_num}</b>\n<i>{fetch_rule(rule_num)}</i>', reply_markup=markup)
+    user_id = message.from_user.id
+    if not has_warnings(user_id):
+        await send_or_reply(message, f'<b><a href="tg:user?id={user_id}">{message.from_user.full_name}'
+                                     f'</a>, правило {rule_num}</b>\n<i>{fetch_rule(rule_num)}</i>',
+                            reply_markup=markup)
+    warn_user(user_id)
+    if warnings_count(user_id) >= 3:
+        logging.info(f"[MUTED] User {user_id} has been muted for violating 6 rule")
+        await mute_for_one_day(message.chat.id, user_id)
     await bot.delete_message(message.chat.id, message.id)
 
 

src/helpers.py

@@ -8,6 +8,22 @@ from telebot import types
 
 from .constants import SPAM, BAN_WORDS, NON_GRATA, ADMIN_ID, ONLY_RUS_LETTERS, ONLY_ENG_LETTERS, RULES_TEXT, SYMBOLS
 
+WARNS: dict[int, int] = {}
+
+
+def has_warnings(user_id: int) -> bool:
+    return user_id in WARNS
+
+
+def warn_user(user_id: int):
+    if user_id not in WARNS:
+        WARNS[user_id] = 0
+    WARNS[user_id] += 1
+
+
+def warnings_count(user_id: int) -> int:
+    return WARNS.get(user_id, 0)
+
 
 def is_url_reachable(url: str) -> bool:
     try:

src/stan.py

@@ -3,6 +3,7 @@ import asyncio
 import html
 import logging
 import random
+from time import time
 
 from telebot import types
 
@@ -12,6 +13,9 @@ from .models import all_chat_quotes, add_quote, is_quote_in_chat, delete_quote_i
 from .robot import bot
 
 TYPING_TIMEOUT = 0.13 / 4  # Reading time is quarter of the same text writing time
+FOR_EVER_TIMEOUT = 0
+ONE_DAY_TIMEOUT = 86400
+ONE_WEEK_TIMEOUT = 86400 * 7
 
 
 def speak(chance_of: int, chat_id: int) -> None | str:
@@ -36,6 +40,23 @@ async def act(message: types.Message):
         await send_quote(len(quote) * 0.13, message, quote)
 
 
+async def mute_forever(chat_id: int, user_id: int):
+    await mute(chat_id, user_id, FOR_EVER_TIMEOUT)
+
+
+async def mute_for_one_day(chat_id: int, user_id: int):
+    await mute(chat_id, user_id, ONE_DAY_TIMEOUT)
+
+
+async def mute_for_one_week(chat_id: int, user_id: int):
+    await mute(chat_id, user_id, ONE_WEEK_TIMEOUT)
+
+
+async def mute(chat_id, user_id, period):
+    logging.info(f"[MUTED] User {user_id} has been muted for: {period} seconds")
+    await bot.restrict_chat_member(chat_id, user_id, until_date=time() + period)
+
+
 @bot.message_handler(func=is_white_id, commands=["add"])
 async def add_stan_quote(message: types.Message):
     if message.reply_to_message and message.reply_to_message.text:

tests/test_bot.py

@@ -511,6 +511,15 @@ class TestBot(IsolatedAsyncioTestCase):
         self.assertEqual(RESULTS[0][1], 'sendMessage')
         self.assertEqual(RESULTS[1][1], 'deleteMessage')
 
+    async def test_invalid_name_user_muted(self):
+        await self.bot.process_new_updates([get_update('новый текст', user_id=111, first_name=None)])
+        await self.bot.process_new_updates([get_update('новый текст', user_id=111, first_name=None)])
+        await self.bot.process_new_updates([get_update('новый текст', user_id=111, first_name=None)])
+        self.assertEqual(RESULTS[1][1], 'deleteMessage')
+        self.assertEqual(RESULTS[2][1], 'deleteMessage')
+        self.assertEqual(RESULTS[4][1], 'deleteMessage')
+        self.assertEqual(RESULTS[3][1], 'restrictChatMember')
+
 
 if __name__ == '__main__':
     main()