#!/usr/bin/python3 -u

import code
import io
import pickle
import sys

import seccomp

f = seccomp.SyscallFilter(defaction=seccomp.KILL)
# Necessary for Python.
f.add_rule(seccomp.ALLOW, "brk")
f.add_rule(seccomp.ALLOW, "exit_group")
f.add_rule(seccomp.ALLOW, "ioctl")
f.add_rule(seccomp.ALLOW, "mmap")
f.add_rule(seccomp.ALLOW, "munmap")
f.add_rule(seccomp.ALLOW, "rt_sigaction")
f.add_rule(seccomp.ALLOW, "rt_sigreturn")

# Mostly harmless.
f.add_rule(seccomp.ALLOW, "mprotect")

# Allow reading from stdin and writing to stdout/stderr.
f.add_rule(seccomp.ALLOW, "read", seccomp.Arg(0, seccomp.EQ, sys.stdin.fileno()))
f.add_rule(seccomp.ALLOW, "write", seccomp.Arg(0, seccomp.EQ, sys.stdout.fileno()))
f.add_rule(seccomp.ALLOW, "write", seccomp.Arg(0, seccomp.EQ, sys.stderr.fileno()))

# Needed for finding source code for exceptions.
f.add_rule(seccomp.ALLOW, "stat")
f.add_rule(seccomp.ALLOW, "open", seccomp.Arg(1, seccomp.MASKED_EQ, 0x3, 0))  # O_RDONLY
f.add_rule(seccomp.ALLOW, "close")
f.add_rule(seccomp.ALLOW, "read")
f.add_rule(seccomp.ALLOW, "fstat")
f.add_rule(seccomp.ALLOW, "lseek")
f.add_rule(seccomp.ALLOW, "fcntl")

# Needed for help().
f.add_rule(seccomp.ALLOW, "openat", seccomp.Arg(2, seccomp.MASKED_EQ, 0x3, 0))  # O_RDONLY
f.add_rule(seccomp.ALLOW, "getdents")
f.add_rule(seccomp.ALLOW, "getrlimit", seccomp.Arg(0, seccomp.EQ, 3))  # RLIMIT_STACK
f.add_rule(seccomp.ALLOW, "getrlimit", seccomp.Arg(0, seccomp.EQ, 7))  # RLIMIT_NOFILE

# Needed for code.InteractiveConsole.
f.add_rule(seccomp.ALLOW, "access")
f.add_rule(seccomp.ALLOW, "select")
f.load()

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)

    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:
                        prompt = '>>> '
                        buffer = []
                except KeyboardInterrupt:
                    prompt = '>>> '
                    buffer = []
                    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()

    MyConsole().interact()