# License: Public Domain
import os
import subprocess

home_dir = os.path.expanduser("~")

try:
    os.makedirs("/tmp/mysandbox_root")
except FileExistsError:
    pass

def sandbox_pyexec(code):
    proc = subprocess.Popen(
        ["bwrap", "--bind", "/tmp/mysandbox_root", "/", "--ro-bind", "/bin", "/bin", # システムファイルをRead-Onlyで共有する
         "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", 
         "--ro-bind", "/usr", "/usr", "--ro-bind", "/etc", "/etc", # あんまり良くない
         "--ro-bind", home_dir + "/.local/lib", home_dir + "/.local/lib", # pip ライブラリをRead-Onlyで共有する
         "--ro-bind", home_dir + "/.local/bin", home_dir + "/.local/bin",
         "--ro-bind", home_dir + "/.local/share", home_dir + "/.local/share",
         "--bind", "/proc", "/proc", "--unshare-all", "--", # なるべくホストと非共有（systemd-nspawnと違ってプロセスリストは共有してるけど）
         "bash", "-c", "python -"],  # Pythonで標準入力のコードを実行
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True  # stdin / stdout を文字列として扱う
    )
    result, err = proc.communicate(code)
    return (result.strip(), err.strip())
    

r, err = sandbox_pyexec( """print("これはテストです")""")

print("stdout:", r)
print("stderr:", err)
# stdout: これはテストです
# stderr: 

r, err = sandbox_pyexec( """import os
os.system("ls /")""")

print("stdout:", r)
print("stderr:", err)
# stdout: bin
# etc
# home
# lib
# lib64
# proc
# usr
# stderr: 

r, err = sandbox_pyexec( """wrong python syntax here""")

print("stdout:", r)
print("stderr:", err)
# stdout: 
# stderr: File "<stdin>", line 1
#     wrong python syntax here
#           ^^^^^^
# SyntaxError: invalid syntax

# License: Public Domain
import os
import pty
import termios
import select
import sys
import signal
import psutil

home_dir = os.path.expanduser("~")
try:
    os.makedirs("/tmp/mysandbox_root")
except FileExistsError:
    pass

def disable_echo(fd):
    attrs = termios.tcgetattr(fd)
    attrs[3] = attrs[3] & ~termios.ECHO  # lflags から ECHO を落とす
    termios.tcsetattr(fd, termios.TCSANOW, attrs)

# IPythonの動作をシミュレート
# python -m pip install ipython --break-system-packages
helper_code = """from IPython.core.ultratb import VerboseTB
import sys
sys.excepthook = VerboseTB(theme_name="Linux", tb_offset=2)
import linecache

file_count = 1

import ast
def _exec_with_expr(code, globals=None, locals=None):
    global file_count
    tree = ast.parse(code, mode="exec")
    if len(tree.body) == 0:
        print("")
        return None
    filename = "/tmp/test-%s.py"%file_count
    file_count += 1
    linecache.cache[filename] = (
        len(code),
        None,
        code.splitlines(True),
        filename,
    )
    last = tree.body[-1]
    if isinstance(last, ast.Expr):
        expr = ast.Expression(last.value)
        tree.body = tree.body[:-1]
        exec(compile(tree, filename, "exec"), globals, locals)
        return eval(compile(expr, filename, "eval"), globals, locals)
    else:
        exec(compile(tree, filename, "exec"), globals, locals)
        return None

""" # + """_exec_with_expr('print("ok")\\n"ok"', globals(), locals())"""

def sandbox_pyenv():
    pid, fd = pty.fork()
    if pid == 0: # child
         # sys.stdout.flush()][
         disable_echo(0)
         os.execvp("bwrap", 
            ["bwrap", "--bind", "/tmp/mysandbox_root", "/", "--ro-bind", "/bin", "/bin", # システムファイルをRead-Onlyで共有する
             "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", 
             "--ro-bind", "/usr", "/usr", "--ro-bind", "/etc", "/etc", # あんまり良くない
             "--ro-bind", home_dir + "/.local/lib", home_dir + "/.local/lib", # pip ライブラリをRead-Onlyで共有する
             "--ro-bind", home_dir + "/.local/bin", home_dir + "/.local/bin",
             "--ro-bind", home_dir + "/.local/share", home_dir + "/.local/share",
             "--bind", "/proc", "/proc", "--unshare-all", "--", # なるべくホストと非共有（systemd-nspawnと違ってプロセスリストは共有してるけど）
             "python", "-i"],  # Pythonのインタラクティブモードで標準入力のコードを実行
         )
         exit()
    else:
        output = b""
        while True:
            r, _, _ = select.select([fd], [], [])
            chunk = os.read(fd, 1024)
            output += chunk

            if output.endswith(b">>> "):
                #print("debug: bunner: %s"%output)
                break

#    print(_sandbox_pyexec((pid, fd), helper_code))
    _sandbox_pyexec((pid, fd), helper_code)
    return (pid, fd)

TIMEOUT = 3.0 # 3秒

def _sandbox_pyexec(senv, code):
    output = b""
    for line in code.split("\n"):
        os.write(senv[1], line.encode("utf-8") + b"\n")

        while True:
            r, _, _ = select.select([senv[1]], [], [], TIMEOUT)
            if not r:
                parent = psutil.Process(senv[0])
                children = parent.children(recursive=True)
                for c in children:
                    if "python" in c.name():
                        c.send_signal(signal.SIGINT)
                output = b""
                while True:
                    r, _, _ = select.select([senv[1]], [], [])
                    chunk = os.read(senv[1], 1024)
                    output += chunk

                    if output.endswith(b">>> "):
                        #print("debug: KeyboardInterrupt: %s"%output)
                        break
                raise TimeoutError("execution timed out")
            chunk = os.read(senv[1], 1024)
            output += chunk

            if output.endswith(b">>> ") or output.endswith(b"... "):
                output = output[:-4] # プロンプトを除去
                break

    text = output.decode(errors="ignore")
    return text

def sandbox_pyexec(senv, code):
    code = "_exec_with_expr('''" + code.replace('\\', '\\\\').replace("'", "\\'") + "''', globals(), locals())\n"
    return _sandbox_pyexec(senv, code)

def sandbox_reset(senv):
    _sandbox_pyexec(senv, """globals().clear()
import gc
gc.collect()
globals().clear()
""" + helper_code)

senv = sandbox_pyenv()
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "a = 100")) # [exec][/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "print(a)")) # [exec]100\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "a")) # [exec]100\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "b")) # [exec]---------------------------------------------------------------------------\nNameError                                 Traceback (most recent call last)\nFile /tmp/test-4.py:1\n----> 1 b\n\nNameError: name 'b' is not defined\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, """print("aaa")\nprint("bbb")""")) # [exec]aaa\nbbb\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, """print("aaa" + \\\n "bbb")""")) # [exec]aaabbb\n[/exec]

sandbox_reset(senv)
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "a"))  # [exec]---------------------------------------------------------------------------\nNameError                                 Traceback (most recent call last)\nFile /tmp/test-1.py:1\n----> 1 a\n\nNameError: name 'a' is not defined\n[/exec]

try:
    print("[exec]%s[/exec]"%sandbox_pyexec(senv, """while True:\n    continue""")) # TimeoutError
except TimeoutError:
    print ("TimeoutError: execution timed out")

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "print(100)")) # [exec]100\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\nprint(100)")) # [exec]100\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\nprint(\"\"\"ok\"\"\")")) # [exec]ok\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\nprint('''ok''')")) # [exec]ok\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\n'ok'")) # [exec]'ok'\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "# comment test")) # [exec]\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "def a(x): return x*x")) # [exec][/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "\n\na('a')")) # [exec]---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\nFile /tmp/test-9.py:3\n----> 3 a('a')\nFile /tmp/test-8.py:1, in a(x='a')\n----> 1 def a(x): return x*x\n        x = 'a'\nTypeError: can't multiply sequence by non-int of type 'str'\n[/exec]

# License: Public Domain
# pip install pyrefly
# export PATH=$HOME/.local/bin:$PATH
import os
import pty
import termios
import select
import sys
import signal
import psutil
import subprocess
import re


home_dir = os.path.expanduser("~")
try:
    os.makedirs("/tmp/mysandbox_root")
except FileExistsError:
    pass

try:
    os.mkfifo("/tmp/mysandbox_pipe.py")
except FileExistsError:
    pass

def disable_echo(fd):
    attrs = termios.tcgetattr(fd)
    attrs[3] = attrs[3] & ~termios.ECHO  # lflags から ECHO を落とす
    termios.tcsetattr(fd, termios.TCSANOW, attrs)

whole_code = ""
whole_code_line = [""]
file_count = 1

# IPythonの動作をシミュレート
# python -m pip install ipython --break-system-packages
helper_code = """from IPython.core.ultratb import VerboseTB
import sys
sys.excepthook = VerboseTB(theme_name="Linux", tb_offset=2)
import linecache

file_count = 1

import ast
def _exec_with_expr(code, globals=None, locals=None):
    global file_count
    tree = ast.parse(code, mode="exec")
    if len(tree.body) == 0:
        print("")
        return None
    filename = "/tmp/test-%s.py"%file_count
    file_count += 1
    linecache.cache[filename] = (
        len(code),
        None,
        code.splitlines(True),
        filename,
    )
    last = tree.body[-1]
    if isinstance(last, ast.Expr):
        expr = ast.Expression(last.value)
        tree.body = tree.body[:-1]
        exec(compile(tree, filename, "exec"), globals, locals)
        return eval(compile(expr, filename, "eval"), globals, locals)
    else:
        exec(compile(tree, filename, "exec"), globals, locals)
        return None

""" # + """_exec_with_expr('print("ok")\\n"ok"', globals(), locals())"""

def sandbox_pyenv():
    pid, fd = pty.fork()
    if pid == 0: # child
         # sys.stdout.flush()][
         disable_echo(0)
         os.execvp("bwrap", 
            ["bwrap", "--bind", "/tmp/mysandbox_root", "/", "--ro-bind", "/bin", "/bin", # システムファイルをRead-Onlyで共有する
             "--ro-bind", "/lib", "/lib", "--ro-bind", "/lib64", "/lib64", 
             "--ro-bind", "/usr", "/usr", "--ro-bind", "/etc", "/etc", # あんまり良くない
             "--ro-bind", home_dir + "/.local/lib", home_dir + "/.local/lib", # pip ライブラリをRead-Onlyで共有する
             "--ro-bind", home_dir + "/.local/bin", home_dir + "/.local/bin",
             "--ro-bind", home_dir + "/.local/share", home_dir + "/.local/share",
             "--bind", "/proc", "/proc", "--unshare-all", "--", # なるべくホストと非共有（systemd-nspawnと違ってプロセスリストは共有してるけど）
             "python", "-i"],  # Pythonのインタラクティブモードで標準入力のコードを実行
         )
         exit()
    else:
        output = b""
        while True:
            r, _, _ = select.select([fd], [], [])
            chunk = os.read(fd, 1024)
            output += chunk

            if output.endswith(b">>> "):
                #print("debug: bunner: %s"%output)
                break

#    print(_sandbox_pyexec((pid, fd), helper_code))
    _sandbox_pyexec((pid, fd), helper_code)
    return (pid, fd)

TIMEOUT = 3.0 # 3秒

def _sandbox_pyexec(senv, code):
    output = b""
    for line in code.split("\n"):
        os.write(senv[1], line.encode("utf-8") + b"\n")

        while True:
            r, _, _ = select.select([senv[1]], [], [], TIMEOUT)
            if not r:
                parent = psutil.Process(senv[0])
                children = parent.children(recursive=True)
                for c in children:
                    if "python" in c.name():
                        c.send_signal(signal.SIGINT)
                output = b""
                while True:
                    r, _, _ = select.select([senv[1]], [], [])
                    chunk = os.read(senv[1], 1024)
                    output += chunk

                    if output.endswith(b">>> "):
                        #print("debug: KeyboardInterrupt: %s"%output)
                        break
                raise TimeoutError("execution timed out")
            chunk = os.read(senv[1], 1024)
            output += chunk

            if output.endswith(b">>> ") or output.endswith(b"... "):
                output = output[:-4] # プロンプトを除去
                break

    text = output.decode(errors="ignore")
    return text

def sandbox_pyexec(senv, code):
    global whole_code
    global whole_code_line
    global file_count

    file_count += 1
    e_code = "_exec_with_expr('''" + code.replace('\\', '\\\\').replace("'", "\\'") + "''', globals(), locals())\n"
    r = _sandbox_pyexec(senv, e_code)

#    print([r[i].encode("utf-8") if i < len(r) else None for i in range(5)])

    # エラーになったコードは履歴に残さない
    if r.startswith("\x1b[91m---------------------------------------------------------------------------"):
        return r

    # lintが通らなかった時も履歴には残す（実際に残ってるため）
    whole_code = whole_code + code + "\n"
    for i, v in enumerate(code.split("\n")): # TODO: 使う側で二分木にした方が良い
        whole_code_line.append("test-%s.py in <cell line: %s>"%(file_count - 1, i))
#    print(whole_code_line)
#    print(whole_code)


    # 返り値が何もない時（関数だけなど）はlintに掛ける
    if r == "":
        proc = subprocess.Popen(["pyrefly", "check", "/tmp/mysandbox_pipe.py"],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True)
        with open("/tmp/mysandbox_pipe.py", "a", encoding="utf-8") as f:
            f.write(whole_code)

        stdout, stderr = proc.communicate()
        def replace_with_array_access(match):
            return whole_code_line[int(match.group(1))] + ":" + match.group(2)
        stdout = re.sub(r'/tmp/mysandbox_pipe.py:(\d+):(\d+)', replace_with_array_access, stdout)
        return stdout

    return r

def sandbox_reset(senv):
    global whole_code
    global whole_code_line
    global file_count
    _sandbox_pyexec(senv, """globals().clear()
import gc
gc.collect()
globals().clear()
""" + helper_code)
    whole_code = ""
    whole_code_line = [""]
    file_count = 1

senv = sandbox_pyenv()
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "a = 100")) # [exec][/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "print(a)")) # [exec]100\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "a")) # [exec]100\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "b")) # [exec]---------------------------------------------------------------------------\nNameError                                 Traceback (most recent call last)\nFile /tmp/test-4.py:1\n----> 1 b\n\nNameError: name 'b' is not defined\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, """print("aaa")\nprint("bbb")""")) # [exec]aaa\nbbb\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, """print("aaa" + \\\n "bbb")""")) # [exec]aaabbb\n[/exec]

sandbox_reset(senv)
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "a"))  # [exec]---------------------------------------------------------------------------\nNameError                                 Traceback (most recent call last)\nFile /tmp/test-1.py:1\n----> 1 a\n\nNameError: name 'a' is not defined\n[/exec]

try:
    print("[exec]%s[/exec]"%sandbox_pyexec(senv, """while True:\n    continue""")) # TimeoutError
except TimeoutError:
    print ("TimeoutError: execution timed out")

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "print(100)")) # [exec]100\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\nprint(100)")) # [exec]100\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\nprint(\"\"\"ok\"\"\")")) # [exec]ok\n[/exec]
print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\nprint('''ok''')")) # [exec]ok\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "if True:\n    pass\n'ok'")) # [exec]'ok'\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "# comment test")) # [exec]\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "def a(x): return x*x")) # [exec][/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "\n\na('a')")) # [exec]---------------------------------------------------------------------------\nTypeError                                 Traceback (most recent call last)\nFile /tmp/test-9.py:3\n----> 3 a('a')\nFile /tmp/test-8.py:1, in a(x='a')\n----> 1 def a(x): return x*x\n        x = 'a'\nTypeError: can't multiply sequence by non-int of type 'str'\n[/exec]

print("[exec]%s[/exec]"%sandbox_pyexec(senv, "import sympy as sp\ndef a(x):\n    sp.undefined()\n    return sp.undefined2()"))
"""
[exec]ERROR No attribute `undefined` in module `sympy` [missing-attribute]
  --> test-11.py in <cell line: 2>:5
   |
18 |     sp.undefined()
   |     ^^^^^^^^^^^^
   |
ERROR No attribute `undefined2` in module `sympy` [missing-attribute]
  --> test-11.py in <cell line: 3>:12
   |
19 |     return sp.undefined2()
   |            ^^^^^^^^^^^^^
   |
[/exec]
"""