diff --git a/autobanbot.py b/autobanbot.py index 8019ece..327c413 100755 --- a/autobanbot.py +++ b/autobanbot.py @@ -16,24 +16,22 @@ def load_config(filename: str = "config.json"): return retval async def new_msg(update, context, regexes, allowed_chats): - if update.message is None: - logging.info(f"Got following update: {update}") + if update.message is None and update.edited_message is None: + logging.info(f"Got following unknown update: {update}") return - if update.message.chat.id not in allowed_chats: - logging.info(f"Got update from not allowed chat: {update.message.chat}. Ignoring") + message = update.message + if update.edited_message is not None: + message = update.edited_message + if message is None: + return + if message.chat.id not in allowed_chats: return - is_spam = False for regex in regexes: - if update.message is not None and \ - update.message.text is not None and \ - regex.search(update.message.text) is not None: - is_spam = True - break - if is_spam: - logging.info(f"Banning {update.message.from_user.name} from {update.message.chat.effective_name}") - await update.message.chat.ban_member(update.message.from_user.id) - await update.message.delete() - logging.info("Banned.") + if message.text is not None and regex.search(message.text) is not None: + logging.info(f"Banning {message.from_user.name} from {message.chat.effective_name} for posting a message matching {regex.pattern}") + await message.chat.ban_member(message.from_user.id) + await message.delete() + return async def handle_error(update, context): logging.error("Exception while handling an update:", exc_info=context.error) diff --git a/tests.py b/tests.py index 87e4580..b21ef9b 100644 --- a/tests.py +++ b/tests.py @@ -43,6 +43,7 @@ class Message: @dataclass class Update: message = Message() + edited_message = None def make_update(msg: str, chat_id: int) -> Update: @@ -51,6 +52,14 @@ def make_update(msg: str, chat_id: int) -> Update: update.message.chat.id = chat_id return update +def make_edited_message(msg: str, chat_id: int) -> Update: + update = Update() + update.message = None + update.edited_message = Message() + update.edited_message.text = msg + update.edited_message.chat.id = chat_id + return update + class TestAutoBanBot(unittest.IsolatedAsyncioTestCase): def setUp(self): @@ -61,13 +70,22 @@ class TestAutoBanBot(unittest.IsolatedAsyncioTestCase): async def test_not_bannable(self): messages = ["not bannable message", "siematest_123elo", "i am not a bot", "t.me, i'm not a bot"] for msg in messages: - with self.subTest(msg=msg): + with self.subTest("new message", message=msg): update = make_update(msg, ALLOWED_CHAT_ID) + self.assertIsNone(update.edited_message) update.message.delete.reset_mock() update.message.chat.ban_member.reset_mock() await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) update.message.delete.assert_not_called() update.message.chat.ban_member.assert_not_called() + with self.subTest("message edit", message=msg): + update = make_edited_message(msg, ALLOWED_CHAT_ID) + self.assertIsNone(update.message) + update.edited_message.delete.reset_mock() + update.edited_message.chat.ban_member.reset_mock() + await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) + update.edited_message.delete.assert_not_called() + update.edited_message.chat.ban_member.assert_not_called() async def test_bannable(self): messages = [ @@ -80,6 +98,7 @@ class TestAutoBanBot(unittest.IsolatedAsyncioTestCase): for msg in messages: with self.subTest("Allowed chat causes a ban", message=msg): update = make_update(msg, ALLOWED_CHAT_ID) + self.assertIsNone(update.edited_message) update.message.delete.reset_mock() update.message.chat.ban_member.reset_mock() await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) @@ -87,11 +106,28 @@ class TestAutoBanBot(unittest.IsolatedAsyncioTestCase): update.message.chat.ban_member.assert_called_once_with(update.message.from_user.id) with self.subTest("Unknown chat is ignored", message=msg): update = make_update(msg, DISALLOWED_CHAT_ID) + self.assertIsNone(update.edited_message) update.message.delete.reset_mock() update.message.chat.ban_member.reset_mock() await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) update.message.delete.assert_not_called() update.message.chat.ban_member.assert_not_called() + with self.subTest("Edit in allowed chat causes a ban", message=msg): + update = make_edited_message(msg, ALLOWED_CHAT_ID) + self.assertIsNone(update.message) + update.edited_message.delete.reset_mock() + update.edited_message.chat.ban_member.reset_mock() + await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) + update.edited_message.delete.assert_called_once() + update.edited_message.chat.ban_member.assert_called_once_with(update.edited_message.from_user.id) + with self.subTest("Edit in unknown chat is ignored", message=msg): + update = make_edited_message(msg, DISALLOWED_CHAT_ID) + self.assertIsNone(update.message) + update.edited_message.delete.reset_mock() + update.edited_message.chat.ban_member.reset_mock() + await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) + update.edited_message.delete.assert_not_called() + update.edited_message.chat.ban_member.assert_not_called() if __name__ == '__main__':