1 files changed, 81 insertions, 40 deletions
diff --git a/python/runner/interpreter.py b/python/runner/interpreter.py
index da60d72..722bc33 100755
--- a/python/runner/interpreter.py
+++ b/python/runner/interpreter.py
@@ -1,6 +1,8 @@
 #!/usr/bin/python3 -u
 
 import code
+import io
+import pickle
 import sys
 
 import seccomp
@@ -43,49 +45,88 @@ f.add_rule(seccomp.ALLOW, "access")
 f.add_rule(seccomp.ALLOW, "select")
 f.load()
 
-class MyConsole(code.InteractiveConsole):
-    def interact(self, banner=None):
-        if banner is not None:
-            self.write('{}\n'.format(banner))
+if len(sys.argv) > 1 and sys.argv[1] == 'run':
+    # Receive a pickled tuple (code, expr, stdin) on standard input, run it,
+    # and print a pickled tuple (result, out, err, exc) on standard output.
+    data = sys.stdin.buffer.read()
+    code, expr, stdin = pickle.loads(data)
 
-        buffer = []
-        prompt = '>>> '
-        while True:
-            try:
-                line = input(prompt)
-                # Assume we are running the user's program; silence the prompt.
-                if line == 'exec("""\\':
-                    self.write('<run>\n')
-                    prompt = ''
+    result, out, err, exc = None, None, None, None
+    try:
+        sys.stdin = io.StringIO(stdin)
+        sys.stdout = io.StringIO()
+        sys.stderr = io.StringIO()
+        env = {}
+        if code:
+            exec(code, env)
+        if expr:
+            result = eval(expr, env)
+    except Exception as ex:
+        # Exception is not JSON serializable, so return traceback as string
+        # (without the first entry, which is this function).
+        import traceback
+        e_type, e_value, e_tb = sys.exc_info()
+        stack = traceback.extract_tb(e_tb)
+        exc = ''.join(
+                ['Traceback (most recent call last):\n'] +
+                ['  line {}, in {}\n'.format(lineno, name) + (line+'\n' if line else '')
+                    for filename, lineno, name, line in stack[1:]] +
+                traceback.format_exception_only(e_type, e_value)
+            ).rstrip()
+    finally:
+        out = sys.stdout.getvalue()
+        err = sys.stderr.getvalue()
+        sys.stdin.close()
+        sys.stdout.close()
+        sys.stderr.close()
+        outdata = pickle.dumps((result, out, err, exc))
+        sys.__stdout__.buffer.write(outdata)
+
+else:
+    # Run an interactive console.
+    class MyConsole(code.InteractiveConsole):
+        def interact(self, banner=None):
+            if banner is not None:
+                self.write('{}\n'.format(banner))
+
+            buffer = []
+            prompt = '>>> '
+            while True:
+                try:
+                    line = input(prompt)
+                    # Assume we are running the user's program; silence the prompt.
+                    if line == 'exec("""\\':
+                        self.write('<run>\n')
+                        prompt = ''
 
-                buffer.append(line)
-                source = '\n'.join(buffer)
-                more = self.runsource(source)
-                if more:
-                    if prompt:
-                        prompt = '... '
-                else:
+                    buffer.append(line)
+                    source = '\n'.join(buffer)
+                    more = self.runsource(source)
+                    if more:
+                        if prompt:
+                            prompt = '... '
+                    else:
+                        prompt = '>>> '
+                        buffer = []
+                except KeyboardInterrupt:
                     prompt = '>>> '
                     buffer = []
-            except KeyboardInterrupt:
-                prompt = '>>> '
-                buffer = []
-                self.write('\n')
-            except ValueError:
-                break
-            except EOFError:
-                break
+                    self.write('\n')
+                except ValueError:
+                    break
+                except EOFError:
+                    break
 
-    def runcode(self, code):
-        try:
-            exec(code, self.locals)
-        except KeyboardInterrupt:
-            self.write('^C')
-            raise
-        except SystemExit as ex:
-            raise
-        except:
-            # Show traceback for all other exceptions.
-            self.showtraceback()
+        def runcode(self, code):
+            try:
+                exec(code, self.locals)
+            except KeyboardInterrupt:
+                self.write('^C')
+                raise
+            except SystemExit as ex:
+                raise
+            except:
+                # Show traceback for all other exceptions.
+                self.showtraceback()
 
-MyConsole().interact()
+    MyConsole().interact()