diff --git a/discordbot/titanembeds/bot.py b/discordbot/titanembeds/bot.py index 1c33940..ccf8bef 100644 --- a/discordbot/titanembeds/bot.py +++ b/discordbot/titanembeds/bot.py @@ -45,20 +45,87 @@ class Titan(discord.Client): print('------') await self.change_presence( - game=discord.Game(name="Get your own @ https://TitanEmbeds.tk/"), status=discord.Status.online + game=discord.Game(name="iFrame your server! Visit https://TitanEmbeds.tk/ today!"), status=discord.Status.online ) try: - await self.database.connect(config["database-uri"]) + await self.database.connect(config["database-uri"] + "?charset=utf8") except Exception: self.logger.error("Unable to connect to specified database!") traceback.print_exc() await self.logout() return + for server in self.servers: + await self.database.update_guild(server) + if server.large: + await request_offline_members(server) + server_bans = await self.get_bans(server) + for member in server.members: + banned = member.id in [u.id for u in server_bans] + await self.database.update_guild_member( + member, + True, + banned + ) + await self.database.flag_unactive_guild_members(server.id, server.members) + await self.database.flag_unactive_bans(server.id, server_bans) + await self.database.remove_unused_guilds(self.servers) + async def on_message(self, message): await self.database.push_message(message) # TODO: Will add command handler + ban/kick command async def on_message_edit(self, message_before, message_after): await self.database.update_message(message_after) + + async def on_message_delete(self, message): + await self.database.delete_message(message) + + async def on_server_join(self, guild): + await self.database.update_guild(guild) + + async def on_server_remove(self, guild): + await self.database.remove_guild(guild) + + async def on_server_update(self, guildbefore, guildafter): + await self.database.update_guild(guildafter) + + async def on_server_role_create(self, role): + if role.name == self.user.name and role.managed: + await asyncio.sleep(2) + await self.database.update_guild(role.server) + + async def on_server_role_delete(self, role): + if role.server.me not in role.server.members: + return + await self.database.update_guild(role.server) + + async def on_server_role_update(self, rolebefore, roleafter): + await self.database.update_guild(roleafter.server) + + async def on_channel_delete(self, channel): + await self.database.update_guild(channel.server) + + async def on_channel_create(self, channel): + await self.database.update_guild(channel.server) + + async def on_channel_update(self, channelbefore, channelafter): + await self.database.update_guild(channelafter.server) + + async def on_member_join(self, member): + await self.database.update_guild_member(member, active=True, banned=False) + + async def on_member_remove(self, member): + await self.database.update_guild_member(member, active=False, banned=False) + + async def on_member_update(self, memberbefore, memberafter): + await self.database.update_guild_member(memberafter) + + async def on_member_ban(self, member): + if role.server.me not in role.server.members: + return + await self.database.update_guild_member(member, active=False, banned=True) + + async def on_member_unban(self, member): + await self.database.update_guild_member(member, active=False, banned=False) diff --git a/discordbot/titanembeds/database/__init__.py b/discordbot/titanembeds/database/__init__.py index dfaa9e4..af01c5a 100644 --- a/discordbot/titanembeds/database/__init__.py +++ b/discordbot/titanembeds/database/__init__.py @@ -6,11 +6,13 @@ from sqlalchemy.orm import sessionmaker, Session from sqlalchemy.ext.declarative import declarative_base import json +import discord Base = declarative_base() from titanembeds.database.guilds import Guilds from titanembeds.database.messages import Messages +from titanembeds.database.guild_members import GuildMembers class DatabaseInterface(object): # Courtesy of https://github.com/SunDwarf/Jokusoramame @@ -22,7 +24,7 @@ class DatabaseInterface(object): async def connect(self, dburi): async with threadpool(): - self.engine = create_engine(dburi, pool_recycle=250) + self.engine = create_engine(dburi, pool_recycle=10) self._sessionmaker = sessionmaker(bind=self.engine, expire_on_commit=False) @contextmanager @@ -74,3 +76,165 @@ class DatabaseInterface(object): msg.mentions = json.dumps(message.mentions) msg.attachments = json.dumps(message.attachments) session.commit() + + async def delete_message(self, message): + if message.server: + async with threadpool(): + with self.get_session() as session: + msg = session.query(Messages) \ + .filter(Messages.guild_id == message.server.id) \ + .filter(Messages.channel_id == message.channel.id) \ + .filter(Messages.message_id == message.id).first() + if msg: + session.delete(msg) + session.commit() + + async def update_guild(self, guild): + async with threadpool(): + with self.get_session() as session: + gui = session.query(Guilds).filter(Guilds.guild_id == guild.id).first() + if not gui: + gui = Guilds( + guild.id, + guild.name, + json.dumps(self.get_roles_list(guild.roles)), + json.dumps(self.get_channels_list(guild.channels)), + guild.owner_id, + guild.icon + ) + session.add(gui) + else: + gui.name = guild.name + gui.roles = json.dumps(self.get_roles_list(guild.roles)) + gui.channels = json.dumps(self.get_channels_list(guild.channels)) + gui.owner_id = guild.owner_id + gui.icon = guild.icon + session.commit() + + def get_roles_list(self, guildroles): + roles = [] + for role in guildroles: + roles.append({ + "id": role.id, + "name": role.name, + "color": role.color.value, + "hoist": role.hoist, + "position": role.position, + }) + return roles + + def get_channels_list(self, guildchannels): + channels = [] + for channel in guildchannels: + if str(channel.type) == "text": + overwrites = [] + for target, overwrite in channel.overwrites: + if isinstance(target, discord.Role): + type = "role" + else: + type = "member" + allow, deny = overwrite.pair() + allow = allow.value + deny = deny.value + overwrites.append({ + "id": target.id, + "type": type, + "allow": allow, + "deny": deny, + }) + + channels.append({ + "id": channel.id, + "name": channel.name, + "topic": channel.topic, + "position": channel.position, + "type": str(channel.type), + "permission_overwrites": overwrites + }) + return channels + + async def remove_unused_guilds(self, guilds): + async with threadpool(): + with self.get_session() as session: + dbguilds = session.query(Guilds).all() + changed = False + for guild in dbguilds: + disguild = discord.utils.get(guilds, id=guild.guild_id) + if not disguild: + changed = True + session.delete(guild) + if changed: + session.commit() + + async def remove_guild(self, guild): + async with threadpool(): + with self.get_session() as session: + gui = session.query(Guilds).filter(Guilds.guild_id == guild.id).first() + if gui: + session.delete(gui) + session.commit() + + async def update_guild_member(self, member, active=True, banned=False): + async with threadpool(): + with self.get_session() as session: + dbmember = session.query(GuildMembers) \ + .filter(GuildMembers.guild_id == member.server.id) \ + .filter(GuildMembers.user_id == member.id).first() + if not dbmember: + dbmember = GuildMembers( + member.server.id, + member.id, + member.name, + member.discriminator, + member.nick, + member.avatar, + active, + banned, + json.dumps(self.list_role_ids(member.roles)) + ) + session.add(dbmember) + else: + dbmember.banned = banned + dbmember.active = active + dbmember.username = member.name + dbmember.discriminator = member.discriminator + dbmember.nick = member.nick + dbmember.avatar = member.avatar + dbmember.roles = json.dumps(self.list_role_ids(member.roles)) + session.commit() + + async def flag_unactive_guild_members(self, guild_id, guild_members): + async with threadpool(): + with self.get_session() as session: + changed = False + dbmembers = session.query(GuildMembers) \ + .filter(GuildMembers.guild_id == guild_id) \ + .filter(GuildMembers.active == True).all() + for member in dbmembers: + dismember = discord.utils.get(guild_members, id=member.user_id) + if not dismember: + changed = True + member.active = False + if changed: + session.commit() + + def list_role_ids(self, usr_roles): + ids = [] + for role in usr_roles: + ids.append(role.id) + return ids + + async def flag_unactive_bans(self, guild_id, guildbans): + async with threadpool(): + with self.get_session() as session: + changed = False + for usr in guildbans: + dbusr = session.query(GuildMembers) \ + .filter(GuildMembers.guild_id == guild_id) \ + .filter(GuildMembers.user_id == usr.id) \ + .filter(GuildMembers.active == False).first() + if dbusr: + changed = True + dbusr.banned = True + if changed: + session.commit() diff --git a/discordbot/titanembeds/database/guild_members.py b/discordbot/titanembeds/database/guild_members.py new file mode 100644 index 0000000..28d5532 --- /dev/null +++ b/discordbot/titanembeds/database/guild_members.py @@ -0,0 +1,28 @@ +from titanembeds.database import db, Base + +class GuildMembers(Base): + __tablename__ = "guild_members" + id = db.Column(db.Integer, primary_key=True) # Auto incremented id + guild_id = db.Column(db.String(255)) # Discord guild id + user_id = db.Column(db.String(255)) # Discord user id + username = db.Column(db.String(255)) # Name + discriminator = db.Column(db.Integer) # User discriminator + nickname = db.Column(db.String(255)) # User nickname + avatar = db.Column(db.String(255)) # The avatar str of the user + active = db.Column(db.Boolean()) # If the user is a member of the guild + banned = db.Column(db.Boolean()) # If the user is banned in the guild + roles = db.Column(db.Text()) # Member roles + + def __init__(self, guild_id, user_id, username, discriminator, nickname, avatar, active, banned, roles): + self.guild_id = guild_id + self.user_id = user_id + self.username = username + self.discriminator = discriminator + self.nickname = nickname + self.avatar = avatar + self.active = active + self.banned = banned + self.roles = roles + + def __repr__(self): + return ''.format(self.id, self.guild_id, self.user_id, self.username, self.discriminator) diff --git a/discordbot/titanembeds/database/guilds.py b/discordbot/titanembeds/database/guilds.py index d80e9f2..f540328 100644 --- a/discordbot/titanembeds/database/guilds.py +++ b/discordbot/titanembeds/database/guilds.py @@ -4,11 +4,21 @@ class Guilds(Base): __tablename__ = "guilds" id = db.Column(db.Integer, primary_key=True) # Auto incremented id guild_id = db.Column(db.String(255)) # Discord guild id + name = db.Column(db.String(255)) # Name unauth_users = db.Column(db.Boolean()) # If allowed unauth users + roles = db.Column(db.Text()) # Guild Roles + channels = db.Column(db.Text()) # Guild channels + owner_id = db.Column(db.String(255)) # Snowflake of the owner + icon = db.Column(db.String(255)) # The icon string, null if none - def __init__(self, guild_id): + def __init__(self, guild_id, name, roles, channels, owner_id, icon): self.guild_id = guild_id + self.name = name self.unauth_users = True # defaults to true + self.roles = roles + self.channels = channels + self.owner_id = owner_id + self.icon = icon def __repr__(self): - return ''.format(self.id, self.guild_id) \ No newline at end of file + return ''.format(self.id, self.guild_id)