Fix gino bugs, reimplement ban and revoke cmds, enforce postgresql

This commit is contained in:
Jeremy Zhang 2018-07-10 19:06:54 +00:00
parent ac225c84f5
commit de94f87c45
4 changed files with 72 additions and 80 deletions

View File

@ -20,6 +20,7 @@ If you happen to have a copy of Ubuntu on your server, you may head onto our [An
# Database installation # Database installation
To set up the database for it to work with the webapp and the discordbot, one must use **alembic** to *migrate* their databases to the current database state. To do so, please follow these instructions. To set up the database for it to work with the webapp and the discordbot, one must use **alembic** to *migrate* their databases to the current database state. To do so, please follow these instructions.
**PostgreSQL supports proper indexing and suitable for Titan needs. For this reason, Titan only supports using a PostgreSQL database.**
1. Install alembic with **Python 3.5's pip** `pip install alembic` 1. Install alembic with **Python 3.5's pip** `pip install alembic`
2. Change your directory to the webapp where the alembic files are located `cd webapp` 2. Change your directory to the webapp where the alembic files are located `cd webapp`
3. Clone `alembic.example.ini` into your own `alembic.ini` file to find and edit the following line `sqlalchemy.url` to equal your database uri. [See here](http://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls) if you need help understanding how database uri works in SQLalchemy. 3. Clone `alembic.example.ini` into your own `alembic.ini` file to find and edit the following line `sqlalchemy.url` to equal your database uri. [See here](http://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls) if you need help understanding how database uri works in SQLalchemy.

View File

@ -23,8 +23,8 @@ sed -i '32s/.*/sqlalchemy.url = postgresql:\/\/\/titan/' ~/workspace/webapp/ale
alembic upgrade head alembic upgrade head
echo "[C9Setup] Setting database uri for discordbot/config.py" echo "[C9Setup] Setting database uri for discordbot/config.py"
#'database-uri': "mysql+psycopg2:///titan?client_encoding=utf8", #'database-uri': "postgresql:///titan",
sed -i "4s/.*/\'database-uri\': \"postgresql+psycopg2:\/\/\/titan?client_encoding=utf8\",/" ~/workspace/discordbot/config.py sed -i "4s/.*/\'database-uri\': \"postgresql:\/\/\/titan\",/" ~/workspace/discordbot/config.py
echo "[C9Setup] Setting database uri and app location for webapp/config.py" echo "[C9Setup] Setting database uri and app location for webapp/config.py"
sed -i "19s/.*/\'database-uri\': \"postgresql+psycopg2:\/\/\/titan?client_encoding=utf8\",/" ~/workspace/webapp/config.py sed -i "19s/.*/\'database-uri\': \"postgresql+psycopg2:\/\/\/titan?client_encoding=utf8\",/" ~/workspace/webapp/config.py

View File

@ -45,7 +45,7 @@ class Titan(discord.AutoShardedClient):
def run(self): def run(self):
try: try:
self.loop.run_until_complete(self.start(config["bot-token"])) self.loop.run_until_complete(self.start())
except discord.errors.LoginFailure: except discord.errors.LoginFailure:
print("Invalid bot token in config!") print("Invalid bot token in config!")
finally: finally:
@ -54,6 +54,10 @@ class Titan(discord.AutoShardedClient):
except Exception as e: except Exception as e:
print("Error in cleanup:", e) print("Error in cleanup:", e)
self.loop.close() self.loop.close()
async def start(self):
await self.database.connect(config["database-uri"])
await super().start(config["bot-token"])
async def on_ready(self): async def on_ready(self):
print('Titan [DiscordBot]') print('Titan [DiscordBot]')
@ -66,14 +70,6 @@ class Titan(discord.AutoShardedClient):
game = discord.Game(name="Embed your Discord server! Visit https://TitanEmbeds.com/") game = discord.Game(name="Embed your Discord server! Visit https://TitanEmbeds.com/")
await self.change_presence(status=discord.Status.online, activity=game) await self.change_presence(status=discord.Status.online, activity=game)
try:
await self.database.connect(config["database-uri"])
except Exception:
self.logger.error("Unable to connect to specified database!")
traceback.print_exc()
await self.logout()
return
self.discordBotsOrg = DiscordBotsOrg(self.user.id, config.get("discord-bots-org-token", None)) self.discordBotsOrg = DiscordBotsOrg(self.user.id, config.get("discord-bots-org-token", None))
self.botsDiscordPw = BotsDiscordPw(self.user.id, config.get("bots-discord-pw-token", None)) self.botsDiscordPw = BotsDiscordPw(self.user.id, config.get("bots-discord-pw-token", None))

View File

@ -1,6 +1,7 @@
from gino import Gino from gino import Gino
import json import json
import discord import discord
import datetime
db = Gino() db = Gino()
@ -21,19 +22,14 @@ class DatabaseInterface(object):
async def push_message(self, message): async def push_message(self, message):
if message.guild: if message.guild:
edit_ts = message.edited_at
if not edit_ts:
edit_ts = None
else:
edit_ts = str(edit_ts)
await Messages.create( await Messages.create(
message_id = int(message.id), message_id = int(message.id),
guild_id = int(message.guild.id), guild_id = int(message.guild.id),
channel_id = int(message.channel.id), channel_id = int(message.channel.id),
content = message.content, content = message.content,
author = json.dumps(get_message_author(message)), author = json.dumps(get_message_author(message)),
timestamp = str(message.created_at), timestamp = message.created_at,
edited_timestamp = edit_ts, edited_timestamp = message.edited_at,
mentions = json.dumps(get_message_mentions(message.mentions)), mentions = json.dumps(get_message_mentions(message.mentions)),
attachments = json.dumps(get_attachments_list(message.attachments)), attachments = json.dumps(get_attachments_list(message.attachments)),
embeds = json.dumps(get_embeds_list(message.embeds)) embeds = json.dumps(get_embeds_list(message.embeds))
@ -41,7 +37,7 @@ class DatabaseInterface(object):
async def update_message(self, message): async def update_message(self, message):
if message.guild: if message.guild:
await Messages.get(int(message.id)).update( await Messages.update.values(
content = message.content, content = message.content,
timestamp = message.created_at, timestamp = message.created_at,
edited_timestamp = message.edited_at, edited_timestamp = message.edited_at,
@ -49,11 +45,11 @@ class DatabaseInterface(object):
attachments = json.dumps(get_attachments_list(message.attachments)), attachments = json.dumps(get_attachments_list(message.attachments)),
embeds = json.dumps(get_embeds_list(message.embeds)), embeds = json.dumps(get_embeds_list(message.embeds)),
author = json.dumps(get_message_author(message)) author = json.dumps(get_message_author(message))
).apply() ).where(Messages.message_id == int(message.id)).gino.status()
async def delete_message(self, message): async def delete_message(self, message):
if message.guild: if message.guild:
await Messages.get(int(message.id)).delete() await Messages.delete.where(Messages.message_id == int(message.id)).gino.status()
async def update_guild(self, guild): async def update_guild(self, guild):
if guild.me.guild_permissions.manage_webhooks: if guild.me.guild_permissions.manage_webhooks:
@ -123,7 +119,7 @@ class DatabaseInterface(object):
guild_id = int(member.guild.id), guild_id = int(member.guild.id),
user_id = int(member.id), user_id = int(member.id),
username = member.name, username = member.name,
discriminator = member.discriminator, discriminator = int(member.discriminator),
nickname = member.nick, nickname = member.nick,
avatar = member.avatar, avatar = member.avatar,
active = active, active = active,
@ -134,23 +130,24 @@ class DatabaseInterface(object):
if len(dbmember) > 1: if len(dbmember) > 1:
for mem in dbmember[1:]: for mem in dbmember[1:]:
await mem.delete() await mem.delete()
dbmember = dbmember[0] dbmember = dbmember[0]
if dbmember.banned != banned or dbmember.active != active or dbmember.username != member.name or dbmember.discriminator != int(member.discriminator) or dbmember.nickname != member.nick or dbmember.avatar != member.avatar or set(json.loads(dbmember.roles)) != set(list_role_ids(member.roles)): if dbmember.banned != banned or dbmember.active != active or dbmember.username != member.name or dbmember.discriminator != int(member.discriminator) or dbmember.nickname != member.nick or dbmember.avatar != member.avatar or set(json.loads(dbmember.roles)) != set(list_role_ids(member.roles)):
await dbmember.update( await dbmember.update(
banned = banned, banned = banned,
active = active, active = active,
username = member.name, username = member.name,
discriminator = member.discriminator, discriminator = int(member.discriminator),
nickname = member.nick, nickname = member.nick,
avatar = member.avatar, avatar = member.avatar,
roles = json.dumps(list_role_ids(member.roles)) roles = json.dumps(list_role_ids(member.roles))
).apply() ).apply()
async def unban_server_user(self, user, server): async def unban_server_user(self, user, server):
await GuildMembers.query \ await GuildMembers.update.values(banned = False) \
.where(GuildMembers.guild_id == int(server.id)) \ .where(GuildMembers.guild_id == int(server.id)) \
.where(GuildMembers.user_id == int(user.id)) \ .where(GuildMembers.user_id == int(user.id)) \
.update(banned = False).apply() .gino.status()
async def flag_unactive_guild_members(self, guild_id, guild_members): async def flag_unactive_guild_members(self, guild_id, guild_members):
async with db.transaction(): async with db.transaction():
@ -174,7 +171,7 @@ class DatabaseInterface(object):
guild_id = int(guild_id), guild_id = int(guild_id),
user_id = int(usr.id), user_id = int(usr.id),
username = usr.name, username = usr.name,
discriminator = usr.discriminator, discriminator = int(usr.discriminator),
nickname = None, nickname = None,
avatar = usr.avatar, avatar = usr.avatar,
active = False, active = False,
@ -183,61 +180,59 @@ class DatabaseInterface(object):
) )
async def ban_unauth_user_by_query(self, guild_id, placer_id, username, discriminator): async def ban_unauth_user_by_query(self, guild_id, placer_id, username, discriminator):
self.bot.loop.run_in_executor(None, self._ban_unauth_user_by_query, guild_id, placer_id, username, discriminator) dbuser = None
if discriminator:
def _ban_unauth_user_by_query(self, guild_id, placer_id, username, discriminator): dbuser = await UnauthenticatedUsers.query \
with self.get_session() as session: .where(UnauthenticatedUsers.guild_id == int(guild_id)) \
dbuser = None .where(UnauthenticatedUsers.username.ilike("%" + username + "%")) \
if discriminator: .where(UnauthenticatedUsers.discriminator == discriminator) \
dbuser = session.query(UnauthenticatedUsers) \ .order_by(UnauthenticatedUsers.id.desc()).gino.first()
.filter(UnauthenticatedUsers.guild_id == int(guild_id)) \ else:
.filter(UnauthenticatedUsers.username.ilike("%" + username + "%")) \ dbuser = await UnauthenticatedUsers.query \
.filter(UnauthenticatedUsers.discriminator == discriminator) \ .where(UnauthenticatedUsers.guild_id == int(guild_id)) \
.order_by(UnauthenticatedUsers.id.desc()).first() .where(UnauthenticatedUsers.username.ilike("%" + username + "%")) \
else: .order_by(UnauthenticatedUsers.id.desc()).gino.first()
dbuser = session.query(UnauthenticatedUsers) \ if not dbuser:
.filter(UnauthenticatedUsers.guild_id == int(guild_id)) \ return "Ban error! Guest user cannot be found."
.filter(UnauthenticatedUsers.username.ilike("%" + username + "%")) \ dbban = await UnauthenticatedBans.query \
.order_by(UnauthenticatedUsers.id.desc()).first() .where(UnauthenticatedBans.guild_id == int(guild_id)) \
if not dbuser: .where(UnauthenticatedBans.last_username == dbuser.username) \
return "Ban error! Guest user cannot be found." .where(UnauthenticatedBans.last_discriminator == dbuser.discriminator).gino.first()
dbban = session.query(UnauthenticatedBans) \ if dbban is not None:
.filter(UnauthenticatedBans.guild_id == int(guild_id)) \ if dbban.lifter_id is None:
.filter(UnauthenticatedBans.last_username == dbuser.username) \ return "Ban error! Guest user, **{}#{}**, has already been banned.".format(dbban.last_username, dbban.last_discriminator)
.filter(UnauthenticatedBans.last_discriminator == dbuser.discriminator).first() await dbban.delete()
if dbban is not None: dbban = await UnauthenticatedBans.create(
if dbban.lifter_id is None: guild_id = int(guild_id),
return "Ban error! Guest user, **{}#{}**, has already been banned.".format(dbban.last_username, dbban.last_discriminator) ip_address = dbuser.ip_address,
session.delete(dbban) last_username = dbuser.username,
dbban = UnauthenticatedBans(int(guild_id), dbuser.ip_address, dbuser.username, dbuser.discriminator, "", int(placer_id)) last_discriminator = dbuser.discriminator,
session.add(dbban) timestamp = datetime.datetime.now(),
session.commit() reason = "",
return "Guest user, **{}#{}**, has successfully been added to the ban list!".format(dbban.last_username, dbban.last_discriminator) lifter_id = None,
placer_id = int(placer_id)
)
return "Guest user, **{}#{}**, has successfully been added to the ban list!".format(dbban.last_username, dbban.last_discriminator)
async def revoke_unauth_user_by_query(self, guild_id, username, discriminator): async def revoke_unauth_user_by_query(self, guild_id, username, discriminator):
self.bot.loop.run_in_executor(None, self._revoke_unauth_user_by_query, guild_id, username, discriminator) dbuser = None
if discriminator:
dbuser = await UnauthenticatedUsers.query \
.where(UnauthenticatedUsers.guild_id == int(guild_id)) \
.where(UnauthenticatedUsers.username.ilike("%" + username + "%")) \
.where(UnauthenticatedUsers.discriminator == discriminator) \
.order_by(UnauthenticatedUsers.id.desc()).gino.first()
else:
dbuser = await UnauthenticatedUsers.query \
.where(UnauthenticatedUsers.guild_id == int(guild_id)) \
.where(UnauthenticatedUsers.username.ilike("%" + username + "%")) \
.order_by(UnauthenticatedUsers.id.desc()).gino.first()
if not dbuser:
return "Kick error! Guest user cannot be found."
elif dbuser.revoked:
return "Kick error! Guest user **{}#{}** has already been kicked!".format(dbuser.username, dbuser.discriminator)
await dbuser.update(revoked = True).apply()
return "Successfully kicked **{}#{}**!".format(dbuser.username, dbuser.discriminator)
def _revoke_unauth_user_by_query(self, guild_id, username, discriminator):
with self.get_session() as session:
dbuser = None
if discriminator:
dbuser = session.query(UnauthenticatedUsers) \
.filter(UnauthenticatedUsers.guild_id == int(guild_id)) \
.filter(UnauthenticatedUsers.username.ilike("%" + username + "%")) \
.filter(UnauthenticatedUsers.discriminator == discriminator) \
.order_by(UnauthenticatedUsers.id.desc()).first()
else:
dbuser = session.query(UnauthenticatedUsers) \
.filter(UnauthenticatedUsers.guild_id == int(guild_id)) \
.filter(UnauthenticatedUsers.username.ilike("%" + username + "%")) \
.order_by(UnauthenticatedUsers.id.desc()).first()
if not dbuser:
return "Kick error! Guest user cannot be found."
elif dbuser.revoked:
return "Kick error! Guest user **{}#{}** has already been kicked!".format(dbuser.username, dbuser.discriminator)
dbuser.revoked = True
session.commit()
return "Successfully kicked **{}#{}**!".format(dbuser.username, dbuser.discriminator)
async def delete_all_messages_from_channel(self, channel_id): async def delete_all_messages_from_channel(self, channel_id):
await Messages.delete.where(Messages.channel_id == int(channel_id)).gino.status() await Messages.delete.where(Messages.channel_id == int(channel_id)).gino.status()