diff options
| -rw-r--r-- | bot/cogs/token_remover.py | 64 | ||||
| -rw-r--r-- | tests/bot/cogs/test_token_remover.py | 45 |
2 files changed, 90 insertions, 19 deletions
diff --git a/bot/cogs/token_remover.py b/bot/cogs/token_remover.py index ef979f222..93ceda6be 100644 --- a/bot/cogs/token_remover.py +++ b/bot/cogs/token_remover.py @@ -18,6 +18,11 @@ LOG_MESSAGE = ( "Censored a seemingly valid token sent by {author} (`{author_id}`) in {channel}, " "token was `{user_id}.{timestamp}.{hmac}`" ) +DECODED_LOG_MESSAGE = "The token user_id decodes into {user_id}." +USER_TOKEN_MESSAGE = ( + "The token user_id decodes into {user_id}, " + "which matches `{user_name}` and means this is a valid USER token." +) DELETION_MESSAGE_TEMPLATE = ( "Hey {mention}! I noticed you posted a seemingly valid Discord API " "token in your message and have removed your message. " @@ -92,7 +97,14 @@ class TokenRemover(Cog): await msg.channel.send(DELETION_MESSAGE_TEMPLATE.format(mention=msg.author.mention)) - log_message = self.format_log_message(msg, found_token) + user_name = None + user_id = self.extract_user_id(found_token.user_id) + user = msg.guild.get_member(user_id) + + if user: + user_name = str(user) + + log_message = self.format_log_message(msg, found_token, user_id, user_name) log.debug(log_message) # Send pretty mod log embed to mod-alerts @@ -103,14 +115,24 @@ class TokenRemover(Cog): text=log_message, thumbnail=msg.author.avatar_url_as(static_format="png"), channel_id=Channels.mod_alerts, + ping_everyone=user_name is not None, ) self.bot.stats.incr("tokens.removed_tokens") @staticmethod - def format_log_message(msg: Message, token: Token) -> str: - """Return the log message to send for `token` being censored in `msg`.""" - return LOG_MESSAGE.format( + def format_log_message( + msg: Message, + token: Token, + user_id: int, + user_name: t.Optional[str] = None, + ) -> str: + """ + Return the log message to send for `token` being censored in `msg`. + + Additonally, mention if the token was decodable into a user id, and if that resolves to a user on the server. + """ + message = LOG_MESSAGE.format( author=msg.author, author_id=msg.author.id, channel=msg.channel.mention, @@ -118,6 +140,11 @@ class TokenRemover(Cog): timestamp=token.timestamp, hmac='x' * len(token.hmac), ) + if user_name: + more = USER_TOKEN_MESSAGE.format(user_id=user_id, user_name=user_name) + else: + more = DECODED_LOG_MESSAGE.format(user_id=user_id) + return message + "\n" + more @classmethod def find_token_in_message(cls, msg: Message) -> t.Optional[Token]: @@ -134,23 +161,34 @@ class TokenRemover(Cog): return @staticmethod - def is_valid_user_id(b64_content: str) -> bool: - """ - Check potential token to see if it contains a valid Discord user ID. - - See: https://discordapp.com/developers/docs/reference#snowflakes - """ + def extract_user_id(b64_content: str) -> t.Optional[int]: + """Return a userid integer from part of a potential token, or None if it couldn't be decoded.""" b64_content = utils.pad_base64(b64_content) try: decoded_bytes = base64.urlsafe_b64decode(b64_content) string = decoded_bytes.decode('utf-8') - - # isdigit on its own would match a lot of other Unicode characters, hence the isascii. - return string.isascii() and string.isdigit() + if not (string.isascii() and string.isdigit()): + # This case triggers if there are fancy unicode digits in the base64 encoding, + # that means it's not a valid user id. + return None + return int(string) except (binascii.Error, ValueError): + return None + + @classmethod + def is_valid_user_id(cls, b64_content: str) -> bool: + """ + Check potential token to see if it contains a valid Discord user ID. + + See: https://discordapp.com/developers/docs/reference#snowflakes + """ + decoded_id = cls.extract_user_id(b64_content) + if not decoded_id: return False + return True + @staticmethod def is_valid_timestamp(b64_content: str) -> bool: """ diff --git a/tests/bot/cogs/test_token_remover.py b/tests/bot/cogs/test_token_remover.py index 3349caa73..275350144 100644 --- a/tests/bot/cogs/test_token_remover.py +++ b/tests/bot/cogs/test_token_remover.py @@ -22,6 +22,7 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase): self.msg = MockMessage(id=555, content="hello world") self.msg.channel.mention = "#lemonade-stand" + self.msg.guild.get_member = MagicMock(return_value="Bob") self.msg.author.__str__ = MagicMock(return_value=self.msg.author.name) self.msg.author.avatar_url_as.return_value = "picture-lemon.png" @@ -230,15 +231,41 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase): results = [match[0] for match in results] self.assertCountEqual((token_1, token_2), results) - @autospec("bot.cogs.token_remover", "LOG_MESSAGE") - def test_format_log_message(self, log_message): + @autospec("bot.cogs.token_remover", "LOG_MESSAGE", "DECODED_LOG_MESSAGE") + def test_format_log_message(self, log_message, decoded_log_message): + """Should correctly format the log message with info from the message and token.""" + token = Token("NDcyMjY1OTQzMDYyNDEzMzMy", "XsySD_", "s45jqDV_Iisn-symw0yDRrk_jf4") + log_message.format.return_value = "Howdy" + decoded_log_message.format.return_value = " Partner" + + return_value = TokenRemover.format_log_message(self.msg, token, 472265943062413332, None) + + self.assertEqual( + return_value, + log_message.format.return_value + "\n" + decoded_log_message.format.return_value, + ) + log_message.format.assert_called_once_with( + author=self.msg.author, + author_id=self.msg.author.id, + channel=self.msg.channel.mention, + user_id=token.user_id, + timestamp=token.timestamp, + hmac="x" * len(token.hmac), + ) + + @autospec("bot.cogs.token_remover", "LOG_MESSAGE", "USER_TOKEN_MESSAGE") + def test_format_log_message_user_token(self, log_message, user_token_message): """Should correctly format the log message with info from the message and token.""" token = Token("NDY3MjIzMjMwNjUwNzc3NjQx", "XsySD_", "s45jqDV_Iisn-symw0yDRrk_jf4") log_message.format.return_value = "Howdy" + user_token_message.format.return_value = "Partner" - return_value = TokenRemover.format_log_message(self.msg, token) + return_value = TokenRemover.format_log_message(self.msg, token, 467223230650777641, "Bob") - self.assertEqual(return_value, log_message.format.return_value) + self.assertEqual( + return_value, + log_message.format.return_value + "\n" + user_token_message.format.return_value, + ) log_message.format.assert_called_once_with( author=self.msg.author, author_id=self.msg.author.id, @@ -247,6 +274,10 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase): timestamp=token.timestamp, hmac="x" * len(token.hmac), ) + user_token_message.format.assert_called_once_with( + user_id=467223230650777641, + user_name="Bob", + ) @mock.patch.object(TokenRemover, "mod_log", new_callable=mock.PropertyMock) @autospec("bot.cogs.token_remover", "log") @@ -256,6 +287,7 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase): cog = TokenRemover(self.bot) mod_log = mock.create_autospec(ModLog, spec_set=True, instance=True) token = mock.create_autospec(Token, spec_set=True, instance=True) + token.user_id = "no-id" log_msg = "testing123" mod_log_property.return_value = mod_log @@ -268,7 +300,7 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase): token_remover.DELETION_MESSAGE_TEMPLATE.format(mention=self.msg.author.mention) ) - format_log_message.assert_called_once_with(self.msg, token) + format_log_message.assert_called_once_with(self.msg, token, None, "Bob") logger.debug.assert_called_with(log_msg) self.bot.stats.incr.assert_called_once_with("tokens.removed_tokens") @@ -279,7 +311,8 @@ class TokenRemoverTests(unittest.IsolatedAsyncioTestCase): title="Token removed!", text=log_msg, thumbnail=self.msg.author.avatar_url_as.return_value, - channel_id=constants.Channels.mod_alerts + channel_id=constants.Channels.mod_alerts, + ping_everyone=True, ) @mock.patch.object(TokenRemover, "mod_log", new_callable=mock.PropertyMock) |