From 8936029dc808289fa46e15b7c42b184ad8b0ab6a Mon Sep 17 00:00:00 2001 From: Ske Date: Tue, 24 Jul 2018 22:47:57 +0200 Subject: [PATCH] Massive refactor of pretty much everything in the bot --- .gitignore | 1 + bot/pluralkit/__init__.py | 1 - bot/pluralkit/bot.py | 122 --- bot/pluralkit/commands.py | 752 ------------------ bot/pluralkit/proxy.py | 213 ----- bot/pluralkit/utils.py | 304 ------- docker-compose.yml | 4 +- bot/requirements.txt => requirements.txt | 0 bot/Dockerfile => src/bot.Dockerfile | 2 +- bot/main.py => src/bot_main.py | 4 +- src/pluralkit/__init__.py | 26 + src/pluralkit/bot/__init__.py | 131 +++ src/pluralkit/bot/channel_logger.py | 109 +++ src/pluralkit/bot/commands/__init__.py | 98 +++ src/pluralkit/bot/commands/import_commands.py | 143 ++++ src/pluralkit/bot/commands/member_commands.py | 150 ++++ .../bot/commands/message_commands.py | 57 ++ src/pluralkit/bot/commands/misc_commands.py | 89 +++ src/pluralkit/bot/commands/mod_commands.py | 24 + src/pluralkit/bot/commands/switch_commands.py | 119 +++ src/pluralkit/bot/commands/system_commands.py | 215 +++++ {bot/pluralkit => src/pluralkit/bot}/help.py | 4 +- src/pluralkit/bot/proxy.py | 297 +++++++ src/pluralkit/bot/utils.py | 219 +++++ {bot => src}/pluralkit/db.py | 156 ++-- {bot => src}/pluralkit/stats.py | 2 - src/requirements.txt | 7 + 27 files changed, 1799 insertions(+), 1450 deletions(-) delete mode 100644 bot/pluralkit/__init__.py delete mode 100644 bot/pluralkit/bot.py delete mode 100644 bot/pluralkit/commands.py delete mode 100644 bot/pluralkit/proxy.py delete mode 100644 bot/pluralkit/utils.py rename bot/requirements.txt => requirements.txt (100%) rename bot/Dockerfile => src/bot.Dockerfile (82%) rename bot/main.py => src/bot_main.py (62%) create mode 100644 src/pluralkit/__init__.py create mode 100644 src/pluralkit/bot/__init__.py create mode 100644 src/pluralkit/bot/channel_logger.py create mode 100644 src/pluralkit/bot/commands/__init__.py create mode 100644 src/pluralkit/bot/commands/import_commands.py create mode 100644 src/pluralkit/bot/commands/member_commands.py create mode 100644 src/pluralkit/bot/commands/message_commands.py create mode 100644 src/pluralkit/bot/commands/misc_commands.py create mode 100644 src/pluralkit/bot/commands/mod_commands.py create mode 100644 src/pluralkit/bot/commands/switch_commands.py create mode 100644 src/pluralkit/bot/commands/system_commands.py rename {bot/pluralkit => src/pluralkit/bot}/help.py (98%) create mode 100644 src/pluralkit/bot/proxy.py create mode 100644 src/pluralkit/bot/utils.py rename {bot => src}/pluralkit/db.py (67%) rename {bot => src}/pluralkit/stats.py (97%) create mode 100644 src/requirements.txt diff --git a/.gitignore b/.gitignore index 39a35e82..d246dc3c 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,3 @@ .env .vscode/ +.idea/ \ No newline at end of file diff --git a/bot/pluralkit/__init__.py b/bot/pluralkit/__init__.py deleted file mode 100644 index 65af38e0..00000000 --- a/bot/pluralkit/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from . import commands, db, proxy, stats \ No newline at end of file diff --git a/bot/pluralkit/bot.py b/bot/pluralkit/bot.py deleted file mode 100644 index f0968839..00000000 --- a/bot/pluralkit/bot.py +++ /dev/null @@ -1,122 +0,0 @@ -import asyncio -from datetime import datetime -import logging -import json -import os -import time - -import discord - -logging.basicConfig(level=logging.INFO, format="[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s") -logging.getLogger("discord").setLevel(logging.INFO) -logging.getLogger("websockets").setLevel(logging.INFO) - -logger = logging.getLogger("pluralkit.bot") -logger.setLevel(logging.INFO) - -client = discord.Client() - - -@client.event -async def on_error(evt, *args, **kwargs): - logger.exception( - "Error while handling event {} with arguments {}:".format(evt, args)) - - -@client.event -async def on_ready(): - # Print status info - logger.info("Connected to Discord.") - logger.info("Account: {}#{}".format( - client.user.name, client.user.discriminator)) - logger.info("User ID: {}".format(client.user.id)) - - -@client.event -async def on_message(message): - # Ignore bot messages - if message.author.bot: - return - - # Split into args. shlex sucks so we don't bother with quotes - args = message.content.split(" ") - - from pluralkit import proxy, utils, stats - - command_items = utils.command_map.items() - command_items = sorted(command_items, key=lambda x: len(x[0]), reverse=True) - - prefix = "pk;" - for command, (func, _, _, _) in command_items: - if message.content.lower().startswith(prefix + command): - args_str = message.content[len(prefix + command):].strip() - args = args_str.split(" ") - - # Splitting on empty string yields one-element array, remove that - if len(args) == 1 and not args[0]: - args = [] - - async with client.pool.acquire() as conn: - time_before = time.perf_counter() - await func(conn, message, args) - time_after = time.perf_counter() - - # Report command time stats - execution_time = time_after - time_before - response_time = (datetime.now() - message.timestamp).total_seconds() - await stats.report_command(command, execution_time, response_time) - return - - # Try doing proxy parsing - async with client.pool.acquire() as conn: - await proxy.handle_proxying(conn, message) - -@client.event -async def on_socket_raw_receive(msg): - # Since on_reaction_add is buggy (only works for messages the bot's already cached, ie. no old messages) - # we parse socket data manually for the reaction add event - if isinstance(msg, str): - try: - msg_data = json.loads(msg) - if msg_data.get("t") == "MESSAGE_REACTION_ADD": - evt_data = msg_data.get("d") - if evt_data: - user_id = evt_data["user_id"] - message_id = evt_data["message_id"] - emoji = evt_data["emoji"]["name"] - - async with client.pool.acquire() as conn: - from pluralkit import proxy - await proxy.handle_reaction(conn, user_id, message_id, emoji) - except ValueError: - pass - -async def periodical_stat_timer(pool): - async with pool.acquire() as conn: - while True: - from pluralkit import stats - await stats.report_periodical_stats(conn) - await asyncio.sleep(30) - -async def run(): - from pluralkit import db, stats - try: - logger.info("Connecting to database...") - pool = await db.connect() - - logger.info("Attempting to create tables...") - async with pool.acquire() as conn: - await db.create_tables(conn) - - logger.info("Connecting to InfluxDB...") - await stats.connect() - - logger.info("Starting periodical stat reporting...") - asyncio.get_event_loop().create_task(periodical_stat_timer(pool)) - - client.pool = pool - logger.info("Connecting to Discord...") - await client.start(os.environ["TOKEN"]) - finally: - logger.info("Logging out from Discord...") - await client.logout() diff --git a/bot/pluralkit/commands.py b/bot/pluralkit/commands.py deleted file mode 100644 index 091f422c..00000000 --- a/bot/pluralkit/commands.py +++ /dev/null @@ -1,752 +0,0 @@ -from datetime import datetime, timezone -import io -import itertools -import json -import os -import re -from urllib.parse import urlparse - -import dateparser -import discord -from discord.utils import oauth_url -import humanize - -from pluralkit import db -from pluralkit.bot import client, logger -from pluralkit.utils import command, generate_hid, generate_member_info_card, generate_system_info_card, member_command, parse_mention, text_input, get_system_fuzzy, get_member_fuzzy, command_map, make_default_embed, parse_channel_mention, bounds_check_member_name, get_fronters, get_fronter_ids, get_front_history - -@command(cmd="system", usage="[system]", description="Shows information about a system.", category="System commands") -async def system_info(conn, message, args): - if len(args) == 0: - # Use sender's system - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - else: - # Look one up - system = await get_system_fuzzy(conn, args[0]) - - if system is None: - return False, "Unable to find system \"{}\".".format(args[0]) - - await client.send_message(message.channel, embed=await generate_system_info_card(conn, system)) - return True - - -@command(cmd="system new", usage="[name]", description="Registers a new system to this account.", category="System commands") -async def new_system(conn, message, args): - system = await db.get_system_by_account(conn, message.author.id) - - if system is not None: - return False, "You already have a system registered. To delete your system, use `pk;system delete`, or to unlink your system from this account, use `pk;system unlink`." - - system_name = None - if len(args) > 0: - system_name = " ".join(args) - - async with conn.transaction(): - # TODO: figure out what to do if this errors out on collision on generate_hid - hid = generate_hid() - - system = await db.create_system(conn, system_name=system_name, system_hid=hid) - - # Link account - await db.link_account(conn, system_id=system["id"], account_id=message.author.id) - return True, "System registered! To begin adding members, use `pk;member new `." - -@command(cmd="system set", usage=" [value]", description="Edits a system property. Leave [value] blank to clear.", category="System commands") -async def system_set(conn, message, args): - if len(args) == 0: - return False - - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - allowed_properties = ["name", "description", "tag", "avatar"] - db_properties = { - "name": "name", - "description": "description", - "tag": "tag", - "avatar": "avatar_url" - } - - prop = args[0] - if prop not in allowed_properties: - return False, "Unknown property {}. Allowed properties are {}.".format(prop, ", ".join(allowed_properties)) - - if len(args) >= 2: - value = " ".join(args[1:]) - - # Sanity checking - if prop == "tag": - # Make sure there are no members which would make the combined length exceed 32 - members_exceeding = await db.get_members_exceeding(conn, system_id=system["id"], length=32 - len(value)) - if len(members_exceeding) > 0: - # If so, error out and warn - member_names = ", ".join([member["name"] - for member in members_exceeding]) - logger.debug("Members exceeding combined length with tag '{}': {}".format(value, member_names)) - return False, "The maximum length of a name plus the system tag is 32 characters. The following members would exceed the limit: {}. Please reduce the length of the tag, or rename the members.".format(member_names) - - if prop == "avatar": - user = await parse_mention(value) - if user: - # Set the avatar to the mentioned user's avatar - # Discord doesn't like webp, but also hosts png alternatives - value = user.avatar_url.replace(".webp", ".png") - else: - # Validate URL - u = urlparse(value) - if u.scheme in ["http", "https"] and u.netloc and u.path: - value = value - else: - return False, "Invalid URL." - else: - # Clear from DB - value = None - - db_prop = db_properties[prop] - await db.update_system_field(conn, system_id=system["id"], field=db_prop, value=value) - - response = make_default_embed("{} system {}.".format("Updated" if value else "Cleared", prop)) - if prop == "avatar" and value: - response.set_image(url=value) - return True, response - -@command(cmd="system link", usage="", description="Links another account to your system.", category="System commands") -async def system_link(conn, message, args): - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - if len(args) == 0: - return False - - # Find account to link - linkee = await parse_mention(args[0]) - if not linkee: - return False, "Account not found." - - # Make sure account doesn't already have a system - account_system = await db.get_system_by_account(conn, linkee.id) - if account_system: - return False, "Account is already linked to a system (`{}`)".format(account_system["hid"]) - - # Send confirmation message - msg = await client.send_message(message.channel, "{}, please confirm the link by clicking the ✅ reaction on this message.".format(linkee.mention)) - await client.add_reaction(msg, "✅") - await client.add_reaction(msg, "❌") - - reaction = await client.wait_for_reaction(emoji=["✅", "❌"], message=msg, user=linkee) - # If account to be linked confirms... - if reaction.reaction.emoji == "✅": - async with conn.transaction(): - # Execute the link - await db.link_account(conn, system_id=system["id"], account_id=linkee.id) - return True, "Account linked to system." - else: - await client.clear_reactions(msg) - return False, "Account link cancelled." - - -@command(cmd="system unlink", description="Unlinks your system from this account. There must be at least one other account linked.", category="System commands") -async def system_unlink(conn, message, args): - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - # Make sure you can't unlink every account - linked_accounts = await db.get_linked_accounts(conn, system_id=system["id"]) - if len(linked_accounts) == 1: - return False, "This is the only account on your system, so you can't unlink it." - - async with conn.transaction(): - await db.unlink_account(conn, system_id=system["id"], account_id=message.author.id) - return True, "Account unlinked." - -@command(cmd="system fronter", usage="[system]", description="Gets the current fronter(s) in the system.", category="Switching commands") -async def system_fronter(conn, message, args): - if len(args) == 0: - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - else: - system = await get_system_fuzzy(conn, args[0]) - - if system is None: - return False, "Can't find system \"{}\".".format(args[0]) - - fronters, timestamp = await get_fronters(conn, system_id=system["id"]) - fronter_names = [member["name"] for member in fronters] - - embed = make_default_embed(None) - - if len(fronter_names) == 0: - embed.add_field(name="Current fronter", value="(no fronter)") - elif len(fronter_names) == 1: - embed.add_field(name="Current fronter", value=fronter_names[0]) - else: - embed.add_field(name="Current fronters", value=", ".join(fronter_names)) - - if timestamp: - embed.add_field(name="Since", value="{} ({})".format(timestamp.isoformat(sep=" ", timespec="seconds"), humanize.naturaltime(timestamp))) - return True, embed - -@command(cmd="system fronthistory", usage="[system]", description="Shows the past 10 switches in the system.", category="Switching commands") -async def system_fronthistory(conn, message, args): - if len(args) == 0: - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - else: - system = await get_system_fuzzy(conn, args[0]) - - if system is None: - return False, "Can't find system \"{}\".".format(args[0]) - - lines = [] - front_history = await get_front_history(conn, system["id"], count=10) - for i, (timestamp, members) in enumerate(front_history): - # Special case when no one's fronting - if len(members) == 0: - name = "(no fronter)" - else: - name = ", ".join([member["name"] for member in members]) - - # Make proper date string - time_text = timestamp.isoformat(sep=" ", timespec="seconds") - rel_text = humanize.naturaltime(timestamp) - - delta_text = "" - if i > 0: - last_switch_time = front_history[i-1][0] - delta_text = ", for {}".format(humanize.naturaldelta(timestamp - last_switch_time)) - lines.append("**{}** ({}, {}{})".format(name, time_text, rel_text, delta_text)) - - embed = make_default_embed("\n".join(lines) or "(none)") - embed.title = "Past switches" - return True, embed - - -@command(cmd="system delete", description="Deletes your system from the database ***permanently***.", category="System commands") -async def system_delete(conn, message, args): - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - await client.send_message(message.channel, "Are you sure you want to delete your system? If so, reply to this message with the system's ID (`{}`).".format(system["hid"])) - - msg = await client.wait_for_message(author=message.author, channel=message.channel, timeout=60.0) - if msg and msg.content == system["hid"]: - await db.remove_system(conn, system_id=system["id"]) - return True, "System deleted." - else: - return True, "System deletion cancelled." - -@member_command(cmd="member", description="Shows information about a system member.", system_only=False, category="Member commands") -async def member_info(conn, message, member, args): - await client.send_message(message.channel, embed=await generate_member_info_card(conn, member)) - return True - -@command(cmd="member new", usage="", description="Adds a new member to your system.", category="Member commands") -async def new_member(conn, message, args): - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - if len(args) == 0: - return False - - name = " ".join(args) - bounds_error = bounds_check_member_name(name, system["tag"]) - if bounds_error: - return False, bounds_error - - async with conn.transaction(): - # TODO: figure out what to do if this errors out on collision on generate_hid - hid = generate_hid() - - # Insert member row - await db.create_member(conn, system_id=system["id"], member_name=name, member_hid=hid) - return True, "Member \"{}\" (`{}`) registered!".format(name, hid) - - -@member_command(cmd="member set", usage=" [value]", description="Edits a member property. Leave [value] blank to clear.", category="Member commands") -async def member_set(conn, message, member, args): - if len(args) == 0: - return False - - allowed_properties = ["name", "description", "color", "pronouns", "birthdate", "avatar"] - db_properties = { - "name": "name", - "description": "description", - "color": "color", - "pronouns": "pronouns", - "birthdate": "birthday", - "avatar": "avatar_url" - } - - prop = args[0] - if prop not in allowed_properties: - return False, "Unknown property {}. Allowed properties are {}.".format(prop, ", ".join(allowed_properties)) - - if len(args) >= 2: - value = " ".join(args[1:]) - - # Sanity/validity checks and type conversions - if prop == "name": - system = await db.get_system(conn, member["system"]) - bounds_error = bounds_check_member_name(value, system["tag"]) - if bounds_error: - return False, bounds_error - - if prop == "color": - match = re.fullmatch("#?([0-9A-Fa-f]{6})", value) - if not match: - return False, "Color must be a valid hex color (eg. #ff0000)" - - value = match.group(1).lower() - - if prop == "birthdate": - try: - value = datetime.strptime(value, "%Y-%m-%d").date() - except ValueError: - try: - # Try again, adding 0001 as a placeholder year - # This is considered a "null year" and will be omitted from the info card - # Useful if you want your birthday to be displayed yearless. - value = value = datetime.strptime("0001-" + value, "%Y-%m-%d").date() - except ValueError: - return False, "Invalid date. Date must be in ISO-8601 format (eg. 1999-07-25)." - - if prop == "avatar": - user = await parse_mention(value) - if user: - # Set the avatar to the mentioned user's avatar - # Discord doesn't like webp, but also hosts png alternatives - value = user.avatar_url.replace(".webp", ".png") - else: - # Validate URL - u = urlparse(value) - if u.scheme in ["http", "https"] and u.netloc and u.path: - value = value - else: - return False, "Invalid URL." - else: - # Can't clear member name - if prop == "name": - return False, "Can't clear member name." - - # Clear from DB - value = None - - db_prop = db_properties[prop] - await db.update_member_field(conn, member_id=member["id"], field=db_prop, value=value) - - response = make_default_embed("{} {}'s {}.".format("Updated" if value else "Cleared", member["name"], prop)) - if prop == "avatar" and value: - response.set_image(url=value) - if prop == "color" and value: - response.colour = int(value, 16) - return True, response - -@member_command(cmd="member proxy", usage="[example]", description="Updates a member's proxy settings. Needs an \"example\" proxied message containing the string \"text\" (eg. [text], |text|, etc).", category="Member commands") -async def member_proxy(conn, message, member, args): - if len(args) == 0: - prefix, suffix = None, None - else: - # Sanity checking - example = " ".join(args) - if "text" not in example: - return False, "Example proxy message must contain the string 'text'." - - if example.count("text") != 1: - return False, "Example proxy message must contain the string 'text' exactly once." - - # Extract prefix and suffix - prefix = example[:example.index("text")].strip() - suffix = example[example.index("text")+4:].strip() - logger.debug( - "Matched prefix '{}' and suffix '{}'".format(prefix, suffix)) - - # DB stores empty strings as None, make that work - if not prefix: - prefix = None - if not suffix: - suffix = None - - async with conn.transaction(): - await db.update_member_field(conn, member_id=member["id"], field="prefix", value=prefix) - await db.update_member_field(conn, member_id=member["id"], field="suffix", value=suffix) - return True, "Proxy settings updated." if prefix or suffix else "Proxy settings cleared." - -@member_command("member delete", description="Deletes a member from your system ***permanently***.", category="Member commands") -async def member_delete(conn, message, member, args): - await client.send_message(message.channel, "Are you sure you want to delete {}? If so, reply to this message with the member's ID (`{}`).".format(member["name"], member["hid"])) - - msg = await client.wait_for_message(author=message.author, channel=message.channel, timeout=60.0) - if msg and msg.content == member["hid"]: - await db.delete_member(conn, member_id=member["id"]) - return True, "Member deleted." - else: - return True, "Member deletion cancelled." - -@command(cmd="message", usage="", description="Shows information about a proxied message. Requires the message ID.", category="Message commands") -async def message_info(conn, message, args): - try: - mid = int(args[0]) - except ValueError: - return False - - # Find the message in the DB - message_row = await db.get_message(conn, mid) - if not message_row: - return False, "Message not found." - - # Get the original sender of the message - original_sender = await client.get_user_info(str(message_row["sender"])) - - # Get sender member and system - member = await db.get_member(conn, message_row["member"]) - system = await db.get_system(conn, member["system"]) - - embed = discord.Embed() - embed.timestamp = discord.utils.snowflake_time(str(mid)) - embed.colour = discord.Colour.blue() - - if system["name"]: - system_value = "{} (`{}`)".format(system["name"], system["hid"]) - else: - system_value = "`{}`".format(system["hid"]) - embed.add_field(name="System", value=system_value) - embed.add_field(name="Member", value="{}: (`{}`)".format( - member["name"], member["hid"])) - embed.add_field(name="Sent by", value="{}#{}".format( - original_sender.name, original_sender.discriminator)) - embed.add_field(name="Content", value=message_row["content"], inline=False) - - embed.set_author(name=member["name"], url=member["avatar_url"]) - - await client.send_message(message.channel, embed=embed) - return True - -@command(cmd="switch", usage=" [name|id]...", description="Registers a switch and changes the current fronter.", category="Switching commands") -async def switch_member(conn, message, args): - if len(args) == 0: - return False - - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - members = [] - for member_name in args: - # Find the member - member = await get_member_fuzzy(conn, system["id"], member_name) - if not member: - return False, "Couldn't find member \"{}\".".format(member_name) - members.append(member) - - # Compare requested switch IDs and existing fronter IDs to check for existing switches - # Lists, because order matters, it makes sense to just swap fronters - member_ids = [member["id"] for member in members] - fronter_ids = (await get_fronter_ids(conn, system["id"]))[0] - if member_ids == fronter_ids: - if len(members) == 1: - return False, "{} is already fronting.".format(members[0]["name"]) - return False, "Members {} are already fronting.".format(", ".join([m["name"] for m in members])) - - # Log the switch - async with conn.transaction(): - switch_id = await db.add_switch(conn, system_id=system["id"]) - for member in members: - await db.add_switch_member(conn, switch_id=switch_id, member_id=member["id"]) - - if len(members) == 1: - return True, "Switch registered. Current fronter is now {}.".format(member["name"]) - else: - return True, "Switch registered. Current fronters are now {}.".format(", ".join([m["name"] for m in members])) - -@command(cmd="switch out", description="Registers a switch with no one in front.", category="Switching commands") -async def switch_out(conn, message, args): - system = await db.get_system_by_account(conn, message.author.id) - - if system is None: - return False, "No system is registered to this account." - - # Get current fronters - fronters, _ = await get_fronter_ids(conn, system_id=system["id"]) - if not fronters: - return False, "There's already no one in front." - - # Log it, and don't log any members - await db.add_switch(conn, system_id=system["id"]) - return True, "Switch-out registered." - -@command(cmd="switch move", usage="