app/sandbox.py

"""
Real Python Sandbox
====================
Executes arbitrary Python code with:
- Auto pip-install on ImportError (retries up to 3 times)
- Persistent package directory across requests
- Captures stdout, stderr, files, images, audio
- Returns everything as structured output
"""
import os, sys, re, json, subprocess, tempfile, shutil, base64, time
from pathlib import Path

# Persistent dirs
HOME       = Path(os.environ.get("HOME", "/home/user"))
PKG_DIR    = HOME / ".praison_pkgs"
WORK_DIR   = HOME / ".praison_work"
PKG_DIR.mkdir(parents=True, exist_ok=True)
WORK_DIR.mkdir(parents=True, exist_ok=True)

# Track installed packages this session
_installed: set = set()

# Common aliases: import name -> pip package name
_PIP_ALIASES = {
    "cv2":               "opencv-python-headless",
    "PIL":               "Pillow",
    "bs4":               "beautifulsoup4",
    "sklearn":           "scikit-learn",
    "yaml":              "pyyaml",
    "dotenv":            "python-dotenv",
    "duckduckgo_search": "duckduckgo-search",
    "telegram":          "pyTelegramBotAPI",
    "googlesearch":      "googlesearch-python",
    "forex_python":      "forex-python",
    "dateutil":          "python-dateutil",
    "Crypto":            "pycryptodome",
    "nacl":              "PyNaCl",
    "gi":                "PyGObject",
    "wx":                "wxPython",
    "usb":               "pyusb",
    "serial":            "pyserial",
    "pynput":            "pynput",
    "pyttsx3":           "pyttsx3",
    "speech_recognition":"SpeechRecognition",
    "wikipedia":         "wikipedia",
    "tweepy":            "tweepy",
    "instaloader":       "instaloader",
    "yt_dlp":            "yt-dlp",
    "pytube":            "pytube",
    "moviepy":           "moviepy",
    "pdf2image":         "pdf2image",
    "docx":              "python-docx",
    "pptx":              "python-pptx",
    "xlrd":              "xlrd",
    "openpyxl":          "openpyxl",
    "qrcode":            "qrcode",
    "barcode":           "python-barcode",
    "cryptography":      "cryptography",
    "paramiko":          "paramiko",
    "ftplib":            "ftplib",
    "imaplib":           "imaplib",
    "smtplib":           "smtplib",
    "win32api":          "pywin32",
    "psutil":            "psutil",
    "GPUtil":            "GPUtil",
    "platform":          "platform",
    "distro":            "distro",
    "netifaces":         "netifaces",
    "scapy":             "scapy",
    "nmap":              "python-nmap",
    "shodan":            "shodan",
    "boto3":             "boto3",
    "google.cloud":      "google-cloud",
    "azure":             "azure",
    "openai":            "openai",
    "anthropic":         "anthropic",
    "langchain":         "langchain",
    "transformers":      "transformers",
    "torch":             "torch",
    "tensorflow":        "tensorflow",
    "keras":             "keras",
    "sklearn":           "scikit-learn",
    "xgboost":           "xgboost",
    "lightgbm":          "lightgbm",
    "catboost":          "catboost",
    "prophet":           "prophet",
    "statsmodels":       "statsmodels",
    "scipy":             "scipy",
    "sympy":             "sympy",
    "networkx":          "networkx",
    "igraph":            "python-igraph",
    "nltk":              "nltk",
    "spacy":             "spacy",
    "textblob":          "textblob",
    "gensim":            "gensim",
    "flair":             "flair",
    "sumy":              "sumy",
    "rake_nltk":         "rake-nltk",
    "wordcloud":         "wordcloud",
    "folium":            "folium",
    "plotly":            "plotly",
    "bokeh":             "bokeh",
    "altair":            "altair",
    "seaborn":           "seaborn",
    "matplotlib":        "matplotlib",
    "pandas":            "pandas",
    "numpy":             "numpy",
    "arrow":             "arrow",
    "pendulum":          "pendulum",
    "pytz":              "pytz",
    "babel":             "Babel",
    "pydantic":          "pydantic",
    "aiohttp":           "aiohttp",
    "httpx":             "httpx",
    "requests":          "requests",
    "flask":             "Flask",
    "fastapi":           "fastapi",
    "celery":            "celery",
    "redis":             "redis",
    "pymongo":           "pymongo",
    "sqlalchemy":        "SQLAlchemy",
    "psycopg2":          "psycopg2-binary",
    "pymysql":           "PyMySQL",
    "sqlite3":           "sqlite3",
    "gtts":              "gTTS",
    "playsound":         "playsound",
    "pydub":             "pydub",
    "librosa":           "librosa",
    "soundfile":         "soundfile",
    "pyaudio":           "PyAudio",
    "yfinance":          "yfinance",
    "alpha_vantage":     "alpha-vantage",
    "fredapi":           "fredapi",
    "quandl":            "quandl",
    "ccxt":              "ccxt",
    "ta":                "ta",
    "backtrader":        "backtrader",
    "zipline":           "zipline-reloaded",
}


def _pip_name(import_name: str) -> str:
    return _PIP_ALIASES.get(import_name, import_name)


def pip_install(packages: list) -> tuple[bool, str]:
    """Actually install packages into PKG_DIR."""
    to_install = []
    for p in packages:
        pip_p = _pip_name(p.strip())
        norm  = pip_p.lower().replace("-","_")
        if norm not in _installed:
            to_install.append(pip_p)
    if not to_install:
        return True, "Already installed"
    cmd = [sys.executable, "-m", "pip", "install", "--quiet",
           "--target", str(PKG_DIR), "--upgrade"] + to_install
    r = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
    if r.returncode == 0:
        for p in to_install:
            _installed.add(p.lower().replace("-","_"))
        return True, f"Installed: {', '.join(to_install)}"
    return False, r.stderr[-600:]


def _extract_missing_module(stderr: str) -> str | None:
    """Extract missing module name from ImportError traceback."""
    patterns = [
        r"No module named '([^']+)'",
        r"ModuleNotFoundError: No module named '([^']+)'",
        r"ImportError: cannot import name .+ from '([^']+)'",
        r"ImportError: No module named ([^\s]+)",
    ]
    for pat in patterns:
        m = re.search(pat, stderr)
        if m:
            # Get root package (e.g. "google.cloud.storage" -> "google")
            return m.group(1).split(".")[0]
    return None


def _make_runner(code: str, work_dir: str) -> str:
    """Wrap code so it runs with PKG_DIR in path and captures structured output."""
    return f'''
import sys, os, json, base64, traceback, io, contextlib
sys.path.insert(0, {repr(str(PKG_DIR))})
os.chdir({repr(work_dir)})

_output_files = []
_stdout_buf = io.StringIO()
_stderr_buf = io.StringIO()

with contextlib.redirect_stdout(_stdout_buf), contextlib.redirect_stderr(_stderr_buf):
    try:
        exec(compile({repr(code)}, "<praison_tool>", "exec"), {{"__name__": "__main__"}})
        _success = True
        _error   = ""
    except SystemExit:
        _success = True
        _error   = ""
    except Exception as _e:
        _success = False
        _error   = traceback.format_exc()

_stdout = _stdout_buf.getvalue()
_stderr = _stderr_buf.getvalue()
if _error:
    _stderr += "\\n" + _error

# Collect any files written to work dir
_files = []
for _f in os.listdir({repr(work_dir)}):
    _fp = os.path.join({repr(work_dir)}, _f)
    if os.path.isfile(_fp):
        try:
            _size = os.path.getsize(_fp)
            if _size < 10_000_000:  # 10MB limit
                with open(_fp, "rb") as _fh:
                    _b64 = base64.b64encode(_fh.read()).decode()
                _ext = _f.rsplit(".", 1)[-1].lower() if "." in _f else ""
                _files.append({{"name": _f, "size": _size, "b64": _b64, "ext": _ext}})
        except Exception:
            pass

print(json.dumps({{
    "ok":      _success,
    "stdout":  _stdout[:8000],
    "stderr":  _stderr[:3000],
    "files":   _files,
}}))
'''


def run(code: str, max_retries: int = 3, timeout: int = 60) -> dict:
    """
    Execute Python code. Auto-installs missing modules and retries.
    Returns:
        ok:      bool
        stdout:  str
        stderr:  str
        files:   list of {name, size, b64, ext}
        installs: list of installed packages
        attempts: int
    """
    installs = []
    work_dir = tempfile.mkdtemp(dir=WORK_DIR)

    try:
        for attempt in range(max_retries):
            script   = _make_runner(code, work_dir)
            tmp_path = tempfile.NamedTemporaryFile(
                mode="w", suffix=".py", delete=False, encoding="utf-8"
            )
            tmp_path.write(script)
            tmp_path.close()

            env = os.environ.copy()
            env["PYTHONPATH"] = str(PKG_DIR) + os.pathsep + env.get("PYTHONPATH","")

            try:
                proc = subprocess.run(
                    [sys.executable, tmp_path.name],
                    capture_output=True, text=True,
                    timeout=timeout, env=env
                )
            except subprocess.TimeoutExpired:
                return {"ok":False,"stdout":"","stderr":f"Timed out after {timeout}s",
                        "files":[],"installs":installs,"attempts":attempt+1}
            finally:
                try: os.unlink(tmp_path.name)
                except: pass

            raw = proc.stdout.strip()
            stderr_raw = proc.stderr.strip()

            # Parse output
            if raw:
                try:
                    last_line = [l for l in raw.split("\n") if l.strip()][-1]
                    result = json.loads(last_line)
                    result["installs"]  = installs
                    result["attempts"]  = attempt + 1
                    result["stderr"]    = result.get("stderr","") + ("\n"+stderr_raw if stderr_raw else "")
                    # Check if we need to install something
                    missing = _extract_missing_module(result.get("stderr",""))
                    if missing and not result["ok"] and attempt < max_retries - 1:
                        ok, msg = pip_install([missing])
                        installs.append({"package": missing, "ok": ok, "msg": msg})
                        continue
                    return result
                except (json.JSONDecodeError, IndexError):
                    pass

            # No JSON output — raw stdout
            combined_err = proc.stderr + "\n" + raw
            missing = _extract_missing_module(combined_err)
            if missing and attempt < max_retries - 1:
                ok, msg = pip_install([missing])
                installs.append({"package": missing, "ok": ok, "msg": msg})
                continue

            return {"ok": proc.returncode == 0,
                    "stdout": raw[:6000],
                    "stderr": combined_err[:2000],
                    "files":  [],
                    "installs": installs,
                    "attempts": attempt + 1}

        return {"ok":False,"stdout":"","stderr":"Max retries reached",
                "files":[],"installs":installs,"attempts":max_retries}
    finally:
        shutil.rmtree(work_dir, ignore_errors=True)