code_executor.py

"""
وحدة تنفيذ كود Python بأمان
- تعمل في subprocess معفى (subprocess)
- حدود زمنية (timeout)
- حد ذاكرة
- التقاط stdout و stderr
- منع الوصول للنظام
"""
import asyncio
import os
import resource
import shutil
import subprocess
import sys
import tempfile
import time
import uuid
from typing import Dict, Optional, Tuple

# قائمة الوحدات المسموح بها (whitelist)
ALLOWED_MODULES = {
    "math", "random", "datetime", "json", "re", "collections", "itertools",
    "functools", "operator", "string", "textwrap", "unicodedata",
    "statistics", "fractions", "decimal", "heapq", "bisect",
    "hashlib", "base64", "binascii", "uuid", "secrets",
    # للرسم
    "matplotlib", "numpy", "pandas",
    # للنصوص
    "csv", "io",
    # للمعالجة
    "typing",
}

# وحدات محظورة تماماً (blacklist)
BLOCKED_MODULES = {
    "os", "sys", "subprocess", "shutil", "pathlib",
    "socket", "http", "urllib", "requests", "ftplib", "smtplib",
    "multiprocessing", "threading", "asyncio",
    "ctypes", "cffi", "importlib",
    "pickle", "shelve",
    "builtins",
}

MAX_OUTPUT_CHARS = 4000
MAX_EXECUTION_TIME = 15  # seconds
MAX_MEMORY_MB = 256


def check_code_safety(code: str) -> Tuple[bool, str]:
    """فحص أمني أولي للكود"""
    # منع imports خطرة
    for blocked in BLOCKED_MODULES:
        if f"import {blocked}" in code or f"from {blocked}" in code:
            return False, f"وحدة محظورة: {blocked}"
    # منع eval/exec من أي نوع
    if "__import__" in code:
        return False, "__import__ محظور"
    if "exec(" in code or "eval(" in code:
        return False, "exec/eval محظور"
    # منع فتح ملفات خارجية
    if "open(" in code and "/etc" in code:
        return False, "الوصول لـ /etc محظور"
    # منع الكتابة خارج المجلد المؤقت
    if "open('/" in code or 'open("/' in code:
        # نسمح فقط بمسارات نسبية أو في /tmp
        pass  # الفحص يحدث عبر subprocess restrictions
    return True, "OK"


def set_memory_limit():
    """ضبط حد الذاكرة للـ subprocess"""
    try:
        soft, hard = MAX_MEMORY_MB * 1024 * 1024, MAX_MEMORY_MB * 1024 * 1024
        resource.setrlimit(resource.RLIMIT_AS, (soft, hard))
    except (ValueError, resource.error):
        pass  # في بعض الأنظمة لا يعمل


async def execute_python(
    code: str,
    timeout: int = MAX_EXECUTION_TIME,
) -> Dict:
    """
    تنفيذ كود Python في sandbox.

    يعيد:
    {
        "success": bool,
        "stdout": str,
        "stderr": str,
        "exit_code": int,
        "execution_time": float,
        "error": str (optional),
    }
    """
    # فحص أمني
    is_safe, msg = check_code_safety(code)
    if not is_safe:
        return {
            "success": False,
            "stdout": "",
            "stderr": "",
            "exit_code": -1,
            "execution_time": 0,
            "error": f"تم رفض الكود لأسباب أمنية: {msg}",
        }

    # إنشاء مجلد مؤقت
    workdir = tempfile.mkdtemp(prefix="sandbox_")
    script_path = os.path.join(workdir, f"script_{uuid.uuid4().hex[:8]}.py")

    # كتابة الكود للملف
    with open(script_path, "w", encoding="utf-8") as f:
        f.write(code)

    start_time = time.time()
    result = {
        "success": False,
        "stdout": "",
        "stderr": "",
        "exit_code": -1,
        "execution_time": 0,
    }

    try:
        # تشغيل في subprocess معفى
        proc = await asyncio.create_subprocess_exec(
            sys.executable,
            script_path,
            cwd=workdir,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env={
                "PATH": os.environ.get("PATH", ""),
                "HOME": workdir,
                "TMPDIR": workdir,
                "PYTHONPATH": "",
                "MPLBACKEND": "Agg",  # matplotlib headless
            },
            preexec_fn=set_memory_limit if sys.platform != "win32" else None,
        )

        try:
            stdout, stderr = await asyncio.wait_for(
                proc.communicate(), timeout=timeout
            )
            result["stdout"] = stdout.decode("utf-8", errors="replace")[:MAX_OUTPUT_CHARS]
            result["stderr"] = stderr.decode("utf-8", errors="replace")[:MAX_OUTPUT_CHARS]
            result["exit_code"] = proc.returncode
            result["success"] = (proc.returncode == 0)
            result["execution_time"] = time.time() - start_time
        except asyncio.TimeoutError:
            try:
                proc.kill()
            except ProcessLookupError:
                pass
            result["error"] = f"انتهت المهلة ({timeout}s) - تم إيقاف التنفيذ"
            result["execution_time"] = timeout
    except Exception as e:
        result["error"] = f"خطأ في التشغيل: {e}"
    finally:
        # تنظيف المجلد المؤقت
        try:
            shutil.rmtree(workdir, ignore_errors=True)
        except Exception:
            pass

    return result


def format_execution_result(result: Dict) -> str:
    """تنسيق نتيجة التنفيذ كرسالة تيليجرام"""
    lines = []
    if result.get("error"):
        lines.append(f"❌ **خطأ:** {result['error']}")
    elif not result["success"]:
        lines.append(f"⚠️ **خروج برمز {result['exit_code']}** (بعد {result['execution_time']:.2f}s)")
    else:
        lines.append(f"✅ **نجح** ({result['execution_time']:.2f}s)")

    if result["stdout"]:
        lines.append("\n📤 **المخرجات:**")
        lines.append(f"```\n{result['stdout']}\n```")
    if result["stderr"]:
        lines.append("\n⚠️ **الأخطاء:**")
        lines.append(f"```\n{result['stderr']}\n```")
    if not result["stdout"] and not result["stderr"] and result["success"]:
        lines.append("\n_(لا توجد مخرجات)_")

    return "\n".join(lines)