refactoring, userdata module

This commit is contained in:
osmarks 2021-02-09 17:17:19 +00:00
parent 8f25254204
commit 7b11e95542
6 changed files with 282 additions and 166 deletions

170
src/commands.py Normal file
View File

@ -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 <https://github.com/osmarks/autobotrobot> - 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: <https://discordapp.com/oauth2/authorize?&client_id=509849474647064576&scope=bot&permissions=68608>
""")
@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)))

View File

@ -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)
);
"""
]

View File

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

View File

@ -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 <https://github.com/osmarks/autobotrobot> - 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: <https://discordapp.com/oauth2/authorize?&client_id=509849474647064576&scope=bot&permissions=68608>
""")
@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

87
src/userdata.py Normal file
View File

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

View File

@ -261,5 +261,7 @@ extensions = (
"telephone",
"achievement",
"heavserver",
"voice"
"voice",
"commands",
"userdata"
)