Run latexify and calc in a secure subprocess (#26)

* Update math.py

Some annoying whitespace left over from the mobile GitHub editor probably

* Update math.py

* Finally fix this

* Update math.py

* Documentation

* Un-inline some comments

* Capitals and trailing spaces

* Fix comment position

* Reorder Popen

* Fix floating comments I guess

* Fix bandit

Signed-off-by: Martmists <[email protected]>

1 files changed, 59 insertions, 27 deletions

diff --git a/bot/cogs/math.py b/bot/cogs/math.py
index 3a8b7c222..0ef045748 100644
--- a/bot/cogs/math.py
+++ b/bot/cogs/math.py
@@ -1,21 +1,49 @@
 #!/usr/bin/env python3
 # -*- coding: utf-8 -*-
 
-
+import asyncio
+import sys
+from contextlib import suppress
 from io import BytesIO
 from re import search
-
+from subprocess import PIPE, Popen, STDOUT, TimeoutExpired  # noqa: B404
 
 from aiohttp import ClientSession
 
 from discord import File
 from discord.ext.commands import command
 
-from sympy import latex
-from sympy.parsing.sympy_parser import parse_expr
 
+LATEX_URL = "https://latex2png.com"
+
+
+async def run_sympy(sympy_code: str, calc: bool = False, timeout: int = 10) -> str:
+    if calc:
+        code_ = "parse_expr(sys.argv[1]).doit()"  # Run the expression
+    else:
+        code_ = "parse_expr(sys.argv[1], evaluate=True)"  # Just latexify it without running
+
+    if "__" in sympy_code:
+        # They're trying to exploit something, raise an error
+        raise TypeError("'__' not allowed in sympy code")
+
+    proc = Popen([  # noqa: B603
+                    sys.executable, "-c",
+                    "import sys,sympy;from sympy.parsing.sympy_parser import parse_expr;"
+                    f"print(sympy.latex({code_}))", sympy_code
+                 ], env={},  # Disable environment variables for security
+                 stdout=PIPE, stderr=STDOUT)  # reroute all to stdout
+
+    for _ in range(timeout*4):  # Check if done every .25 seconds for `timeout` seconds
+        await asyncio.sleep(1/4)
 
-LATEX_URL = "http://latex2png.com"
+        # Ignore TimeoutExpired...
+        with suppress(TimeoutExpired):
+            proc.wait(0)
+            break  # ... But stop the loop when not raised
+
+    proc.kill()  # Kill the process regardless of whether it finished or not
+    return proc.returncode, proc.stdout.read().decode().strip()
 
 
 class Math:
@@ -28,18 +56,24 @@ class Math:
         Return the LaTex output for a mathematical expression
         """
 
-        fixed_expr = expr.replace('^', '**').strip('`').replace("__", "")
+        fixed_expr = expr.replace('^', '**').strip('`')  # Syntax fixes
         try:
-            parsed = parse_expr(fixed_expr, evaluate=False)
+            retcode, parsed = await run_sympy(fixed_expr)  # Run the sympy code
 
-        except SyntaxError:
-            await ctx.send("Invalid expression!")
+        except TypeError as e:  # Exploit was tried
+            await ctx.send(e.args[0])
 
         else:
-            ltx = latex(parsed)
-
+            if retcode != 0:  # ERROR
+                await ctx.send(f"Error:\n```{parsed}```")
+                return
+            elif not parsed:  # Timeout
+                await ctx.send("Code did not return or took too long")
+                return
+
+            # Send LaTeX to website to get image
             data = {
-                "latex": ltx,
+                "latex": parsed,
                 "res": 300,
                 "color": 808080
             }
@@ -61,28 +95,26 @@ class Math:
     async def calc(self, ctx, *, expr: str):
         """
         Return the LaTex output for the solution to a mathematical expression
-
-        Note that exponentials are disabled to avoid abuse
         """
 
-        fixed_expr = expr.replace('^', '**').strip('`')
-
-        if any(x in fixed_expr for x in ("**", "__")):
-            return await ctx.send(
-                "You used an expression that has been disabled for security, our apologies")
-
+        fixed_expr = expr.replace('^', '**').strip('`')  # Syntax fixes
         try:
-            parsed = parse_expr(fixed_expr, evaluate=False)
-            result = parsed.doit()
+            retcode, parsed = await run_sympy(fixed_expr, calc=True)  # Run sympy
 
-        except SyntaxError:
-            await ctx.send("Invalid expression!")
+        except TypeError as e:  # Exploitation tried
+            await ctx.send(e.args[0])
 
         else:
-            ltx = latex(result)
-
+            if retcode != 0:  # ERROR
+                await ctx.send(f"Error:\n```{parsed}```")
+                return
+            elif not parsed:  # Timeout
+                await ctx.send("Code did not return or took too long")
+                return
+
+            # Send LaTeX to website to get image
             data = {
-                "latex": ltx,
+                "latex": parsed,
                 "res": 300,
                 "color": 808080
             }