diff options
author | Mango0x45 <thomasvoss@live.com> | 2021-01-04 23:12:01 +0000 |
---|---|---|
committer | Mango0x45 <thomasvoss@live.com> | 2021-01-04 23:12:01 +0000 |
commit | dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9 (patch) | |
tree | 6b796516bf79900958df23f50a2a1e29c3b862e3 /cogs/utils.py | |
parent | b27570db9eb993bc74e47eab168f55a071b62870 (diff) | |
download | steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.tar steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.tar.gz steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.tar.bz2 steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.tar.lz steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.tar.xz steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.tar.zst steve-bot-dc4b9ce16987f59ade2b6e030f1bc27cf911a7a9.zip |
Fix security vulnerability in math and add timeout
Diffstat (limited to 'cogs/utils.py')
-rwxr-xr-x | cogs/utils.py | 36 |
1 files changed, 22 insertions, 14 deletions
diff --git a/cogs/utils.py b/cogs/utils.py index fdc7a9f..1c1a758 100755 --- a/cogs/utils.py +++ b/cogs/utils.py @@ -2,6 +2,7 @@ import asyncio import datetime import functools import json +import os import subprocess from collections import namedtuple from datetime import timedelta @@ -32,6 +33,23 @@ from pytz import exceptions, timezone # driver.set_window_size(*window_size) +async def bc_calc(ctx, eqn: str): + try: + # Allow for proper absolute value notation + pipes = eqn.count("|") + eqn = eqn.replace("|", "abs(", pipes // 2).replace("|", ")", pipes // 2) + + with open("temp.txt", "w") as f: + f.write(eqn) + + result = subprocess.check_output(f"bc bc_funcs/* temp.txt", shell=True) + os.remove("temp.txt") + await ctx.send(result.decode("utf-8").replace("\\\n", "").strip()) + except subprocess.CalledProcessError as err: + print(err) + await ctx.send("Something went wrong") + + async def reportStuff(self, ctx, message): channel = self.bot.get_channel( int(self.bot.config[str(ctx.message.guild.id)]["report_channel"]) @@ -548,20 +566,10 @@ class Utils(commands.Cog): output = ", ".join([*commands]) await ctx.send(f"```List of custom commands:\n{output}```") - # @commands.command(aliases=["calc"]) - # async def math(self, ctx, *, eqn: str): - # try: - # # Allow for proper absolute value notation - # pipes = eqn.count("|") - # eqn = eqn.replace("|", "abs(", pipes // 2).replace("|", ")", pipes // 2) - # - # result = subprocess.check_output( - # f"echo 'scale = 10; {eqn}' | bc bc_funcs/*", shell=True - # ) - # await ctx.send(result.decode("utf-8").replace("\\\n", "").strip()) - # except subprocess.CalledProcessError as err: - # print(err) - # await ctx.send("Something went wrong") + @commands.command(aliases=["calc"]) + async def math(self, ctx, *, eqn: str): + # 10 second timeout + await asyncio.wait_for(bc_calc(ctx, eqn), 10) @commands.command() async def retime(self, ctx, start_sec, end_sec, frames=0, framerate=30): |