Format code

This commit is contained in:
Ske
2018-07-12 00:49:02 +02:00
parent b81b768b04
commit 902c14ef65
6 changed files with 157 additions and 65 deletions

View File

@@ -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)