diff --git a/database/pet.png b/database/pet.png new file mode 100644 index 00000000..439cb94f Binary files /dev/null and b/database/pet.png differ diff --git a/database/sommie.png b/database/sommie.png new file mode 100644 index 00000000..f83062d3 Binary files /dev/null and b/database/sommie.png differ diff --git a/docker-compose.yaml b/docker-compose.yaml index cc63556e..2b8da00a 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -2,11 +2,31 @@ services: ninian: image: registry.froginata.com/ninian:${IMAGE_TAG} restart: unless-stopped + depends_on: + - database environment: + DATABASE_CONNECTION: "postgresql://ninian:${POSTGRES_PASSWORD}@database/bot_db" OWNER_ID: ${OWNER_ID} TOKEN: ${TOKEN} labels: - "com.centurylinklabs.watchtower.scope=ninian-${IMAGE_TAG}" + networks: + - nininet + + database: + image: postgres:17 + restart: unless-stopped + environment: + POSTGRES_DB: bot_db + POSTGRES_USER: ninian + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD} + networks: + - nininet + ports: + - 5432:5432 + volumes: + - db:/var/lib/postgresql/data + watchtower: image: containrrr/watchtower restart: unless-stopped @@ -19,3 +39,12 @@ services: - "com.centurylinklabs.watchtower.enable=false" volumes: - /var/run/docker.sock:/var/run/docker.sock + +networks: + nininet: + name: "nininet-${IMAGE_TAG}" + driver: bridge + +volumes: + db: + name: "ninian-db-${IMAGE_TAG}" diff --git a/main.py b/main.py index 87f7a8a4..ca016314 100644 --- a/main.py +++ b/main.py @@ -88,19 +88,9 @@ async def random_command(ctx): class SommieButton(maji.Button): def action(self, interaction): - user = oifey.user.get(-1) - - if not "sommie" in user: - user["sommie"] = 0 - - user["sommie"] += 1 - - oifey.user.update(-1, user) - + oifey.sql.User(interaction.user.id).pet_sommie() edit = self.parent.pages[1] - - edit.description = edit.description.format(user["sommie"]) - + edit.description = edit.description.format(oifey.sql.sommie_pets()) super().action(interaction) async def sommie(ctx): @@ -115,17 +105,6 @@ async def sommie(ctx): await multi.send(ctx) -@maji.commands.classic("sommie_set") -async def sommie_set(ctx): - if ctx.author.id != client.owner: - return - - user = oifey.user.get(-1) - - user["sommie"] = int(ctx.content) - - oifey.user.update(-1, user) - engage.special["sommie"] = sommie # ninian @@ -133,13 +112,10 @@ async def sommie_set(ctx): async def ninian(ctx): embed = maji.Embed() - embed.attach("image", "database/ninian.png") - return embed fe6.special["ninian"] = ninian - fe6.special["you"] = ninian # lorenz @@ -167,12 +143,10 @@ async def b2(ctx): b2 = oifey.modules["b2"] async def lorenz(ctx): - if ctx.content.lower().strip() == "b2 lorenz" or ctx.content.lower().strip() == 'lorenz': - embed = maji.Embed() - - embed.attach("image", "database/lorenz.gif") - + embed = maji.Embed() + embed.attach("image", "database/lorenz.gif") return embed + fe3.special["b2 lorenz"] = lorenz b2.special["lorenz"] = lorenz diff --git a/oifey/__init__.py b/oifey/__init__.py index 4a4f4cf7..f7ae8df0 100644 --- a/oifey/__init__.py +++ b/oifey/__init__.py @@ -4,5 +4,3 @@ import oifey.alias as alias modules = module.modules -user = sql.user -server = sql.server \ No newline at end of file diff --git a/oifey/alias.py b/oifey/alias.py index cb0c737d..5c3428fb 100644 --- a/oifey/alias.py +++ b/oifey/alias.py @@ -7,117 +7,78 @@ search_text = util.text.search_text -character_min = 3 -character_max = 30 +character_min = 1 +character_max = 64 -alias_limit = 30 +alias_limit = 100 # ADD -async def add_alias(ctx, table, name, og, alias): - row = table.get(name) - if not "alias" in row: row["alias"] = [] +async def add_alias(ctx, author_id, og, alias): + user = sql.User(author_id) - if len(row["alias"]) >= alias_limit: + if len(user.custom_aliases) >= alias_limit: embed = maji.Embed(title = "Alias limit reached!", desc = "Maybe try removing old ones?", color = embed_color) - await ctx.send(embed=embed) - return - - def remove(data): - for x in sql.banned_words: - data = data.replace(x, "") - - return data - - og = remove(og) - alias = remove(alias) def check(data): - clean = search_text(data, ignore_space = True) - - # character size limits - return len(clean) >= character_min and len(clean) <= character_max - + return character_min <= len(data) <= character_max + if check(og) and check(alias): - add = { - "og": {"clean": search_text(og, ignore_space = True), "display": og}, - "alias": {"clean": search_text(alias, ignore_space = True), "display": alias}, - } - - row["alias"].append(add) - - table.update(name, row) - - embed = maji.Embed(title = "Success", desc = "Alias added!", color = embed_color) - await ctx.send(embed=embed) - + try: + user.add_alias(alias, og) + embed = maji.Embed(title = "Success", desc = "Alias added!", color = embed_color) + except Exception as e: + print(e) + embed = maji.Embed(title="Error", desc="Your alias couldn't be added. Please report this to staff!", color=embed_color) else: - embed = maji.Embed(title = "Error", desc = f"Your alias couldn't be added. Please check if it is at least {character_min} characters, and most special characters are ignored!", color = embed_color) - - await ctx.send(embed=embed) + embed = maji.Embed(title = "Error", desc = f"Your alias couldn't be added. Please keep it under {character_max} characters long.", color = embed_color) + + await ctx.send(embed=embed) # SHOW -async def show_alias(ctx, table, name): - row = table.get(name) - row = row.get("alias") or [] +async def show_alias(ctx, author_id): + user = sql.User(author_id) - if not row: - embed = maji.Embed(title = "No aliases yet!", desc = "There's still no aliases here.\nIf you're looking at adding alias for yourself, use /alias user.\nIf you're looking at adding aliases for a server, only people that can manage messages can add/remove them, with the command being /alias server.", color = embed_color) - - await ctx.send(embed=embed) - + if not user.custom_aliases: + embed = maji.Embed(title = "No aliases yet!", desc = "There's still no aliases here.\nIf you're looking at adding alias for yourself, use /alias user.\nIf you're looking at adding aliases for a server, you can't, yet.", color = embed_color) else: text = "" - i = 1 - for value in row: - og = value["og"]["display"] - alias = value["alias"]["display"] - + for alias, og in user.custom_aliases: text += f"**{i}.** `{alias}` is an alias to `{og}`\n" - i += 1 - - embed = maji.Embed(title = f"Alias {len(row)}/{alias_limit}", desc = text, footer = "Type / if you want to add more or remove them!", color = embed_color) + embed = maji.Embed(title = f"Alias usage {len(user.custom_aliases)}/{alias_limit}", desc = text, footer = "Type / if you want to add more or remove them!", color = embed_color) - await ctx.send(embed=embed) + await ctx.send(embed=embed) # REMOVE -async def remove_alias(ctx, table, name, index = 0): - row = table.get(name) - if not "alias" in row: row["alias"] = [] - - alias = row["alias"] - - index = max(index - 1, 0) - - if alias and index < len(alias): - alias.pop(index) - - table.update(name, row) - - embed = maji.Embed(title = "Success", desc = "Alias removed!", color = embed_color) - - await ctx.send(embed=embed) - +async def remove_alias(ctx, author_id, alias): + user = sql.User(author_id) + if alias not in user.custom_aliases: + embed = maji.Embed(title="Error", desc="Could not find that alias to remove!", color=embed_color) else: - embed = maji.Embed(title = "Invalid index!", desc = "It's too high, or maybe you didn't even add any aliases yet.", color = embed_color) - - await ctx.send(embed=embed) + try: + user.remove_alias(alias) + embed = maji.Embed(title = "Success", desc = "Alias removed!", color = embed_color) + except Exception as e: + print(e) + embed = maji.Embed(title="Error", desc="Your alias couldn't be removed. Please report this to staff!", color=embed_color) + await ctx.send(embed=embed) # Command funcs # user async def user_add(ctx, options = {}): - await add_alias(ctx, sql.user, ctx.author.id, options["name"], options["alias"]) + await add_alias(ctx, ctx.author.id, options["name"], options["alias"]) -async def user_show(ctx, options = {}): - await show_alias(ctx, sql.user, ctx.author.id) +async def user_show(ctx): + await show_alias(ctx, ctx.author.id) async def user_remove(ctx, options = {}): - await remove_alias(ctx, sql.user, ctx.author.id, options["index"]) + await remove_alias(ctx, ctx.author.id, options["alias"]) # server +""" async def server_add(ctx, options = {}): if ctx.guild is None or not ctx.channel.permissions_for(ctx.author).manage_messages: return @@ -132,7 +93,7 @@ async def server_remove(ctx, options = {}): if ctx.guild is None or not ctx.channel.permissions_for(ctx.author).manage_messages: return await remove_alias(ctx, sql.server, ctx.guild.id, options["index"]) - +""" # Options add_options = [ @@ -152,8 +113,8 @@ async def server_remove(ctx, options = {}): remove_options = [ { - "name": "index", - "description": "Which alias to remove (you can see the index in the show command).", + "name": "alias", + "description": "Which alias to remove.", "type": 4, "required": True } @@ -161,11 +122,9 @@ async def server_remove(ctx, options = {}): # Add Commands # user -#maji.commands.add_slash(user_add, "add", desc="User | Adds an alias for something.", group="alias", subgroup="user", options=add_options) - -#maji.commands.add_slash(user_show, "show", desc="User | Shows all active aliases.", group="alias", subgroup="user") - -#maji.commands.add_slash(user_remove, "remove", desc="User | Removes an alias.", group="alias", subgroup="user", options=remove_options) +maji.commands.add_slash(user_add, "add", desc="User | Adds an alias for something.", group="alias", subgroup="user", options=add_options) +maji.commands.add_slash(user_show, "show", desc="User | Shows all active aliases.", group="alias", subgroup="user") +maji.commands.add_slash(user_remove, "remove", desc="User | Removes an alias.", group="alias", subgroup="user", options=remove_options) # server #maji.commands.add_slash(server_add, "add", desc="Server | Adds an alias for something.", group="alias", subgroup="server", options=add_options) diff --git a/oifey/pool.py b/oifey/pool.py index 3872d012..2b87dbac 100644 --- a/oifey/pool.py +++ b/oifey/pool.py @@ -235,56 +235,53 @@ def append_alt(k, v): def search(self, og_text, ctx = None) -> SearchResult: # if context exists to get aliases from if ctx: - text = util.text.search_text(og_text, ignore_space = True) - + text = util.text.search_text(og_text) + # organize aliases - aliases = sql.user.get(ctx.author.id).get("alias") or [] - + aliases = sql.User(ctx.author.id).custom_aliases + + """ if ctx.guild: guild = sql.server.get(ctx.guild.id).get("alias") or [] - + aliases = aliases + guild - - for value in aliases: # loop in order - original = value["og"]["clean"] #original - alias = value["alias"]["clean"] #alias - + """ + + for k, v in aliases: + alias = util.text.search_text(k) + original = util.text.search_text(v) + # find alias that is between whitespace or/and # has 0-2 numbers in the end regex = r"\b" + alias + r"\d{0,2}\b" - regex = re.search(regex, text) - + # ^ regex jank to find if the name is in the text - # if it finds it, replace it with the alias if regex: new = regex.group(0).replace(alias, original) - text = text[:regex.start()] + new + text[regex.end():] - - text = util.text.search_text(text) - + else: text = util.text.search_text(og_text) - + finds = [] - + # Pick a random entry - if text.startswith("random") and (not(ctx) or ctx.allow_random): + if text.startswith("random") and (not ctx or ctx.allow_random): finds.append(self.random_pick()) - + # try searching if it's a hard alt name elif text in self.hard: finds.append(self.hard[text]) - + elif text: for key, value in self.alt.items(): for name in value: if text in name: finds.append(self.items[key]) break - + return SearchResult(finds, self, og_text) def random_pick(self): diff --git a/oifey/sql.py b/oifey/sql.py index 24296a8b..a70e3129 100644 --- a/oifey/sql.py +++ b/oifey/sql.py @@ -1,67 +1,113 @@ -import sqlite3, json +from datetime import datetime, timedelta +import os +import psycopg +from psycopg.types import TypeInfo +from psycopg.types.hstore import register_hstore -import oifey.util as util +connection = psycopg.connect(os.getenv('DATABASE_CONNECTION')) +type_info = TypeInfo.fetch(connection, "hstore") +register_hstore(type_info, connection) +with connection.cursor() as cursor: + cursor.execute(""" + create table if not exists users ( + id bigint primary key not null, + sommie_pets int default 0, + favorite_hero varchar(64), + custom_aliases hstore default hstore(''), + banned_until timestamp, + ban_message text + );""") + connection.commit() -util.file.create_folder(".oifey") +def sommie_pets(): + with connection.cursor() as db: + db.execute("select sum(sommie_pets) from users;") + return db.fetchone()[0] -connection = sqlite3.connect(".oifey/lite.db") -cursor = connection.cursor() -banned_words = ['"', "'", "\\", "\n"] - -class Table: - def __init__(self, key): - self.key = key - - # try pulling from table - # if not, create one - +class User: + def __init__(self, user_id): try: - self.data = self.select() - - except sqlite3.OperationalError: - cursor.execute(f"CREATE TABLE {self.key}(name, data)") - - self.data = self.select() - - def select(self): - "Gets all of the data in the table and stores in a dict for easy use" - current = cursor.execute(f"SELECT * FROM '{self.key}'") - - result = {} - - for value in current.fetchall(): - result[value[0]] = json.loads(value[1]) + self.id = int(user_id) + with connection.cursor() as db: + db.execute(f"select * from users where id = {self.id};") + database_row = db.fetchone() + if database_row is None: # Never seen this person before. + db.execute(f"insert into users(id) values ({self.id});") + connection.commit() + self.sommie_pets = 0 + self.favorite_hero = None + self.custom_aliases = {} + self.banned_until = None + self.ban_message = None + else: # Unwrap a tuple fetched from the DB into a nice python object! + self.sommie_pets = database_row[1] + self.favorite_hero = database_row[2] + self.custom_aliases = database_row[3] or {} + self.banned_until = database_row[4] + self.ban_message = database_row[5] + except Exception as e: + print(f"ERROR: Couldn't make a User for ID {self.id}.\n{e}") + + # You may read from this class's variables all you want, but please only modify them using the below functions! + # Is one missing? Add it here! This way, we can keep the database synced and only updated in intentional ways. + def pet_sommie(self): + with connection.cursor() as db: + db.execute(f"update users set sommie_pets = sommie_pets + 1 where id = {self.id};") + connection.commit() + self.sommie_pets += 1 + + def set_favorite_hero(self, hero: str): + if self.favorite_hero == hero: + return + with connection.cursor() as db: + db.execute( + f"update users set favorite_hero = (%s::text) where id = {self.id};", + [hero] # This silly little thing should fight SQL injection. + ) # I trust self.id because in init we make sure it's an int. + connection.commit() + self.favorite_hero = hero + + def add_alias(self, alias: str, real_name: str): + if self.custom_aliases.get(alias) == real_name: + return + with connection.cursor() as db: + db.execute( + f"update users set custom_aliases = custom_aliases || (%s::hstore) where id = {self.id};", + [f'"{alias}"=>"{real_name}"'] # god it took forever to get the quotes and typing right + ) + connection.commit() + self.custom_aliases[alias] = real_name + + def remove_alias(self, alias: str): + if alias not in self.custom_aliases: + return + with connection.cursor() as db: + db.execute( + f"update users set custom_aliases = custom_aliases - (%s::text) where id = {self.id};", + [alias] + ) + connection.commit() + del self.custom_aliases[alias] + + def ban(self, duration=timedelta(weeks=2871), reason=None): # default: 400 years + ban_timestamp = datetime.now() + duration + with connection.cursor() as db: + db.execute( + f"update users set banned_until = (%s::timestamp), ban_message = (%s::text) where id = {self.id};", + [ban_timestamp, reason] + ) + connection.commit() + self.banned_until = ban_timestamp + self.ban_message = reason - return result - - def get(self, name): - if not name in self.data: - return {} - - else: - return self.data[name] - - def update(self, name, value): - text = json.dumps(value) - - if not name in self.data: - cursor.execute(f"INSERT INTO '{self.key}' VALUES ({name}, '{text}');") - - else: - cursor.execute(f"UPDATE '{self.key}' SET data = '{text}' WHERE name = {name}") - - connection.commit() + def unban(self): + with connection.cursor() as db: + db.execute(f"update users set banned_until = null where id = {self.id};") + connection.commit() + self.banned_until = None - self.data[name] = value - - def drop(self, name): - if name in self.data: self.data.pop(name) - - cursor.execute(f"DELETE FROM '{self.key}' WHERE name = {name}") - - connection.commit() - -# common used tables -user = Table("user") -server = Table("server") \ No newline at end of file + def is_banned(self) -> bool: + if self.banned_until is None: + return False + return self.banned_until > datetime.now() diff --git a/requirements.txt b/requirements.txt index 0bd64501..55547a67 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,6 @@ # these are needed to run the bot, but not to use the webscraper discord.py~=2.4 lupa +psycopg[binary] python-dotenv -unidecode \ No newline at end of file +unidecode