Format code
This commit is contained in:
Ske
parent
b81b768b04
commit
902c14ef65
@ -1 +1 @@
|
||||
from . import commands, db, proxy
|
||||
from . import commands, db, proxy
|
||||
|
||||
@ -12,23 +12,28 @@ logger.setLevel(logging.DEBUG)
|
||||
|
||||
client = discord.Client()
|
||||
|
||||
|
||||
@client.event
|
||||
async def on_error(evt, *args, **kwargs):
|
||||
logger.exception("Error while handling event {} with arguments {}:".format(evt, args))
|
||||
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("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(" ")
|
||||
|
||||
@ -54,6 +59,7 @@ async def on_message(message):
|
||||
async with client.pool.acquire() as conn:
|
||||
await proxy.handle_proxying(conn, message)
|
||||
|
||||
|
||||
@client.event
|
||||
async def on_reaction_add(reaction, user):
|
||||
from pluralkit import proxy
|
||||
@ -62,6 +68,7 @@ async def on_reaction_add(reaction, user):
|
||||
async with client.pool.acquire() as conn:
|
||||
await proxy.handle_reaction(conn, reaction, user)
|
||||
|
||||
|
||||
async def run():
|
||||
from pluralkit import db
|
||||
try:
|
||||
@ -79,4 +86,4 @@ async def run():
|
||||
await client.start(os.environ["TOKEN"])
|
||||
finally:
|
||||
logger.info("Logging out from Discord...")
|
||||
await client.logout()
|
||||
await client.logout()
|
||||
|
||||
@ -8,6 +8,7 @@ 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
|
||||
|
||||
|
||||
@command(cmd="pk;system", subcommand=None, description="Shows information about your system.")
|
||||
async def this_system_info(conn, message, args):
|
||||
system = await db.get_system_by_account(conn, message.author.id)
|
||||
@ -18,6 +19,7 @@ async def this_system_info(conn, message, args):
|
||||
await client.send_message(message.channel, embed=await generate_system_info_card(conn, system))
|
||||
return True
|
||||
|
||||
|
||||
@command(cmd="pk;system", subcommand="new", usage="[name]", description="Registers a new system to this account.")
|
||||
async def new_system(conn, message, args):
|
||||
system = await db.get_system_by_account(conn, message.author.id)
|
||||
@ -39,12 +41,13 @@ async def new_system(conn, message, args):
|
||||
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 <name>`."
|
||||
|
||||
|
||||
@command(cmd="pk;system", subcommand="info", usage="[system]", description="Shows information about a system.")
|
||||
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:
|
||||
@ -53,14 +56,15 @@ async def system_info(conn, message, args):
|
||||
|
||||
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="pk;system", subcommand="name", usage="[name]", description="Renames your system. Leave blank to clear.")
|
||||
async def system_name(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."
|
||||
|
||||
@ -73,10 +77,11 @@ async def system_name(conn, message, args):
|
||||
await db.update_system_field(conn, system_id=system["id"], field="name", value=new_name)
|
||||
return True, "Name updated to {}.".format(new_name) if new_name else "Name cleared."
|
||||
|
||||
|
||||
@command(cmd="pk;system", subcommand="description", usage="[clear]", description="Updates your system description. Add \"clear\" to clear.")
|
||||
async def system_description(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."
|
||||
|
||||
@ -93,10 +98,11 @@ async def system_description(conn, message, args):
|
||||
await db.update_system_field(conn, system_id=system["id"], field="description", value=new_description)
|
||||
return True, "Description set." if new_description else "Description cleared."
|
||||
|
||||
|
||||
@command(cmd="pk;system", subcommand="tag", usage="[tag]", description="Updates your system tag. Leave blank to clear.")
|
||||
async def system_tag(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."
|
||||
|
||||
@ -110,24 +116,27 @@ async def system_tag(conn, message, args):
|
||||
members_exceeding = await db.get_members_exceeding(conn, system_id=system["id"], length=max_length - len(tag))
|
||||
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(tag, member_names))
|
||||
member_names = ", ".join([member["name"]
|
||||
for member in members_exceeding])
|
||||
logger.debug("Members exceeding combined length with tag '{}': {}".format(
|
||||
tag, 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)
|
||||
|
||||
async with conn.transaction():
|
||||
await db.update_system_field(conn, system_id=system["id"], field="tag", value=tag)
|
||||
|
||||
|
||||
return True, "Tag updated to {}.".format(tag) if tag else "Tag cleared."
|
||||
|
||||
|
||||
@command(cmd="pk;system", subcommand="remove", description="Removes your system ***permanently***.")
|
||||
async def system_remove(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 remove 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)
|
||||
if msg.content == system["hid"]:
|
||||
await db.remove_system(conn, system_id=system["id"])
|
||||
@ -139,7 +148,7 @@ async def system_remove(conn, message, args):
|
||||
@command(cmd="pk;system", subcommand="link", usage="<account>", description="Links another account to your system.")
|
||||
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."
|
||||
|
||||
@ -176,7 +185,7 @@ async def system_link(conn, message, args):
|
||||
@command(cmd="pk;system", subcommand="unlink", description="Unlinks your system from this account. There must be at least one other account linked.")
|
||||
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."
|
||||
|
||||
@ -184,15 +193,16 @@ async def system_unlink(conn, message, args):
|
||||
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="pk;member", subcommand="new", usage="<name>", description="Adds a new member to your system.")
|
||||
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."
|
||||
|
||||
@ -208,11 +218,13 @@ async def new_member(conn, message, args):
|
||||
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="pk;member", subcommand="info", description="Shows information about a system member.", system_only=False)
|
||||
async def member_info(conn, message, member, args):
|
||||
await client.send_message(message.channel, embed=await generate_member_info_card(conn, member))
|
||||
return True
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="color", usage="[color]", description="Updates a member's associated color. Leave blank to clear.")
|
||||
async def member_color(conn, message, member, args):
|
||||
if len(args) == 0:
|
||||
@ -227,7 +239,8 @@ async def member_color(conn, message, member, args):
|
||||
async with conn.transaction():
|
||||
await db.update_member_field(conn, member_id=member["id"], field="color", value=color)
|
||||
return True, "Color updated to #{}.".format(color) if color else "Color cleared."
|
||||
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="pronouns", usage="[pronouns]", description="Updates a member's pronouns. Leave blank to clear.")
|
||||
async def member_pronouns(conn, message, member, args):
|
||||
if len(args) == 0:
|
||||
@ -239,12 +252,13 @@ async def member_pronouns(conn, message, member, args):
|
||||
await db.update_member_field(conn, member_id=member["id"], field="pronouns", value=pronouns)
|
||||
return True, "Pronouns set to {}".format(pronouns) if pronouns else "Pronouns cleared."
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="birthdate", usage="[birthdate]", description="Updates a member's birthdate. Must be in ISO-8601 format (eg. 1999-07-25). Leave blank to clear.")
|
||||
async def member_birthday(conn, message, member, args):
|
||||
if len(args) == 0:
|
||||
new_date = None
|
||||
else:
|
||||
# Parse date
|
||||
# Parse date
|
||||
try:
|
||||
new_date = datetime.strptime(args[0], "%Y-%m-%d").date()
|
||||
except ValueError:
|
||||
@ -254,6 +268,7 @@ async def member_birthday(conn, message, member, args):
|
||||
await db.update_member_field(conn, member_id=member["id"], field="birthday", value=new_date)
|
||||
return True, "Birthdate set to {}".format(new_date) if new_date else "Birthdate cleared."
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="description", description="Updates a member's description. Add \"clear\" to clear.")
|
||||
async def member_description(conn, message, member, args):
|
||||
if len(args) > 0 and args[0] == "clear":
|
||||
@ -267,11 +282,12 @@ async def member_description(conn, message, member, args):
|
||||
async with conn.transaction():
|
||||
await db.update_member_field(conn, member_id=member["id"], field="description", value=new_description)
|
||||
return True, "Description set." if new_description else "Description cleared."
|
||||
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="remove", description="Removes a member from your system.")
|
||||
async def member_remove(conn, message, member, args):
|
||||
await client.send_message(message.channel, "Are you sure you want to remove {}? If so, reply to this message with the member's name.".format(member["name"]))
|
||||
|
||||
|
||||
msg = await client.wait_for_message(author=message.author, channel=message.channel)
|
||||
if msg.content == member["name"]:
|
||||
await db.delete_member(conn, member_id=member["id"])
|
||||
@ -279,6 +295,7 @@ async def member_remove(conn, message, member, args):
|
||||
else:
|
||||
return True, "Member removal cancelled."
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="avatar", usage="[user|url]", description="Updates a member's avatar. Can be an account mention (which will use that account's avatar), or a link to an image. Leave blank to clear.")
|
||||
async def member_avatar(conn, message, member, args):
|
||||
if len(args) == 0:
|
||||
@ -296,11 +313,11 @@ async def member_avatar(conn, message, member, args):
|
||||
avatar_url = args[0]
|
||||
else:
|
||||
return False, "Invalid URL."
|
||||
|
||||
|
||||
async with conn.transaction():
|
||||
await db.update_member_field(conn, member_id=member["id"], field="avatar_url", value=avatar_url)
|
||||
return True, "Avatar set." if avatar_url else "Avatar cleared."
|
||||
|
||||
|
||||
|
||||
@member_command(cmd="pk;member", subcommand="proxy", usage="[example]", description="Updates a member's proxy settings. Needs an \"example\" proxied message containing the string \"text\" (eg. [text], |text|, etc).")
|
||||
async def member_proxy(conn, message, member, args):
|
||||
@ -318,7 +335,8 @@ async def member_proxy(conn, message, member, args):
|
||||
# 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))
|
||||
logger.debug(
|
||||
"Matched prefix '{}' and suffix '{}'".format(prefix, suffix))
|
||||
|
||||
# DB stores empty strings as None, make that work
|
||||
if not prefix:
|
||||
@ -331,6 +349,7 @@ async def member_proxy(conn, message, member, args):
|
||||
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."
|
||||
|
||||
|
||||
@command(cmd="pk;message", subcommand=None, usage="<id>", description="Shows information about a proxied message. Requires the message ID.")
|
||||
async def message_info(conn, message, args):
|
||||
try:
|
||||
@ -357,14 +376,16 @@ async def message_info(conn, message, args):
|
||||
embed = discord.Embed()
|
||||
embed.timestamp = message.timestamp
|
||||
embed.colour = discord.Colour.blue()
|
||||
|
||||
|
||||
if system["name"]:
|
||||
system_value = "`{}`: {}".format(system["hid"], system["name"])
|
||||
else:
|
||||
system_value = "`{}`".format(system["hid"])
|
||||
embed.add_field(name="System", value=system_value)
|
||||
embed.add_field(name="Member", value="`{}`: {}".format(member["hid"], member["name"]))
|
||||
embed.add_field(name="Sent by", value="{}#{}".format(original_sender.name, original_sender.discriminator))
|
||||
embed.add_field(name="Member", value="`{}`: {}".format(
|
||||
member["hid"], member["name"]))
|
||||
embed.add_field(name="Sent by", value="{}#{}".format(
|
||||
original_sender.name, original_sender.discriminator))
|
||||
embed.add_field(name="Content", value=message.clean_content, inline=False)
|
||||
|
||||
embed.set_author(name=member["name"], url=member["avatar_url"])
|
||||
@ -372,12 +393,14 @@ async def message_info(conn, message, args):
|
||||
await client.send_message(message.channel, embed=embed)
|
||||
return True
|
||||
|
||||
|
||||
@command(cmd="pk;help", subcommand=None, usage="[system|member|message]", description="Shows this help message.")
|
||||
async def show_help(conn, message, args):
|
||||
embed = discord.Embed()
|
||||
embed.colour = discord.Colour.blue()
|
||||
embed.title = "PluralKit Help"
|
||||
embed.set_footer(text="<> denotes mandatory arguments, [] denotes optional arguments")
|
||||
embed.set_footer(
|
||||
text="<> denotes mandatory arguments, [] denotes optional arguments")
|
||||
|
||||
if len(args) > 0 and ("pk;" + args[0]) in command_map:
|
||||
cmds = ["", ("pk;" + args[0], command_map["pk;" + args[0]])]
|
||||
@ -386,7 +409,8 @@ async def show_help(conn, message, args):
|
||||
|
||||
for cmd, subcommands in cmds:
|
||||
for subcmd, (_, usage, description) in subcommands.items():
|
||||
embed.add_field(name="{} {} {}".format(cmd, subcmd or "", usage or ""), value=description, inline=False)
|
||||
|
||||
embed.add_field(name="{} {} {}".format(
|
||||
cmd, subcmd or "", usage or ""), value=description, inline=False)
|
||||
|
||||
await client.send_message(message.channel, embed=embed)
|
||||
return True
|
||||
return True
|
||||
|
||||
@ -5,6 +5,7 @@ import asyncpg.exceptions
|
||||
|
||||
from pluralkit.bot import logger
|
||||
|
||||
|
||||
async def connect():
|
||||
while True:
|
||||
try:
|
||||
@ -12,32 +13,40 @@ async def connect():
|
||||
except (ConnectionError, asyncpg.exceptions.CannotConnectNowError):
|
||||
pass
|
||||
|
||||
|
||||
def db_wrap(func):
|
||||
async def inner(*args, **kwargs):
|
||||
before = time.perf_counter()
|
||||
res = await func(*args, **kwargs)
|
||||
after = time.perf_counter()
|
||||
|
||||
|
||||
logger.debug(" - DB took {:.2f} ms".format((after - before) * 1000))
|
||||
return res
|
||||
return inner
|
||||
|
||||
|
||||
webhook_cache = {}
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def create_system(conn, system_name: str, system_hid: str):
|
||||
logger.debug("Creating system (name={}, hid={})".format(system_name, system_hid))
|
||||
logger.debug("Creating system (name={}, hid={})".format(
|
||||
system_name, system_hid))
|
||||
return await conn.fetchrow("insert into systems (name, hid) values ($1, $2) returning *", system_name, system_hid)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def remove_system(conn, system_id: int):
|
||||
logger.debug("Deleting system (id={})".format(system_id))
|
||||
await conn.execute("delete from systems where id = $1", system_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def create_member(conn, system_id: int, member_name: str, member_hid: str):
|
||||
logger.debug("Creating member (system={}, name={}, hid={})".format(system_id, member_name, member_hid))
|
||||
return await conn.fetchrow("insert into members (name, system, hid) values ($1, $2, $3) returning *", member_name, system_id, member_hid)
|
||||
logger.debug("Creating member (system={}, name={}, hid={})".format(
|
||||
system_id, member_name, member_hid))
|
||||
return await conn.fetchrow("insert into members (name, system, hid) values ($1, $2, $3) returning *", member_name, system_id, member_hid)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def delete_member(conn, member_id: int):
|
||||
@ -45,70 +54,90 @@ async def delete_member(conn, member_id: int):
|
||||
await conn.execute("update switches set member = null, member_del = true where member = $1", member_id)
|
||||
await conn.execute("delete from members where id = $1", member_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def link_account(conn, system_id: int, account_id: str):
|
||||
logger.debug("Linking account (account_id={}, system_id={})".format(account_id, system_id))
|
||||
logger.debug("Linking account (account_id={}, system_id={})".format(
|
||||
account_id, system_id))
|
||||
await conn.execute("insert into accounts (uid, system) values ($1, $2)", int(account_id), system_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def unlink_account(conn, system_id: int, account_id: str):
|
||||
logger.debug("Unlinking account (account_id={}, system_id={})".format(account_id, system_id))
|
||||
logger.debug("Unlinking account (account_id={}, system_id={})".format(
|
||||
account_id, system_id))
|
||||
await conn.execute("delete from accounts where uid = $1 and system = $2", int(account_id), system_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_linked_accounts(conn, system_id: int):
|
||||
return [row["uid"] for row in await conn.fetch("select uid from accounts where system = $1", system_id)]
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_system_by_account(conn, account_id: str):
|
||||
return await conn.fetchrow("select systems.* from systems, accounts where accounts.uid = $1 and accounts.system = systems.id", int(account_id))
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_system_by_hid(conn, system_hid: str):
|
||||
return await conn.fetchrow("select * from systems where hid = $1", system_hid)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_system(conn, system_id: int):
|
||||
return await conn.fetchrow("select * from systems where id = $1", system_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_member_by_name(conn, system_id: int, member_name: str):
|
||||
return await conn.fetchrow("select * from members where system = $1 and name = $2", system_id, member_name)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_member_by_hid_in_system(conn, system_id: int, member_hid: str):
|
||||
return await conn.fetchrow("select * from members where system = $1 and hid = $2", system_id, member_hid)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_member_by_hid(conn, member_hid: str):
|
||||
return await conn.fetchrow("select * from members where hid = $1", member_hid)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_member(conn, member_id: int):
|
||||
return await conn.fetchrow("select * from members where id = $1", member_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_message(conn, message_id: str):
|
||||
return await conn.fetchrow("select * from messages where mid = $1", message_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def update_system_field(conn, system_id: int, field: str, value):
|
||||
logger.debug("Updating system field (id={}, {}={})".format(system_id, field, value))
|
||||
logger.debug("Updating system field (id={}, {}={})".format(
|
||||
system_id, field, value))
|
||||
await conn.execute("update systems set {} = $1 where id = $2".format(field), value, system_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def update_member_field(conn, member_id: int, field: str, value):
|
||||
logger.debug("Updating member field (id={}, {}={})".format(member_id, field, value))
|
||||
logger.debug("Updating member field (id={}, {}={})".format(
|
||||
member_id, field, value))
|
||||
await conn.execute("update members set {} = $1 where id = $2".format(field), value, member_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_all_members(conn, system_id: int):
|
||||
return await conn.fetch("select * from members where system = $1", system_id)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_members_exceeding(conn, system_id: int, length: int):
|
||||
return await conn.fetch("select * from members where system = $1 and length(name) >= $2", system_id, length)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_webhook(conn, channel_id: str):
|
||||
if channel_id in webhook_cache:
|
||||
@ -117,30 +146,38 @@ async def get_webhook(conn, channel_id: str):
|
||||
webhook_cache[channel_id] = res
|
||||
return res
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def add_webhook(conn, channel_id: str, webhook_id: str, webhook_token: str):
|
||||
logger.debug("Adding new webhook (channel={}, webhook={}, token={})".format(channel_id, webhook_id, webhook_token))
|
||||
logger.debug("Adding new webhook (channel={}, webhook={}, token={})".format(
|
||||
channel_id, webhook_id, webhook_token))
|
||||
await conn.execute("insert into webhooks (channel, webhook, token) values ($1, $2, $3)", int(channel_id), int(webhook_id), webhook_token)
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def add_message(conn, message_id: str, channel_id: str, member_id: int, sender_id: str):
|
||||
logger.debug("Adding new message (id={}, channel={}, member={}, sender={})".format(message_id, channel_id, member_id, sender_id))
|
||||
logger.debug("Adding new message (id={}, channel={}, member={}, sender={})".format(
|
||||
message_id, channel_id, member_id, sender_id))
|
||||
await conn.execute("insert into messages (mid, channel, member, sender) values ($1, $2, $3, $4)", int(message_id), int(channel_id), member_id, int(sender_id))
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_members_by_account(conn, account_id: str):
|
||||
# Returns a "chimera" object
|
||||
return await conn.fetch("select members.id, members.hid, members.prefix, members.suffix, members.name, members.avatar_url, systems.tag from systems, members, accounts where accounts.uid = $1 and systems.id = accounts.system and members.system = systems.id", int(account_id))
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def get_message_by_sender_and_id(conn, message_id: str, sender_id: str):
|
||||
await conn.fetchrow("select * from messages where mid = $1 and sender = $2", int(message_id), int(sender_id))
|
||||
|
||||
|
||||
@db_wrap
|
||||
async def delete_message(conn, message_id: str):
|
||||
logger.debug("Deleting message (id={})".format(message_id))
|
||||
await conn.execute("delete from messages where mid = $1", int(message_id))
|
||||
|
||||
|
||||
async def create_tables(conn):
|
||||
await conn.execute("""create table if not exists systems (
|
||||
id serial primary key,
|
||||
@ -190,4 +227,4 @@ async def create_tables(conn):
|
||||
id bigint primary key,
|
||||
cmd_chans bigint[],
|
||||
proxy_chans bigint[]
|
||||
)""")
|
||||
)""")
|
||||
|
||||
@ -6,6 +6,7 @@ import aiohttp
|
||||
from pluralkit import db
|
||||
from pluralkit.bot import client, logger
|
||||
|
||||
|
||||
async def get_webhook(conn, channel):
|
||||
async with conn.transaction():
|
||||
# Try to find an existing webhook
|
||||
@ -14,7 +15,8 @@ async def get_webhook(conn, channel):
|
||||
if not hook_row:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
req_data = {"name": "PluralKit Proxy Webhook"}
|
||||
req_headers = {"Authorization": "Bot {}".format(os.environ["TOKEN"])}
|
||||
req_headers = {
|
||||
"Authorization": "Bot {}".format(os.environ["TOKEN"])}
|
||||
|
||||
async with session.post("https://discordapp.com/api/v6/channels/{}/webhooks".format(channel.id), json=req_data, headers=req_headers) as resp:
|
||||
data = await resp.json()
|
||||
@ -24,17 +26,19 @@ async def get_webhook(conn, channel):
|
||||
# Insert new hook into DB
|
||||
await db.add_webhook(conn, channel_id=channel.id, webhook_id=hook_id, webhook_token=token)
|
||||
return hook_id, token
|
||||
|
||||
|
||||
return hook_row["webhook"], hook_row["token"]
|
||||
|
||||
|
||||
async def proxy_message(conn, member, message, inner):
|
||||
logger.debug("Proxying message '{}' for member {}".format(inner, member["hid"]))
|
||||
logger.debug("Proxying message '{}' for member {}".format(
|
||||
inner, member["hid"]))
|
||||
# Delete the original message
|
||||
await client.delete_message(message)
|
||||
|
||||
# Get the webhook details
|
||||
hook_id, hook_token = await get_webhook(conn, message.channel)
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
req_data = {
|
||||
"username": "{} {}".format(member["name"], member["tag"] or "").strip(),
|
||||
"avatar_url": member["avatar_url"],
|
||||
@ -49,13 +53,15 @@ async def proxy_message(conn, member, message, inner):
|
||||
# Insert new message details into the DB
|
||||
await db.add_message(conn, message_id=resp_data["id"], channel_id=message.channel.id, member_id=member["id"], sender_id=message.author.id)
|
||||
|
||||
|
||||
async def handle_proxying(conn, message):
|
||||
# Big fat query to find every member associated with this account
|
||||
# Returned member object has a few more keys (system tag, for example)
|
||||
members = await db.get_members_by_account(conn, account_id=message.author.id)
|
||||
|
||||
# Sort by specificity (members with both prefix and suffix go higher)
|
||||
members = sorted(members, key=lambda x: int(bool(x["prefix"])) + int(bool(x["suffix"])), reverse=True)
|
||||
members = sorted(members, key=lambda x: int(
|
||||
bool(x["prefix"])) + int(bool(x["suffix"])), reverse=True)
|
||||
|
||||
msg = message.content
|
||||
for member in members:
|
||||
@ -71,14 +77,14 @@ async def handle_proxying(conn, message):
|
||||
if msg.startswith(prefix) and msg.endswith(suffix):
|
||||
# Extract the actual message contents sans tags
|
||||
if suffix:
|
||||
inner_message = message.content[len(prefix):-len(suffix)].strip()
|
||||
inner_message = message.content[len(
|
||||
prefix):-len(suffix)].strip()
|
||||
else:
|
||||
# Slicing to -0 breaks, don't do that
|
||||
inner_message = message.content[len(prefix):].strip()
|
||||
|
||||
await proxy_message(conn, member, message, inner_message)
|
||||
break
|
||||
|
||||
|
||||
|
||||
async def handle_reaction(conn, reaction, user):
|
||||
@ -90,4 +96,4 @@ async def handle_reaction(conn, reaction, user):
|
||||
if message:
|
||||
# If so, delete the message and remove it from the DB
|
||||
await db.delete_message(conn, message["mid"])
|
||||
await client.delete_message(reaction.message)
|
||||
await client.delete_message(reaction.message)
|
||||
|
||||
@ -9,9 +9,11 @@ import discord
|
||||
from pluralkit import db
|
||||
from pluralkit.bot import client, logger
|
||||
|
||||
|
||||
def generate_hid() -> str:
|
||||
return "".join(random.choices(string.ascii_lowercase, k=5))
|
||||
|
||||
|
||||
async def parse_mention(mention: str) -> discord.User:
|
||||
# First try matching mention format
|
||||
match = re.fullmatch("<@!?(\\d+)>", mention)
|
||||
@ -27,10 +29,11 @@ async def parse_mention(mention: str) -> discord.User:
|
||||
except (ValueError, discord.NotFound):
|
||||
return None
|
||||
|
||||
|
||||
async def get_system_fuzzy(conn, key) -> asyncpg.Record:
|
||||
if isinstance(key, discord.User):
|
||||
return await db.get_system_by_account(conn, account_id=key.id)
|
||||
|
||||
|
||||
if isinstance(key, str) and len(key) == 5:
|
||||
return await db.get_system_by_hid(conn, system_hid=key)
|
||||
|
||||
@ -39,7 +42,8 @@ async def get_system_fuzzy(conn, key) -> asyncpg.Record:
|
||||
if system:
|
||||
return system
|
||||
return None
|
||||
|
||||
|
||||
|
||||
async def get_member_fuzzy(conn, system_id: int, key: str, system_only=True) -> asyncpg.Record:
|
||||
# First search by hid
|
||||
if system_only:
|
||||
@ -60,6 +64,8 @@ command_map = {}
|
||||
# Command wrapper
|
||||
# Return True for success, return False for failure
|
||||
# Second parameter is the message it'll send. If just False, will print usage
|
||||
|
||||
|
||||
def command(cmd, subcommand, usage=None, description=None):
|
||||
def wrap(func):
|
||||
async def wrapper(conn, message, args):
|
||||
@ -70,12 +76,13 @@ def command(cmd, subcommand, usage=None, description=None):
|
||||
success, msg = res, None
|
||||
else:
|
||||
success, msg = res
|
||||
|
||||
|
||||
if not success and not msg:
|
||||
# Failure, no message, print usage
|
||||
usage_embed = discord.Embed()
|
||||
usage_embed.colour = discord.Colour.blue()
|
||||
usage_embed.add_field(name="Usage", value=usage, inline=False)
|
||||
usage_embed.add_field(
|
||||
name="Usage", value=usage, inline=False)
|
||||
|
||||
await client.send_message(message.channel, embed=usage_embed)
|
||||
elif not success:
|
||||
@ -103,6 +110,8 @@ def command(cmd, subcommand, usage=None, description=None):
|
||||
# Member command wrapper
|
||||
# Tries to find member by first argument
|
||||
# If system_only=False, allows members from other systems by hid
|
||||
|
||||
|
||||
def member_command(cmd, subcommand, usage=None, description=None, system_only=True):
|
||||
def wrap(func):
|
||||
async def wrapper(conn, message, args):
|
||||
@ -122,11 +131,12 @@ def member_command(cmd, subcommand, usage=None, description=None, system_only=Tr
|
||||
|
||||
if member is None:
|
||||
return False, "Can't find member \"{}\".".format(args[0])
|
||||
|
||||
|
||||
return await func(conn, message, member, args[1:])
|
||||
return command(cmd=cmd, subcommand=subcommand, usage=usage, description=description)(wrapper)
|
||||
return wrap
|
||||
|
||||
|
||||
async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Embed:
|
||||
card = discord.Embed()
|
||||
|
||||
@ -134,7 +144,8 @@ async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Emb
|
||||
card.title = system["name"]
|
||||
|
||||
if system["description"]:
|
||||
card.add_field(name="Description", value=system["description"], inline=False)
|
||||
card.add_field(name="Description",
|
||||
value=system["description"], inline=False)
|
||||
|
||||
if system["tag"]:
|
||||
card.add_field(name="Tag", value=system["tag"])
|
||||
@ -158,11 +169,13 @@ async def generate_system_info_card(conn, system: asyncpg.Record) -> discord.Emb
|
||||
member_texts.append("`{}`: {}".format(member["hid"], member["name"]))
|
||||
|
||||
if len(member_texts) > 0:
|
||||
card.add_field(name="Members", value="\n".join(member_texts), inline=False)
|
||||
card.add_field(name="Members", value="\n".join(
|
||||
member_texts), inline=False)
|
||||
|
||||
card.set_footer(text="System ID: {}".format(system["hid"]))
|
||||
return card
|
||||
|
||||
|
||||
async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Embed:
|
||||
card = discord.Embed()
|
||||
card.set_author(name=member["name"], icon_url=member["avatar_url"])
|
||||
@ -171,18 +184,21 @@ async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Emb
|
||||
card.colour = int(member["color"], 16)
|
||||
|
||||
if member["birthday"]:
|
||||
card.add_field(name="Birthdate", value=member["birthday"].strftime("%b %d, %Y"))
|
||||
|
||||
card.add_field(name="Birthdate",
|
||||
value=member["birthday"].strftime("%b %d, %Y"))
|
||||
|
||||
if member["pronouns"]:
|
||||
card.add_field(name="Pronouns", value=member["pronouns"])
|
||||
|
||||
if member["prefix"] or member["suffix"]:
|
||||
prefix = member["prefix"] or ""
|
||||
suffix = member["suffix"] or ""
|
||||
card.add_field(name="Proxy Tags", value="{}text{}".format(prefix, suffix))
|
||||
card.add_field(name="Proxy Tags",
|
||||
value="{}text{}".format(prefix, suffix))
|
||||
|
||||
if member["description"]:
|
||||
card.add_field(name="Description", value=member["description"], inline=False)
|
||||
card.add_field(name="Description",
|
||||
value=member["description"], inline=False)
|
||||
|
||||
# Get system name and hid
|
||||
system = await db.get_system(conn, system_id=member["system"])
|
||||
@ -192,9 +208,11 @@ async def generate_member_info_card(conn, member: asyncpg.Record) -> discord.Emb
|
||||
system_value = "`{}`".format(system["hid"])
|
||||
card.add_field(name="System", value=system_value, inline=False)
|
||||
|
||||
card.set_footer(text="System ID: {} | Member ID: {}".format(system["hid"], member["hid"]))
|
||||
card.set_footer(text="System ID: {} | Member ID: {}".format(
|
||||
system["hid"], member["hid"]))
|
||||
return card
|
||||
|
||||
|
||||
async def text_input(message, subject):
|
||||
await client.send_message(message.channel, "Reply in this channel with the new description you want to set for {}.".format(subject))
|
||||
msg = await client.wait_for_message(author=message.author, channel=message.channel)
|
||||
@ -208,4 +226,4 @@ async def text_input(message, subject):
|
||||
return msg.content
|
||||
else:
|
||||
await client.clear_reactions(msg)
|
||||
return None
|
||||
return None
|
||||
|
||||
Loading…
Reference in New Issue
Block a user