Handle edited messages

This commit is contained in:
Michał Rudowicz 2024-06-20 16:42:48 +02:00
parent 284311e9d5
commit 9a52cdf3cf
2 changed files with 50 additions and 16 deletions

View File

@ -16,24 +16,22 @@ def load_config(filename: str = "config.json"):
return retval return retval
async def new_msg(update, context, regexes, allowed_chats): async def new_msg(update, context, regexes, allowed_chats):
if update.message is None: if update.message is None and update.edited_message is None:
logging.info(f"Got following update: {update}") logging.info(f"Got following unknown update: {update}")
return return
if update.message.chat.id not in allowed_chats: message = update.message
logging.info(f"Got update from not allowed chat: {update.message.chat}. Ignoring") 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 return
is_spam = False
for regex in regexes: for regex in regexes:
if update.message is not None and \ if message.text is not None and regex.search(message.text) is not None:
update.message.text is not None and \ logging.info(f"Banning {message.from_user.name} from {message.chat.effective_name} for posting a message matching {regex.pattern}")
regex.search(update.message.text) is not None: await message.chat.ban_member(message.from_user.id)
is_spam = True await message.delete()
break return
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.")
async def handle_error(update, context): async def handle_error(update, context):
logging.error("Exception while handling an update:", exc_info=context.error) logging.error("Exception while handling an update:", exc_info=context.error)

View File

@ -43,6 +43,7 @@ class Message:
@dataclass @dataclass
class Update: class Update:
message = Message() message = Message()
edited_message = None
def make_update(msg: str, chat_id: int) -> Update: 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 update.message.chat.id = chat_id
return update 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): class TestAutoBanBot(unittest.IsolatedAsyncioTestCase):
def setUp(self): def setUp(self):
@ -61,13 +70,22 @@ class TestAutoBanBot(unittest.IsolatedAsyncioTestCase):
async def test_not_bannable(self): async def test_not_bannable(self):
messages = ["not bannable message", "siematest_123elo", "i am not a bot", "t.me, i'm not a bot"] messages = ["not bannable message", "siematest_123elo", "i am not a bot", "t.me, i'm not a bot"]
for msg in messages: for msg in messages:
with self.subTest(msg=msg): with self.subTest("new message", message=msg):
update = make_update(msg, ALLOWED_CHAT_ID) update = make_update(msg, ALLOWED_CHAT_ID)
self.assertIsNone(update.edited_message)
update.message.delete.reset_mock() update.message.delete.reset_mock()
update.message.chat.ban_member.reset_mock() update.message.chat.ban_member.reset_mock()
await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
update.message.delete.assert_not_called() update.message.delete.assert_not_called()
update.message.chat.ban_member.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): async def test_bannable(self):
messages = [ messages = [
@ -80,6 +98,7 @@ class TestAutoBanBot(unittest.IsolatedAsyncioTestCase):
for msg in messages: for msg in messages:
with self.subTest("Allowed chat causes a ban", message=msg): with self.subTest("Allowed chat causes a ban", message=msg):
update = make_update(msg, ALLOWED_CHAT_ID) update = make_update(msg, ALLOWED_CHAT_ID)
self.assertIsNone(update.edited_message)
update.message.delete.reset_mock() update.message.delete.reset_mock()
update.message.chat.ban_member.reset_mock() update.message.chat.ban_member.reset_mock()
await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) 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) update.message.chat.ban_member.assert_called_once_with(update.message.from_user.id)
with self.subTest("Unknown chat is ignored", message=msg): with self.subTest("Unknown chat is ignored", message=msg):
update = make_update(msg, DISALLOWED_CHAT_ID) update = make_update(msg, DISALLOWED_CHAT_ID)
self.assertIsNone(update.edited_message)
update.message.delete.reset_mock() update.message.delete.reset_mock()
update.message.chat.ban_member.reset_mock() update.message.chat.ban_member.reset_mock()
await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats) await autobanbot.new_msg(update, None, self.regexes, self.allowed_chats)
update.message.delete.assert_not_called() update.message.delete.assert_not_called()
update.message.chat.ban_member.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__': if __name__ == '__main__':