sanitize inputs for math

1 files changed, 53 insertions, 5 deletions

diff --git a/cogs/math_commands.py b/cogs/math_commands.py
index bf164b3..e30c92b 100644
--- a/cogs/math_commands.py
+++ b/cogs/math_commands.py
@@ -5,12 +5,45 @@
 
 import logging
 import asyncio
+import re
 
 from discord import ApplicationContext
 from discord.ext import commands
 
 logger = logging.getLogger(__name__)
 
+MATH_SANITIZE_PATTERN = re.compile(r'^[\d\s\+\-\*\/\^\%\(\)\.\,\=\\<\>\!\|\&\~\?\:\'\"\[\]\{\}a-zA-Z]+$')
+
+
+def sanitize_math_input(equation: str) -> str | None:
+    if not equation or len(equation) > 1000:
+        return None
+    equation = equation.strip()
+    if not equation:
+        return None
+    if not MATH_SANITIZE_PATTERN.match(equation):
+        return None
+    for char in equation:
+        if ord(char) < 32 or ord(char) > 126:
+            return None
+    if equation.count('(') != equation.count(')'):
+        return None
+    if equation.count('[') != equation.count(']'):
+        return None
+    if equation.count('{') != equation.count('}'):
+        return None
+    dangerous_patterns = [
+        r'eval', r'exec', r'__import__', r'os\.', r'sys\.',
+        r'subprocess', r'open(', r'file', r'shutil',
+        r'lambda:', r'compile', r'dir(', r'vars(',
+        r'\|\|', r'&&', r';', r'`', r'\$', r'\\n', r'\\r',
+        r'\x', r'\0', r'\a', r'\b', r'\f', r'\v',
+    ]
+    for pattern in dangerous_patterns:
+        if re.search(pattern, equation, re.IGNORECASE):
+            return None
+    return equation
+
 
 class Math(commands.Cog):
     def __init__(self, bot):
@@ -21,13 +54,28 @@ class Math(commands.Cog):
         description="wrug szyskih matfisuf"
     )
     async def math(self, ctx: ApplicationContext, equation: str):
+        sanitized = sanitize_math_input(equation)
+        if sanitized is None:
+            await ctx.respond("Invalid equation. Only basic math expressions allowed.", ephemeral=True)
+            logger.warning(f"Math command rejected invalid input: {equation}")
+            return
         try:
-            math = await asyncio.create_subprocess_exec(f"qalc -e '{equation}'")
-            result = await math.wait()
-            await ctx.respond(f"{equation} = {result.decode('utf-8').strip()}")
-            logger.info(f"Math command executed: {equation}")
+            math_proc = await asyncio.create_subprocess_exec(
+                "qalc", "-e", sanitized,
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE
+            )
+            stdout, stderr = await math_proc.communicate()
+            if math_proc.returncode != 0:
+                error_output = stderr.decode('utf-8', errors='replace').strip()
+                await ctx.respond(f"Calculation error: {error_output}", ephemeral=True)
+                logger.error(f"qalc error: {error_output}")
+                return
+            result = stdout.decode('utf-8', errors='replace').strip()
+            await ctx.respond(f"{sanitized} = {result}")
+            logger.info(f"Math command executed: {sanitized}")
         except Exception as e:
-            await ctx.respond(e, ephemeral=True)
+            await ctx.respond(f"Error: {str(e)}", ephemeral=True)
             logger.error(f"Error in math command: {e}", exc_info=True)