executor.py

"""
Code Execution Engine - Sandboxed via Docker.

Handles execution of Python and shell code with proper sandboxing.
"""

import discord
import asyncio
import subprocess
import traceback
import textwrap
import logging
from datetime import datetime, timedelta

from .config import DOCKER_IMAGE, DOCKER_TIMEOUT, INDEXER_NETWORKS
from .allowlist import check_indexer_status, get_indexer_logs, restart_indexer

log = logging.getLogger(__name__)


async def execute_in_docker(command: list[str], timeout: int = DOCKER_TIMEOUT) -> tuple[bool, str]:
    """
    Execute a command inside a Docker container.

    The container:
    - Runs as non-root user
    - Has no network access (--network none)
    - Has limited memory (512MB)
    - Has limited CPU (1 core)
    - Auto-removes after execution

    Returns (success, output_or_error)
    """
    docker_cmd = [
        "docker", "run",
        "--rm",                          # Remove container after execution
        "--network", "none",             # No network access
        "--memory", "512m",              # Memory limit
        "--cpus", "1",                   # CPU limit
        "--user", "sandbox",             # Run as non-root
        "--read-only",                   # Read-only filesystem
        "--tmpfs", "/tmp:rw,noexec,nosuid,size=100m",  # Writable /tmp
        DOCKER_IMAGE,
    ] + command

    try:
        result = subprocess.run(
            docker_cmd,
            capture_output=True,
            text=True,
            timeout=timeout
        )

        output = result.stdout
        if result.stderr:
            output += f"\n[stderr]: {result.stderr}"

        if result.returncode != 0:
            return False, f"Exit code {result.returncode}\n{output}"

        return True, output.strip() if output.strip() else "Done"

    except subprocess.TimeoutExpired:
        return False, f"Command timed out after {timeout} seconds"
    except FileNotFoundError:
        return False, "Docker not found. Is Docker running?"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


async def execute_shell(command: str, shell_type: str = "auto") -> tuple[bool, str]:
    """
    Execute a shell command inside a Docker sandbox.

    Returns (success, output_or_error)
    """
    return await execute_in_docker(["bash", "-c", command])


# Alias for compatibility
async def execute_bash(command: str) -> tuple[bool, str]:
    return await execute_shell(command)


async def execute_python_sandboxed(code: str) -> tuple[bool, str]:
    """
    Execute Python code inside a Docker sandbox.

    This is for general Python code that doesn't need Discord access.
    Returns (success, output_or_error)
    """
    wrapper = f'''
import sys
try:
    result = None
    output = []

{code}

    if result is not None:
        print(result)
    elif output:
        print("\\n".join(str(x) for x in output))
    else:
        print("Done")
except Exception as e:
    import traceback
    print(f"Error: {{type(e).__name__}}: {{e}}", file=sys.stderr)
    traceback.print_exc()
    sys.exit(1)
'''
    return await execute_in_docker(["python", "-c", wrapper])


async def execute_code(code: str, context: dict) -> tuple[bool, str]:
    """
    Execute Python code with access to Discord objects.

    This runs on the HOST (not sandboxed) because it needs Discord API access.
    Only admin users can trigger this.

    Context contains:
        - message: the triggering discord.Message
        - channel: the channel
        - guild: the guild
        - client: the Discord client
        - recent_messages: list of recent Message objects
        - discord: the discord module
        - asyncio: the asyncio module
        - datetime, timedelta: for time operations
        - Allowlisted ops: check_indexer_status, get_indexer_logs, restart_indexer

    Returns (success, output_or_error)
    """
    # Check if the code looks like it needs Discord or host access
    discord_keywords = ['message', 'channel', 'guild', 'client', 'discord', 'recent_messages', 'sent_message']
    allowlist_keywords = ['check_indexer_status', 'get_indexer_logs', 'restart_indexer']

    needs_discord = any(kw in code for kw in discord_keywords)
    needs_allowlist = any(kw in code for kw in allowlist_keywords)

    # If no special access needed, run in sandbox
    if not needs_discord and not needs_allowlist:
        log.info("Code doesn't need special access - running in Docker sandbox")
        return await execute_python_sandboxed(code)

    # Otherwise, run with full context (on host)
    log.info("Code needs special access - running on host with context")

    # Build execution namespace
    namespace = {
        **context,
        'discord': discord,
        'asyncio': asyncio,
        'datetime': datetime,
        'timedelta': timedelta,
        'log': log,
        'result': None,
        'output': [],
        # Allowlisted operations
        'check_indexer_status': check_indexer_status,
        'get_indexer_logs': get_indexer_logs,
        'restart_indexer': restart_indexer,
        'INDEXER_NETWORKS': INDEXER_NETWORKS,
    }

    try:
        wrapped_code = f"""
async def __exec_code():
    global result, output
{textwrap.indent(code, '    ')}
    return result if result is not None else ('\\n'.join(output) if output else 'Done')
"""
        exec(wrapped_code, namespace)
        exec_result = await namespace['__exec_code']()
        return True, str(exec_result) if exec_result else "Done"

    except Exception as e:
        error_msg = f"{type(e).__name__}: {e}\n{traceback.format_exc()}"
        return False, error_msg