From 7b11e95542f383bf817f91fb0b72c30ce9d19bd9 Mon Sep 17 00:00:00 2001 From: osmarks Date: Tue, 9 Feb 2021 17:17:19 +0000 Subject: [PATCH] refactoring, userdata module --- src/commands.py | 170 ++++++++++++++++++++++++++++++++++++++++++++++++ src/db.py | 9 +++ src/debug.py | 19 ++++-- src/main.py | 159 +------------------------------------------- src/userdata.py | 87 +++++++++++++++++++++++++ src/util.py | 4 +- 6 files changed, 282 insertions(+), 166 deletions(-) create mode 100644 src/commands.py create mode 100644 src/userdata.py diff --git a/src/commands.py b/src/commands.py new file mode 100644 index 0000000..cad0b67 --- /dev/null +++ b/src/commands.py @@ -0,0 +1,170 @@ +import subprocess +import asyncio +import argparse +import random +from numpy.random import default_rng +import re +import discord.ext.commands + +import tio +import util + +def setup(bot): + cleaner = discord.ext.commands.clean_content() + def clean(ctx, text): + return cleaner.convert(ctx, text) + + @bot.command(help="Gives you a random fortune as generated by `fortune`.") + async def fortune(ctx): + await ctx.send(subprocess.run(["fortune"], stdout=subprocess.PIPE, encoding="UTF-8").stdout) + + @bot.command(help="Generates an apioform type.") + async def apioform(ctx): + await ctx.send(util.apioform()) + + @bot.command(help="Says Pong.") + async def ping(ctx): + await ctx.send("Pong.") + + @bot.command(help="Deletes the specified target.", rest_is_raw=True) + async def delete(ctx, *, raw_target): + target = await clean(ctx, raw_target.strip().replace("\n", " ")) + if len(target) > 256: + await ctx.send(embed=util.error_embed("Deletion target must be max 256 chars")) + return + async with ctx.typing(): + await ctx.send(f"Deleting {target}...") + await asyncio.sleep(1) + await bot.database.execute("INSERT INTO deleted_items (timestamp, item) VALUES (?, ?)", (util.timestamp(), target)) + await bot.database.commit() + await ctx.send(f"Deleted {target} successfully.") + + @bot.command(help="View recently deleted things, optionally matching a filter.") + async def list_deleted(ctx, search=None): + acc = "Recently deleted:\n" + if search: acc = f"Recently deleted (matching {search}):\n" + csr = None + if search: + csr = bot.database.execute("SELECT * FROM deleted_items WHERE item LIKE ? ORDER BY timestamp DESC LIMIT 100", (f"%{search}%",)) + else: + csr = bot.database.execute("SELECT * FROM deleted_items ORDER BY timestamp DESC LIMIT 100") + async with csr as cursor: + async for row in cursor: + to_add = "- " + row[2].replace("```", "[REDACTED]") + "\n" + if len(acc + to_add) > 2000: + break + acc += to_add + await ctx.send(acc) + + # Python, for some *very intelligent reason*, makes the default ArgumentParser exit the program on error. + # This is obviously undesirable behavior in a Discord bot, so we override this. + class NonExitingArgumentParser(argparse.ArgumentParser): + def exit(self, status=0, message=None): + if status: + raise Exception(f'Flag parse error: {message}') + exit(status) + + EXEC_REGEX = "^(.*)```([a-zA-Z0-9_\\-+]+)?\n(.*)```$" + + exec_flag_parser = NonExitingArgumentParser(add_help=False) + exec_flag_parser.add_argument("--verbose", "-v", action="store_true") + exec_flag_parser.add_argument("--language", "-L") + + @bot.command(rest_is_raw=True, help="Execute provided code (in a codeblock) using TIO.run.") + async def exec(ctx, *, arg): + match = re.match(EXEC_REGEX, arg, flags=re.DOTALL) + if match == None: + await ctx.send(embed=util.error_embed("Invalid format. Expected a codeblock.")) + return + flags_raw = match.group(1) + flags = exec_flag_parser.parse_args(flags_raw.split()) + lang = flags.language or match.group(2) + if not lang: + await ctx.send(embed=util.error_embed("No language specified. Use the -L flag or add a language to your codeblock.")) + return + lang = lang.strip() + code = match.group(3) + + async with ctx.typing(): + ok, real_lang, result, debug = await tio.run(lang, code) + if not ok: + await ctx.send(embed=util.error_embed(util.gen_codeblock(result), "Execution failed")) + else: + out = result + if flags.verbose: + debug_block = "\n" + util.gen_codeblock(f"""{debug}\nLanguage: {real_lang}""") + out = out[:2000 - len(debug_block)] + debug_block + else: + out = out[:2000] + await ctx.send(out) + + @bot.command(help="List supported languages, optionally matching a filter.") + async def supported_langs(ctx, search=None): + langs = sorted(tio.languages()) + acc = "" + for lang in langs: + if len(acc + lang) > 2000: + await ctx.send(acc) + acc = "" + if search == None or search in lang: acc += lang + " " + if acc == "": acc = "No results." + await ctx.send(acc) + + @bot.command(help="Get some information about the bot.", aliases=["invite"]) + async def about(ctx): + await ctx.send("""**AutoBotRobot: The least useful Discord bot ever designed.** + AutoBotRobot has many features, but not necessarily any practical ones. + It can execute code via TIO.run, do reminders, print fortunes, and not any more! + AutoBotRobot is open source - the code is available at - and you could run your own instance if you wanted to and could get around the complete lack of user guide or documentation. + You can also invite it to your server: + """) + + @bot.command(help="Roll simulated dice (basic NdX syntax only, no + etc., N <= 50, X <= 1000000).") + async def roll(ctx, dice): + match = re.match("([-0-9]*)d([0-9]+)", dice) + if not match: raise ValueError("Invalid dice notation") + n, x = match.groups() + if n == "": n = 1 + n, x = int(n), int(x) + if n > 50 or x > 1e6: raise ValueError("N or X exceeds limit") + rolls = [ random.randint(1, x) for _ in range(n) ] + await ctx.send(f"{sum(rolls)} ({' '.join(map(str, sorted(rolls)))})") + + def weight(thing): + lthing = thing.lower() + weight = 1.0 + if lthing == "c": weight *= 0.3 + for bad_thing in util.config["autobias"]["bad_things"]: + if bad_thing in lthing: weight *= 0.5 + for good_thing in util.config["autobias"]["good_things"]: + if good_thing in lthing: weight *= 2.0 + for negation in util.config["autobias"]["negations"]: + for _ in range(lthing.count(negation)): weight = 1 / weight + return weight + + rng = default_rng() + + @bot.command(help="'Randomly' choose between the specified options.", name="choice", aliases=["choose"]) + async def random_choice(ctx, *choices): + choices = list(choices) + samples = 1 + # apparently doing typing.Optional[int] doesn't work properly with this, so just bodge around it + try: + samples = int(choices[0]) + choices.pop(0) + except: pass + + if samples > 9223372036854775807 or samples < 1 or len(choices) < 1: + await ctx.send("No.") + return + + # because of python weirdness, using sum() on the bare map iterator consumes it, which means we have to actually make a list + weights = list(map(weight, choices)) + + if samples == 1: return await ctx.send(random.choices(choices, weights=weights, k=1)[0]) + + total = sum(weights) + probabilities = list(map(lambda x: x / total, weights)) + results = map(lambda t: (choices[t[0]], t[1]), enumerate(rng.multinomial(samples, list(probabilities)))) + + await ctx.send("\n".join(map(lambda x: f"{x[0]} x{x[1]}", results))) \ No newline at end of file diff --git a/src/db.py b/src/db.py index 7719bfd..6570ac2 100644 --- a/src/db.py +++ b/src/db.py @@ -67,6 +67,15 @@ CREATE TABLE assets ( identifier TEXT PRIMARY KEY, url TEXT NOT NULL ); +""", +""" +CREATE TABLE user_data ( + user_id INTEGER NOT NULL, + guild_id INTEGER, + key TEXT NOT NULL, + value TEXT NOT NULL, + UNIQUE (user_id, guild_id, key) +); """ ] diff --git a/src/debug.py b/src/debug.py index 280da2e..d1da8da 100644 --- a/src/debug.py +++ b/src/debug.py @@ -6,14 +6,15 @@ from discord.ext import commands import util def setup(bot): - @bot.group() + @bot.group(help="Debug/random messing around utilities. Owner-only.") @commands.check(util.admin_check) async def magic(ctx): if ctx.invoked_subcommand == None: return await ctx.send("Invalid magic command.") - @magic.command(rest_is_raw=True) + @magic.command(rest_is_raw=True, brief="Execute Python.") async def py(ctx, *, code): + "Executes Python. You may supply a codeblock. Comments in the form #timeout:([0-9]+) will be used as a timeout specifier. React with :x: to stop, probably." timeout = 5.0 timeout_match = re.search("#timeout:([0-9]+)", code, re.IGNORECASE) if timeout_match: @@ -49,8 +50,9 @@ def setup(bot): except BaseException as e: await ctx.send(embed=util.error_embed(util.gen_codeblock(traceback.format_exc()))) - @magic.command(rest_is_raw=True) + @magic.command(rest_is_raw=True, help="Execute SQL code against the database.") async def sql(ctx, *, code): + "Executes SQL (and commits). You may use a codeblock, similarly to with py." code = util.extract_codeblock(code) try: csr = bot.database.execute(code) @@ -63,11 +65,14 @@ def setup(bot): except Exception as e: await ctx.send(embed=util.error_embed(util.gen_codeblock(traceback.format_exc()))) - @magic.command() + @magic.command(help="Reload configuration file.") async def reload_config(ctx): util.load_config() ctx.send("Done!") - @magic.command() - async def reload_ext(ctx): - for ext in util.extensions: bot.reload_extension(ext) \ No newline at end of file + @magic.command(help="Reload extensions (all or the specified one).") + async def reload_ext(ctx, ext="all"): + if ext == "all": + for ext in util.extensions: bot.reload_extension(ext) + else: bot.reload_extension(ext) + await ctx.send("Done!") \ No newline at end of file diff --git a/src/main.py b/src/main.py index 6431d75..29e0509 100644 --- a/src/main.py +++ b/src/main.py @@ -1,21 +1,17 @@ import discord import toml import logging -import subprocess import discord.ext.commands as commands import discord.ext.tasks as tasks import re import asyncio import json -import argparse import traceback import random -import rolldice import collections import prometheus_client import prometheus_async.aio import typing -from numpy.random import default_rng import tio import db @@ -33,10 +29,6 @@ bot = commands.Bot(command_prefix=config["prefix"], description="AutoBotRobot, t case_insensitive=True, allowed_mentions=discord.AllowedMentions(everyone=False, users=True, roles=True)) bot._skip_check = lambda x, y: False -cleaner = discord.ext.commands.clean_content() -def clean(ctx, text): - return cleaner.convert(ctx, text) - messages = prometheus_client.Counter("abr_messages", "Messages seen/handled by bot") command_invocations = prometheus_client.Counter("abr_command_invocations", "Total commands invoked (includes failed)") @bot.event @@ -57,7 +49,7 @@ command_errors = prometheus_client.Counter("abr_errors", "Count of errors encoun async def on_command_error(ctx, err): #print(ctx, err) if isinstance(err, (commands.CommandNotFound, commands.CheckFailure)): return - if isinstance(err, commands.MissingRequiredArgument): return await ctx.send(embed=util.error_embed(str(err))) + if isinstance(err, (commands.MissingRequiredArgument, ValueError)): return await ctx.send(embed=util.error_embed(str(err))) try: command_errors.inc() trace = re.sub("\n\n+", "\n", "\n".join(traceback.format_exception(err, err, err.__traceback__))) @@ -66,155 +58,6 @@ async def on_command_error(ctx, err): await achievement.achieve(ctx.bot, ctx.message, "error") except Exception as e: print("meta-error:", e) -@bot.command(help="Gives you a random fortune as generated by `fortune`.") -async def fortune(ctx): - await ctx.send(subprocess.run(["fortune"], stdout=subprocess.PIPE, encoding="UTF-8").stdout) - -@bot.command(help="Generates an apioform type.") -async def apioform(ctx): - await ctx.send(util.apioform()) - -@bot.command(help="Says Pong.") -async def ping(ctx): - await ctx.send("Pong.") - -@bot.command(help="Deletes the specified target.", rest_is_raw=True) -async def delete(ctx, *, raw_target): - target = await clean(ctx, raw_target.strip().replace("\n", " ")) - if len(target) > 256: - await ctx.send(embed=util.error_embed("Deletion target must be max 256 chars")) - return - async with ctx.typing(): - await ctx.send(f"Deleting {target}...") - await asyncio.sleep(1) - await bot.database.execute("INSERT INTO deleted_items (timestamp, item) VALUES (?, ?)", (util.timestamp(), target)) - await bot.database.commit() - await ctx.send(f"Deleted {target} successfully.") - -@bot.command(help="View recently deleted things, optionally matching a filter.") -async def list_deleted(ctx, search=None): - acc = "Recently deleted:\n" - if search: acc = f"Recently deleted (matching {search}):\n" - csr = None - if search: - csr = bot.database.execute("SELECT * FROM deleted_items WHERE item LIKE ? ORDER BY timestamp DESC LIMIT 100", (f"%{search}%",)) - else: - csr = bot.database.execute("SELECT * FROM deleted_items ORDER BY timestamp DESC LIMIT 100") - async with csr as cursor: - async for row in cursor: - to_add = "- " + row[2].replace("```", "[REDACTED]") + "\n" - if len(acc + to_add) > 2000: - break - acc += to_add - await ctx.send(acc) - -# Python, for some *very intelligent reason*, makes the default ArgumentParser exit the program on error. -# This is obviously undesirable behavior in a Discord bot, so we override this. -class NonExitingArgumentParser(argparse.ArgumentParser): - def exit(self, status=0, message=None): - if status: - raise Exception(f'Flag parse error: {message}') - exit(status) - -EXEC_REGEX = "^(.*)```([a-zA-Z0-9_\\-+]+)?\n(.*)```$" - -exec_flag_parser = NonExitingArgumentParser(add_help=False) -exec_flag_parser.add_argument("--verbose", "-v", action="store_true") -exec_flag_parser.add_argument("--language", "-L") - -@bot.command(rest_is_raw=True, help="Execute provided code (in a codeblock) using TIO.run.") -async def exec(ctx, *, arg): - match = re.match(EXEC_REGEX, arg, flags=re.DOTALL) - if match == None: - await ctx.send(embed=util.error_embed("Invalid format. Expected a codeblock.")) - return - flags_raw = match.group(1) - flags = exec_flag_parser.parse_args(flags_raw.split()) - lang = flags.language or match.group(2) - if not lang: - await ctx.send(embed=util.error_embed("No language specified. Use the -L flag or add a language to your codeblock.")) - return - lang = lang.strip() - code = match.group(3) - - async with ctx.typing(): - ok, real_lang, result, debug = await tio.run(lang, code) - if not ok: - await ctx.send(embed=util.error_embed(util.gen_codeblock(result), "Execution failed")) - else: - out = result - if flags.verbose: - debug_block = "\n" + util.gen_codeblock(f"""{debug}\nLanguage: {real_lang}""") - out = out[:2000 - len(debug_block)] + debug_block - else: - out = out[:2000] - await ctx.send(out) - -@bot.command(help="List supported languages, optionally matching a filter.") -async def supported_langs(ctx, search=None): - langs = sorted(tio.languages()) - acc = "" - for lang in langs: - if len(acc + lang) > 2000: - await ctx.send(acc) - acc = "" - if search == None or search in lang: acc += lang + " " - if acc == "": acc = "No results." - await ctx.send(acc) - -@bot.command(help="Get some information about the bot.", aliases=["invite"]) -async def about(ctx): - await ctx.send("""**AutoBotRobot: The least useful Discord bot ever designed.** -AutoBotRobot has many features, but not necessarily any practical ones. -It can execute code via TIO.run, do reminders, print fortunes, and not any more! -AutoBotRobot is open source - the code is available at - and you could run your own instance if you wanted to and could get around the complete lack of user guide or documentation. -You can also invite it to your server: -""") - -@bot.command(help="Randomly generate an integer using dice syntax.", name="random", rest_is_raw=True) -async def random_int(ctx, *, dice): - await ctx.send("Disabled until CPU use restrictions can be done.") - #await ctx.send(rolldice.roll_dice(dice)[0]) - -def weight(thing): - lthing = thing.lower() - weight = 1.0 - if lthing == "c": weight *= 0.3 - for bad_thing in util.config["autobias"]["bad_things"]: - if bad_thing in lthing: weight *= 0.5 - for good_thing in util.config["autobias"]["good_things"]: - if good_thing in lthing: weight *= 2.0 - for negation in util.config["autobias"]["negations"]: - for _ in range(lthing.count(negation)): weight = 1 / weight - return weight - -rng = default_rng() - -@bot.command(help="'Randomly' choose between the specified options.", name="choice", aliases=["choose"]) -async def random_choice(ctx, *choices): - choices = list(choices) - samples = 1 - # apparently doing typing.Optional[int] doesn't work properly with this, so just bodge around it - try: - samples = int(choices[0]) - choices.pop(0) - except: pass - - if samples > 9223372036854775807 or samples < 1 or len(choices) < 1: - await ctx.send("No.") - return - - # because of python weirdness, using sum() on the bare map iterator consumes it, which means we have to actually make a list - weights = list(map(weight, choices)) - - if samples == 1: return await ctx.send(random.choices(choices, weights=weights, k=1)[0]) - - total = sum(weights) - probabilities = list(map(lambda x: x / total, weights)) - results = map(lambda t: (choices[t[0]], t[1]), enumerate(rng.multinomial(samples, list(probabilities)))) - - await ctx.send("\n".join(map(lambda x: f"{x[0]} x{x[1]}", results))) - @bot.check async def andrew_bad(ctx): return ctx.message.author.id != 543131534685765673 diff --git a/src/userdata.py b/src/userdata.py new file mode 100644 index 0000000..b26c40d --- /dev/null +++ b/src/userdata.py @@ -0,0 +1,87 @@ +import util + +def setup(bot): + @bot.group(name="userdata", aliases=["data"], brief="Store guild/user-localized data AND retrieve it later!") + async def userdata(ctx): pass + + async def get_userdata(db, user, guild, key): + return (await db.execute_fetchone("SELECT * FROM user_data WHERE user_id = ? AND guild_id = ? AND key = ?", (user, guild, key)) + or await db.execute_fetchone("SELECT * FROM user_data WHERE user_id = ? AND guild_id IS NULL AND key = ?", (user, key))) + async def set_userdata(db, user, guild, key, value): + await db.execute("INSERT OR REPLACE INTO user_data VALUES (?, ?, ?, ?)", (user, guild, key, value)) + await bot.database.commit() + + @userdata.command(help="Get a userdata key. Checks guild first, then global.") + async def get(ctx, *, key): + no_header = False + if key.startswith("noheader "): + key = key[9:] + no_header = True + row = await get_userdata(bot.database, ctx.author.id, ctx.guild.id, key) + if not row: + raise ValueError("No such key") + if no_header: + await ctx.send(row["value"]) + else: + await ctx.send(f"**{key}**: {row['value']}") + + @userdata.command(name="list", help="List userdata keys in a given scope matching a query. Can also show associated values.") + async def list_cmd(ctx, query="%", scope="guild", show_values: bool = False): + if scope == "global": + rows = await bot.database.execute_fetchall("SELECT * FROM user_data WHERE user_id = ? AND guild_id IS NULL AND key LIKE ?", (ctx.author.id, query)) + else: + rows = await bot.database.execute_fetchall("SELECT * FROM user_data WHERE user_id = ? AND guild_id = ? AND key LIKE ?", (ctx.author.id, ctx.guild.id, query)) + out = [] + for row in rows: + if show_values: + out.append(f"**{row['key']}**: {row['value']}") + else: + out.append(row["key"]) + if len(out) == 0: return await ctx.send("No data") + await ctx.send(("\n" if show_values else " ").join(out)[:2000]) # TODO: split better + + def check_key(key): + if len(key) > 128: raise ValueError("Key too long") + + def preprocess_value(value): + value = value.replace("\n", "").strip() + if len(value) > 256: raise ValueError("Value too long") + return value + + @userdata.command(name="set", help="Set a userdata key in the guild scope.") + async def set_cmd(ctx, key, *, value): + check_key(key) + value = preprocess_value(value) + await set_userdata(bot.database, ctx.author.id, ctx.guild.id, key, value) + await ctx.send(f"**{key}** set (scope guild)") + + @userdata.command(help="Set a userdata key in the global scope.") + async def set_global(ctx, key, *, value): + check_key(key) + value = preprocess_value(value) + await set_userdata(bot.database, ctx.author.id, None, key, value) + await ctx.send(f"**{key}** set (scope global)") + + @userdata.command() + async def inc(ctx, key, by: int = 1): + check_key(key) + row = await get_userdata(bot.database, ctx.author.id, ctx.guild.id, key) + if not row: + value = 0 + guild = ctx.guild.id + else: + value = int(row["value"]) + guild = row["guild_id"] + new_value = value + by + await set_userdata(bot.database, ctx.author.id, guild, key, str(new_value)) + await ctx.send(f"**{key}** set to {new_value}") + + @userdata.command() + async def delete(ctx, *keys): + for key in keys: + row = await get_userdata(bot.database, ctx.author.id, ctx.guild.id, key) + if not row: + return await ctx.send(embed=util.error_embed(f"No such key {key}")) + await bot.database.execute("DELETE FROM user_data WHERE user_id = ? AND guild_id = ? AND key = ?", (ctx.author.id, row["guild_id"], key)) + await bot.database.commit() + await ctx.send(f"**{key}** deleted") \ No newline at end of file diff --git a/src/util.py b/src/util.py index 01eee85..bd1feb4 100644 --- a/src/util.py +++ b/src/util.py @@ -261,5 +261,7 @@ extensions = ( "telephone", "achievement", "heavserver", - "voice" + "voice", + "commands", + "userdata" ) \ No newline at end of file