diff --git a/bot/db_utils.py b/bot/db_utils.py index 802e3f0..94e915e 100644 --- a/bot/db_utils.py +++ b/bot/db_utils.py @@ -90,6 +90,9 @@ def delete_player(username: str) -> bool: ) player = CURSOR.fetchone() + if not player: + return False + player_id = player[0] # Delete pulls @@ -107,10 +110,40 @@ def delete_player(username: str) -> bool: return True -def is_player_administrator(player_id: int) -> bool: +def ban_player(username: str) -> bool: + '''Adds a player to the ban list.''' + try: + CURSOR.execute( + 'INSERT INTO banned_players (handle) VALUES (?)', + (username,) + ) + return True + except sqlite3.IntegrityError: + return False + + +def unban_player(username: str) -> bool: + '''Removes a player from the ban list.''' CURSOR.execute( - 'SELECT is_administrator FROM PLAYERS WHERE id = ? LIMIT 1', - (player_id,) + 'DELETE FROM banned_players WHERE handle = ?', + (username,) + ) + return CURSOR.rowcount > 0 + + +def is_player_banned(username: str) -> bool: + CURSOR.execute( + 'SELECT * FROM banned_players WHERE handle = ?', + (username,) + ) + row = CURSOR.fetchone() + return row is not None + + +def is_player_administrator(username: str) -> bool: + CURSOR.execute( + 'SELECT is_administrator FROM players WHERE username = ? LIMIT 1', + (username,) ) row = CURSOR.fetchone() return row[0] if row else False @@ -151,7 +184,8 @@ def add_to_whitelist(instance: str) -> bool: present''' try: CURSOR.execute( - 'INSERT INTO instance_whitelist (tld) VALUES (?)', (instance,)) + 'INSERT INTO instance_whitelist (tld) VALUES (?)', (instance,) + ) return True except sqlite3.IntegrityError: return False diff --git a/bot/notification.py b/bot/notification.py index faac1d3..deb8ec6 100644 --- a/bot/notification.py +++ b/bot/notification.py @@ -6,7 +6,7 @@ from misskey.exceptions import MisskeyAPIException from config import NOTIFICATION_BATCH_SIZE, USE_WHITELIST from parsing import parse_notification -from db_utils import get_config, set_config, is_whitelisted +from db_utils import get_config, set_config, is_whitelisted, is_player_banned from response import generate_response from custom_types import BotResponse @@ -44,6 +44,11 @@ def process_notification( if not parsed_notification: return + author = parsed_notification['author'] + if is_player_banned(author): + print(f'⚠️ Blocked notification from banned player: {author}') + return + # Get the note Id to reply to note_id = notification.get('note', {}).get('id') diff --git a/bot/parsing.py b/bot/parsing.py index eece077..e1e8583 100644 --- a/bot/parsing.py +++ b/bot/parsing.py @@ -24,6 +24,8 @@ def parse_notification( note_id = note_obj.get("id") note = note_text.strip().lower() if note_text else "" + # Split words into tokens + parts = note.split() # Check for both short and fully-qualified name mentions username_variants = [ @@ -31,18 +33,16 @@ def parse_notification( f'@{config.USER.split("@")[1]}' ] - # Make sure the notification text explicitly mentions the bot - if not any(variant in note for variant in username_variants): + # Notifs must consist of the initial mention and at least one other token + if len(parts) <= 1: return None - # Find command and arguments after the mention - # Removes all mentions - # regex = mentions that start with @ and may contain @domain - cleaned_text = re.sub(r"@\w+(?:@\S+)?", "", note).strip() - parts = cleaned_text.split() + # Make sure the first token is a mention to the bot + if not parts[0] in username_variants: + return None - command = parts[0].lower() if parts else None - arguments = parts[1:] if len(parts) > 1 else [] + command = parts[1].lower() + arguments = parts[2:] if len(parts) > 2 else [] return { 'author': full_user, diff --git a/bot/response.py b/bot/response.py index 16bd05e..f21d7b0 100644 --- a/bot/response.py +++ b/bot/response.py @@ -211,6 +211,52 @@ the whitelist', } +def do_ban(author: str, args: list[str]) -> BotResponse: + if len(args) == 0: + return { + 'message': f'{author} Please specify a user to ban', + 'attachment_urls': None + } + + if db.is_player_administrator(args[0]): + return { + 'message': f'{author} Cannot ban other administrators.', + 'attachment_urls': None + } + + if db.ban_player(args[0]): + # Delete banned player's account + db.delete_player(args[0]) + return { + 'message': f'{author} 🔨 **BONK!** Get banned, {args[0]}!', + 'attachment_urls': None + } + else: + return { + 'message': f'{author} Player is already banned: {args[0]}', + 'attachment_urls': None + } + + +def do_unban(author: str, args: list[str]) -> BotResponse: + if len(args) == 0: + return { + 'message': f'{author} Please specify a user to unban', + 'attachment_urls': None + } + + if db.unban_player(args[0]): + return { + 'message': f'{author} Player unbanned: {args[0]}!', + 'attachment_urls': None + } + else: + return { + 'message': f'{author} Player was not banned: {args[0]}', + 'attachment_urls': None + } + + def generate_response(notification: ParsedNotification) -> BotResponse | None: '''Given a command with arguments, processes the game state and returns a response''' @@ -252,7 +298,7 @@ def generate_response(notification: ParsedNotification) -> BotResponse | None: pass # Commands beyond this point require the user to be an administrator - if not db.is_player_administrator(player_id): + if not db.is_player_administrator(author): return res # Admin commands @@ -261,10 +307,10 @@ def generate_response(notification: ParsedNotification) -> BotResponse | None: res = do_whitelist(author, notification['arguments']) case 'unwhitelist': res = do_unwhitelist(author, notification['arguments']) - # case 'ban': - # res = do_ban(author, notification['arguments']) - # case 'unban': - # res = do_unban(author, notification['arguments']) + case 'ban': + res = do_ban(author, notification['arguments']) + case 'unban': + res = do_unban(author, notification['arguments']) case _: pass diff --git a/migrations/0005_add_whitelist.sql b/migrations/0005_add_whitelist.sql new file mode 100644 index 0000000..d24f2e3 --- /dev/null +++ b/migrations/0005_add_whitelist.sql @@ -0,0 +1,7 @@ +CREATE TABLE IF NOT EXISTS instance_whitelist ( + tld TEXT UNIQUE PRIMARY KEY +); + +CREATE TABLE IF NOT EXISTS banned_players ( + handle TEXT UNIQUE PRIMARY KEY +);