From a079db8be02dd1f6ed3d9498bceb85d61a40ee0b Mon Sep 17 00:00:00 2001 From: Ske Date: Sun, 9 Sep 2018 20:38:57 +0200 Subject: [PATCH] Move various system functionality to system.py --- src/pluralkit/bot/commands/__init__.py | 4 +- src/pluralkit/bot/commands/import_commands.py | 5 +- src/pluralkit/bot/commands/member_commands.py | 7 +- src/pluralkit/bot/commands/system_commands.py | 135 ++++++------------ src/pluralkit/bot/utils.py | 5 - src/pluralkit/errors.py | 36 +++++ src/pluralkit/system.py | 70 ++++++++- src/pluralkit/utils.py | 22 +++ 8 files changed, 182 insertions(+), 102 deletions(-) create mode 100644 src/pluralkit/errors.py diff --git a/src/pluralkit/bot/commands/__init__.py b/src/pluralkit/bot/commands/__init__.py index 58c4a0c4..46d40327 100644 --- a/src/pluralkit/bot/commands/__init__.py +++ b/src/pluralkit/bot/commands/__init__.py @@ -4,7 +4,7 @@ import discord import logging import re import traceback -from typing import Tuple, Optional +from typing import Tuple, Optional, Union from pluralkit import db from pluralkit.system import System @@ -110,7 +110,7 @@ class CommandContext: async def reply(self, content=None, embed=None): return await self.client.send_message(self.message.channel, content=content, embed=embed) - async def confirm_react(self, user: discord.Member, message: str): + async def confirm_react(self, user: Union[discord.Member, discord.User], message: str): message = await self.reply(message) await self.client.add_reaction(message, "✅") diff --git a/src/pluralkit/bot/commands/import_commands.py b/src/pluralkit/bot/commands/import_commands.py index 0ec16f12..e74805e7 100644 --- a/src/pluralkit/bot/commands/import_commands.py +++ b/src/pluralkit/bot/commands/import_commands.py @@ -1,6 +1,7 @@ import asyncio from datetime import datetime +import pluralkit.utils from pluralkit.bot.commands import * logger = logging.getLogger("pluralkit.commands") @@ -100,7 +101,7 @@ async def import_tupperware(ctx: CommandContext): # Create new (nameless) system if there isn't any registered system = await ctx.get_system() if system is None: - hid = utils.generate_hid() + hid = pluralkit.utils.generate_hid() logger.debug("Creating new system (hid={})...".format(hid)) system = await db.create_system(ctx.conn, system_name=None, system_hid=hid) await db.link_account(ctx.conn, system_id=system.id, account_id=ctx.message.author.id) @@ -140,7 +141,7 @@ async def import_tupperware(ctx: CommandContext): existing_member = await db.get_member_by_name(ctx.conn, system_id=system.id, member_name=name) if not existing_member: # Or create a new member - hid = utils.generate_hid() + hid = pluralkit.utils.generate_hid() logger.debug("Creating new member {} (hid={})...".format(name, hid)) existing_member = await db.create_member(ctx.conn, system_id=system.id, member_name=name, member_hid=hid) diff --git a/src/pluralkit/bot/commands/member_commands.py b/src/pluralkit/bot/commands/member_commands.py index 8f817951..7ad7a9e2 100644 --- a/src/pluralkit/bot/commands/member_commands.py +++ b/src/pluralkit/bot/commands/member_commands.py @@ -1,10 +1,9 @@ -import re from datetime import datetime -from typing import List from urllib.parse import urlparse -from pluralkit.bot.commands import * +import pluralkit.utils from pluralkit.bot import help +from pluralkit.bot.commands import * logger = logging.getLogger("pluralkit.commands") @@ -26,7 +25,7 @@ async def new_member(ctx: CommandContext): return CommandError(bounds_error) # TODO: figure out what to do if this errors out on collision on generate_hid - hid = utils.generate_hid() + hid = pluralkit.utils.generate_hid() # Insert member row await db.create_member(ctx.conn, system_id=system.id, member_name=name, member_hid=hid) diff --git a/src/pluralkit/bot/commands/system_commands.py b/src/pluralkit/bot/commands/system_commands.py index 56df651d..8daf10fc 100644 --- a/src/pluralkit/bot/commands/system_commands.py +++ b/src/pluralkit/bot/commands/system_commands.py @@ -1,11 +1,12 @@ import dateparser import humanize from datetime import datetime -from urllib.parse import urlparse import pluralkit.utils from pluralkit.bot import help from pluralkit.bot.commands import * +from pluralkit.errors import ExistingSystemError, DescriptionTooLongError, TagTooLongError, TagTooLongWithMembersError, \ + InvalidAvatarURLError, UnlinkingLastAccountError logger = logging.getLogger("pluralkit.commands") @@ -20,90 +21,58 @@ async def system_info(ctx: CommandContext): async def new_system(ctx: CommandContext): - system = await ctx.get_system() - if system: + system_name = ctx.remaining() or None + + try: + await System.create_system(ctx.conn, ctx.message.author.id, system_name) + except ExistingSystemError: return CommandError( "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 = ctx.remaining() or None - - async with ctx.conn.transaction(): - # TODO: figure out what to do if this errors out on collision on generate_hid - hid = utils.generate_hid() - - system = await db.create_system(ctx.conn, system_name=system_name, system_hid=hid) - - # Link account - await db.link_account(ctx.conn, system_id=system.id, account_id=ctx.message.author.id) - return CommandSuccess("System registered! To begin adding members, use `pk;member new `.") + return CommandSuccess("System registered! To begin adding members, use `pk;member new `.") async def system_set(ctx: CommandContext): system = await ctx.ensure_system() - prop = ctx.pop_str(CommandError("You must pass a property name to set.", help=help.edit_system)) + property_name = ctx.pop_str(CommandError("You must pass a property name to set.", help=help.edit_system)) - allowed_properties = ["name", "description", "tag", "avatar"] - db_properties = { - "name": "name", - "description": "description", - "tag": "tag", - "avatar": "avatar_url" + async def avatar_setter(conn, url): + user = await utils.parse_mention(ctx.client, url) + if user: + # Set the avatar to the mentioned user's avatar + # Discord pushes webp by default, which isn't supported by webhooks, but also hosts png alternatives + url = user.avatar_url.replace(".webp", ".png") + + await system.set_avatar(conn, url) + + properties = { + "name": system.set_name, + "description": system.set_description, + "tag": system.set_tag, + "avatar": avatar_setter } - if prop not in allowed_properties: + if property_name not in properties: return CommandError( - "Unknown property {}. Allowed properties are {}.".format(prop, ", ".join(allowed_properties)), + "Unknown property {}. Allowed properties are {}.".format(property_name, ", ".join(allowed_properties)), help=help.edit_system) - if ctx.has_next(): - value = ctx.remaining() - # Sanity checking - if prop == "description": - if len(value) > 1024: - return CommandError("You can't have a description longer than 1024 characters.") + value = ctx.remaining() or None - if prop == "tag": - if len(value) > 32: - return CommandError("You can't have a system tag longer than 32 characters.") + try: + await properties[property_name](ctx.conn, value) + except DescriptionTooLongError: + return CommandError("You can't have a description longer than 1024 characters.") + except TagTooLongError: + return CommandError("You can't have a system tag longer than 32 characters.") + except TagTooLongWithMembersError as e: + return CommandError("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(", ".join(e.member_names))) + except InvalidAvatarURLError: + return CommandError("Invalid image URL.") - if re.search("", value): - return CommandError("Due to a Discord limitation, custom emojis aren't supported. Please use a standard emoji instead.") - - # Make sure there are no members which would make the combined length exceed 32 - members_exceeding = await db.get_members_exceeding(ctx.conn, system_id=system.id, - length=32 - len(value) - 1) - 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 CommandError( - "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 utils.parse_mention(ctx.client, 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 CommandError("Invalid image URL.") - else: - # Clear from DB - value = None - - db_prop = db_properties[prop] - await db.update_system_field(ctx.conn, system_id=system.id, field=db_prop, value=value) - - response = CommandSuccess("{} system {}.".format("Updated" if value else "Cleared", prop)) - #if prop == "avatar" and value: + response = CommandSuccess("{} system {}.".format("Updated" if value else "Cleared", property_name)) + # if prop == "avatar" and value: # response.set_image(url=value) return response @@ -118,36 +87,25 @@ async def system_link(ctx: CommandContext): return CommandError("Account not found.") # Make sure account doesn't already have a system - account_system = await db.get_system_by_account(ctx.conn, linkee.id) + account_system = await System.get_by_account(ctx.conn, linkee.id) if account_system: return CommandError("The mentioned account is already linked to a system (`{}`)".format(account_system.hid)) - # Send confirmation message - msg = await ctx.reply( - "{}, please confirm the link by clicking the ✅ reaction on this message.".format(linkee.mention)) - await ctx.client.add_reaction(msg, "✅") - await ctx.client.add_reaction(msg, "❌") - - reaction = await ctx.client.wait_for_reaction(emoji=["✅", "❌"], message=msg, user=linkee, timeout=60.0 * 5) - # If account to be linked confirms... - if not reaction: - return CommandError("Account link timed out.") - if not reaction.reaction.emoji == "✅": + if not await ctx.confirm_react(linkee, "{}, please confirm the link by clicking the ✅ reaction on this message.".format(linkee.mention)): return CommandError("Account link cancelled.") - await db.link_account(ctx.conn, system_id=system.id, account_id=linkee.id) + await system.link_account(ctx.conn, linkee.id) return CommandSuccess("Account linked to system.") async def system_unlink(ctx: CommandContext): system = await ctx.ensure_system() - # Make sure you can't unlink every account - linked_accounts = await db.get_linked_accounts(ctx.conn, system_id=system.id) - if len(linked_accounts) == 1: + try: + await system.unlink_account(ctx.conn, ctx.message.author.id) + except UnlinkingLastAccountError: return CommandError("This is the only account on your system, so you can't unlink it.") - await db.unlink_account(ctx.conn, system_id=system.id, account_id=ctx.message.author.id) return CommandSuccess("Account unlinked.") @@ -208,11 +166,12 @@ async def system_fronthistory(ctx: CommandContext): async def system_delete(ctx: CommandContext): system = await ctx.ensure_system() - delete_confirm_msg = "Are you sure you want to delete your system? If so, reply to this message with the system's ID (`{}`).".format(system.hid) + delete_confirm_msg = "Are you sure you want to delete your system? If so, reply to this message with the system's ID (`{}`).".format( + system.hid) if not await ctx.confirm_text(ctx.message.author, ctx.message.channel, system.hid, delete_confirm_msg): return CommandError("System deletion cancelled.") - await db.remove_system(ctx.conn, system_id=system.id) + await system.delete(ctx.conn) return CommandSuccess("System deleted.") diff --git a/src/pluralkit/bot/utils.py b/src/pluralkit/bot/utils.py index 53b26ac0..f74cb21a 100644 --- a/src/pluralkit/bot/utils.py +++ b/src/pluralkit/bot/utils.py @@ -1,7 +1,5 @@ import logging -import random import re -import string import discord import humanize @@ -16,9 +14,6 @@ logger = logging.getLogger("pluralkit.utils") def escape(s): return s.replace("`", "\\`") -def generate_hid() -> str: - return "".join(random.choices(string.ascii_lowercase, k=5)) - def bounds_check_member_name(new_name, system_tag): if len(new_name) > 32: return "Name cannot be longer than 32 characters." diff --git a/src/pluralkit/errors.py b/src/pluralkit/errors.py new file mode 100644 index 00000000..24d6f222 --- /dev/null +++ b/src/pluralkit/errors.py @@ -0,0 +1,36 @@ +class PluralKitError(Exception): + pass + + +class ExistingSystemError(PluralKitError): + pass + + +class DescriptionTooLongError(PluralKitError): + pass + + +class TagTooLongError(PluralKitError): + pass + + +class TagTooLongWithMembersError(PluralKitError): + def __init__(self, member_names): + self.member_names = member_names + + +class CustomEmojiError(PluralKitError): + pass + + +class InvalidAvatarURLError(PluralKitError): + pass + + +class AccountAlreadyLinkedError(PluralKitError): + def __init__(self, existing_system): + self.existing_system = existing_system + + +class UnlinkingLastAccountError(PluralKitError): + pass diff --git a/src/pluralkit/system.py b/src/pluralkit/system.py index 93592645..62fbf0bd 100644 --- a/src/pluralkit/system.py +++ b/src/pluralkit/system.py @@ -1,6 +1,10 @@ from datetime import datetime from collections.__init__ import namedtuple +from typing import Optional + +from pluralkit import db, errors +from pluralkit.utils import generate_hid, contains_custom_emoji, validate_avatar_url_or_raise class System(namedtuple("System", ["id", "hid", "name", "description", "tag", "avatar_url", "created"])): @@ -12,6 +16,70 @@ class System(namedtuple("System", ["id", "hid", "name", "description", "tag", "a avatar_url: str created: datetime + @staticmethod + async def get_by_account(conn, account_id: str) -> "System": + return await db.get_system_by_account(conn, account_id) + + @staticmethod + async def create_system(conn, account_id: str, system_name: Optional[str] = None) -> "System": + existing_system = await System.get_by_account(conn, account_id) + if existing_system: + raise errors.ExistingSystemError() + + new_hid = generate_hid() + + async with conn.transaction(): + new_system = await db.create_system(conn, system_name, new_hid) + await db.link_account(conn, new_system.id, account_id) + + return new_system + + async def set_name(self, conn, new_name: Optional[str]): + await db.update_system_field(conn, self.id, "name", new_name) + + async def set_description(self, conn, new_description: Optional[str]): + if new_description and len(new_description) > 1024: + raise errors.DescriptionTooLongError() + + await db.update_system_field(conn, self.id, "description", new_description) + + async def set_tag(self, conn, new_tag: Optional[str]): + if new_tag: + if len(new_tag) > 32: + raise errors.TagTooLongError() + + if contains_custom_emoji(new_tag): + raise errors.CustomEmojiError() + + members_exceeding = await db.get_members_exceeding(conn, system_id=self.id, length=32 - len(new_tag) - 1) + if len(members_exceeding) > 0: + raise errors.TagTooLongWithMembersError([member.name for member in members_exceeding]) + + await db.update_system_field(conn, self.id, "tag", new_tag) + + async def set_avatar(self, conn, new_avatar_url: Optional[str]): + if new_avatar_url: + validate_avatar_url_or_raise(new_avatar_url) + + await db.update_system_field(conn, self.id, "avatar", new_avatar_url) + + async def link_account(self, conn, new_account_id: str): + existing_system = await System.get_by_account(conn, new_account_id) + if existing_system: + raise errors.AccountAlreadyLinkedError(existing_system) + + await db.link_account(conn, self.id, new_account_id) + + async def unlink_account(self, conn, account_id: str): + linked_accounts = await db.get_linked_accounts(conn, self.id) + if len(linked_accounts) == 1: + raise errors.UnlinkingLastAccountError() + + await db.unlink_account(conn, self.id, account_id) + + async def delete(self, conn): + await db.remove_system(conn, self.id) + def to_json(self): return { "id": self.hid, @@ -19,4 +87,4 @@ class System(namedtuple("System", ["id", "hid", "name", "description", "tag", "a "description": self.description, "tag": self.tag, "avatar_url": self.avatar_url - } \ No newline at end of file + } diff --git a/src/pluralkit/utils.py b/src/pluralkit/utils.py index 7d459655..4e33ca86 100644 --- a/src/pluralkit/utils.py +++ b/src/pluralkit/utils.py @@ -1,7 +1,13 @@ +import re + +import random +import string from datetime import datetime, timezone from typing import List, Tuple +from urllib.parse import urlparse from pluralkit import db +from pluralkit.errors import InvalidAvatarURLError from pluralkit.member import Member @@ -48,3 +54,19 @@ async def get_front_history(conn, system_id, count) -> List[Tuple[datetime, List members = [all_members[id] for id in switch["members"]] out.append((timestamp, members)) return out + + +def generate_hid() -> str: + return "".join(random.choices(string.ascii_lowercase, k=5)) + + +def contains_custom_emoji(value): + return bool(re.search("", value)) + + +def validate_avatar_url_or_raise(url): + u = urlparse(url) + if not (u.scheme in ["http", "https"] and u.netloc and u.path): + raise InvalidAvatarURLError() + + # TODO: check file type and size of image \ No newline at end of file