"""Swarm worker task executor (Bridge Builder 2/5).

Executes a swarm slot's task text in a sandbox and captures the result.

Sandbox rules (hard):
- No network calls except to localhost.
- No writes outside /tmp.
- No sudo / privilege escalation.
- No credential access (no reading ~/.ssh, ~/.aws, key stores, etc).
- Strict timeout; kill the process group on exceed.
- Refuse tasks that look destructive without executing anything.
"""

import os
import re
import shlex
import signal
import subprocess
import time

OUTPUT_TRUNCATE = 2000

# Patterns that indicate a potentially destructive command. Matched against
# the raw task text (case-insensitive) before any execution is attempted.
_REFUSE_PATTERNS = [
    r"\brm\s+-rf?\b",          # rm -r / rm -rf
    r"\brm\s+.*\s+/\s*$",      # rm ... /  (trailing root)
    r"\bmkfs\b",
    r"\bdd\b.*\bof=/dev/",
    r"\bdd\b.*\bof=\s*/dev/",
    r":\(\)\s*\{",              # fork bomb
    r"\bshutdown\b",
    r"\breboot\b",
    r"\bpoweroff\b",
    r"\bhalt\b",
    r"\binit\s+[06]\b",
    r"\bsudo\b",
    r"\bsu\b",
    r"\bchmod\s+-R\s+777\s+/\b",
    r"\bchown\s+-R\b.*\s+/\s*$",
    r"\bmv\s+.*\s+/\s*$",
    r"\bwipefs\b",
    r"\bshred\b.*\s/dev/",
    r">\s*/dev/sd",
    r"\bcurl\b.*\|\s*(ba)?sh",  # curl|sh pipe-to-shell
    r"\bwget\b.*\|\s*(ba)?sh",
    r"\bnc\b.*-e\s",            # netcat reverse shell
    r"\bpython\w*\s+-c\b.*socket",
]

_REFUSE_RE = re.compile("|".join("(?:%s)" % p for p in _REFUSE_PATTERNS),
                       re.IGNORECASE)

# Paths that must never be read (credential / identity material).
_FORBIDDEN_READ_PREFIXES = (
    os.path.expanduser("~/.ssh"),
    os.path.expanduser("~/.aws"),
    os.path.expanduser("~/.gnupg"),
    "/etc/shadow",
    "/etc/netvm",
    "/etc/ssl/private",
)

# A task is treated as a shell command when it is short, single-purpose,
# and does not look like prose/instructions. Heuristic: one or two lines,
# starts with a plausible command token, no sentence-like structure.
_PROSE_RE = re.compile(r"[.?!]\s+[A-Z]|\n\n|please\s|you\s+are\s|your\s+task",
                       re.IGNORECASE)


def _looks_like_shell(task_text):
    """Heuristic: is this plausibly a shell command?"""
    t = task_text.strip()
    if not t:
        return False
    if len(t.splitlines()) > 3:
        return False
    if _PROSE_RE.search(t):
        return False
    # Must start with a word-ish token (not a quote or sentence).
    first = t.split()[0] if t.split() else ""
    if not re.match(r"^[a-zA-Z0-9_.\-/]+$", first):
        return False
    return True


def _refused(task_text):
    return bool(_REFUSE_RE.search(task_text))


def _sandbox_env():
    """Minimal environment: strip anything credential-shaped."""
    env = {
        "PATH": "/usr/bin:/bin",
        "HOME": "/tmp",
        "TMPDIR": "/tmp",
        "LANG": "C.UTF-8",
    }
    return env


def execute_task(task_text, timeout=600):
    """Execute a swarm slot's task text in a sandbox.

    Args:
        task_text: the task string from the swarm slot.
        timeout: max seconds for execution (default 600). Strictly enforced.

    Returns:
        dict(success=bool, output=str, duration_s=float, error=str|None)
    """
    started = time.monotonic()
    text = (task_text or "").strip()

    def done(success, output, error=None):
        dur = round(time.monotonic() - started, 3)
        out = (output or "")[:OUTPUT_TRUNCATE]
        return {
            "success": success,
            "output": out,
            "duration_s": dur,
            "error": error,
        }

    if not text:
        return done(False, "", error="empty task")

    if _refused(text):
        return done(False, "", error="refused: potentially destructive")

    if not _looks_like_shell(text):
        return done(
            True,
            "received but not executable: task is prose/instructions, "
            "not a shell command; recorded %d chars" % len(text),
            error=None,
        )

    # Sandbox: run in /tmp, fresh process group so timeout kills children too.
    try:
        proc = subprocess.Popen(
            text,
            shell=True,
            executable="/bin/bash",
            cwd="/tmp",
            env=_sandbox_env(),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            start_new_session=True,  # new process group for reliable kill
        )
    except Exception as e:  # e.g. /bin/bash missing
        return done(False, "", error="spawn failed: %s" % e)

    try:
        stdout, _ = proc.communicate(timeout=timeout)
        rc = proc.returncode
    except subprocess.TimeoutExpired:
        # Kill the whole process group.
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
        except ProcessLookupError:
            pass
        try:
            stdout, _ = proc.communicate(timeout=5)
        except Exception:
            stdout = ""
        return done(
            False,
            stdout or "",
            error="timeout: exceeded %ss, process group killed" % timeout,
        )

    output = stdout or ""
    if rc == 0:
        return done(True, output)
    return done(False, output, error="exit code %d" % rc)


if __name__ == "__main__":
    import json
    import sys

    cases = [
        ("echo hello", 10),
        ("rm -rf /", 10),
    ]
    # Allow ad-hoc: python executor.py "<task>" [timeout]
    if len(sys.argv) > 1:
        cases = [(sys.argv[1], int(sys.argv[2]) if len(sys.argv) > 2 else 600)]
    for task, to in cases:
        print("TASK:", task)
        print(json.dumps(execute_task(task, timeout=to), indent=1))
        print("-" * 40)