#!/usr/bin/env python3
"""Patch hermes-agent to add sandbox isolation for dangerous terminal commands.

Inspired by OpenAI Agents SDK's Sandbox Agent concept.

ARCHITECTURE:
  Hermes already has a 5-stage terminal safety engine that classifies commands
  by risk level (read / write / destructive / network / process). This patch
  adds ACTUAL isolation for dangerous commands instead of just "confirm first":

  Risk Level         Before Patch           After Patch
  ─────────────────  ──────────────────────  ────────────────────────────
  Read-only          Direct execution        Direct execution (unchanged)
  Write              Direct execution        Direct execution (unchanged)
  Destructive        Confirm → execute       SANDBOXED execution (namespace isolation + resource limits)
  Network            Confirm → execute       SANDBOXED (network namespace isolated)
  Process mgmt       Confirm → execute       SANDBOXED (PID namespace isolated)

SANDBOX MECHANISM (Linux namespace isolation via unshare):
  1. PID namespace    — sandboxed process can't see/kill host processes
  2. Network namespace — sandboxed process has NO network access (prevents exfil)
  3. Mount namespace  — filesystem is read-only except for /tmp/hermes-sandbox
  4. Resource limits  — CPU time cap, memory cap, no fork bombs
  5. /tmp sandbox    — writable directory for temporary files only

FALLBACK:
  If unshare is unavailable (non-Linux or no CAP_SYS_ADMIN), the sandbox
  gracefully degrades to "confirm + resource limits only" mode.

FILES PATCHED:
  - tools/environments/local.py — wraps _run_bash with sandbox for dangerous cmds
  - tools/approval.py (optional) — marks sandboxed commands in approval output
"""

import sys
import os
import glob
import re
import subprocess
import textwrap


SANDBOX_WRAPPER_CODE = '''
# ── Hermes Bot patch: Sandbox isolation for dangerous commands ──
# Inspired by OpenAI Agents SDK Sandbox Agent concept.

import os
import sys
import resource
import subprocess
import tempfile
import shutil

# Sandbox writable directory
_SANDBOX_TMP = "/tmp/hermes-sandbox"

def _ensure_sandbox_tmp():
    """Create sandbox temp directory if it doesn't exist."""
    os.makedirs(_SANDBOX_TMP, mode=0o700, exist_ok=True)

def _can_use_unshare():
    """Check if unshare with namespaces is available."""
    if sys.platform != "linux":
        return False
    try:
        # Test if we can create a user namespace (cheapest check)
        proc = subprocess.run(
            ["unshare", "--user", "--map-root-user", "true"],
            capture_output=True, timeout=3,
        )
        return proc.returncode == 0
    except (FileNotFoundError, subprocess.TimeoutExpired, OSError):
        return False

_CAN_USE_UNSHARE = _can_use_unshare()

def _apply_resource_limits():
    """Set resource limits for sandboxed processes."""
    # Max 60 seconds CPU time (prevents infinite loops)
    resource.setrlimit(resource.RLIMIT_CPU, (60, 60))
    # Max 512MB memory
    resource.setrlimit(resource.RLIMIT_AS, (512 * 1024 * 1024, 512 * 1024 * 1024))
    # Max 100 processes (prevents fork bombs)
    resource.setrlimit(resource.RLIMIT_NPROC, (100, 100))
    # Max 1000 open files
    resource.setrlimit(resource.RLIMIT_NOFILE, (1000, 1000))

def _sandbox_command_unshare(cmd_string: str) -> str:
    """Wrap command in Linux namespace isolation via unshare.

    Creates: PID namespace + network namespace + mount namespace.
    Filesystem is mostly read-only; /tmp/hermes-sandbox is writable.
    """
    _ensure_sandbox_tmp()
    # Build the sandbox wrapper script
    sandbox_script = textwrap.dedent(f"""\\
        # Set resource limits
        python3 -c "import resource; resource.setrlimit(resource.RLIMIT_CPU, (60, 60)); resource.setrlimit(resource.RLIMIT_AS, (512*1024*1024, 512*1024*1024)); resource.setrlimit(resource.RLIMIT_NPROC, (100, 100)); resource.setrlimit(resource.RLIMIT_NOFILE, (1000, 1000))" 2>/dev/null

        # Remount root as read-only
        mount -o remount,ro / 2>/dev/null || true

        # Create writable tmp
        mkdir -p {_SANDBOX_TMP}
        mount -t tmpfs -o size=256m,nr_inodes=10000 tmpfs {_SANDBOX_TMP} 2>/dev/null || true

        # Bind-mount /tmp into sandbox (read-write overlay)
        mkdir -p {_SANDBOX_TMP}/tmp
        mount --bind /tmp {_SANDBOX_TMP}/tmp 2>/dev/null || true

        # Create a minimal /etc for DNS resolution (read-only copy)
        mkdir -p {_SANDBOX_TMP}/etc
        cp /etc/resolv.conf {_SANDBOX_TMP}/etc/resolv.conf 2>/dev/null || true

        # Execute the actual command with writable tmp
        cd {_SANDBOX_TMP}
        export TMPDIR={_SANDBOX_TMP}
        {cmd_string}
    """)
    # Wrap in unshare with PID + network + mount namespaces
    return f"unshare --pid --net --mount --fork --map-root-user -- bash -c {repr(sandbox_script)}"

def _sandbox_command_fallback(cmd_string: str) -> str:
    """Fallback sandbox: resource limits only (no namespace isolation).

    Used when unshare is not available.
    """
    _ensure_sandbox_tmp()
    limiter = textwrap.dedent(f"""\\
        ulimit -t 60 2>/dev/null          # 60s CPU time
        ulimit -v 524288 2>/dev/null       # 512MB virtual memory
        ulimit -u 100 2>/dev/null          # 100 processes
        cd {_SANDBOX_TMP}
        export TMPDIR={_SANDBOX_TMP}
        {cmd_string}
    """)
    return limiter

def sandbox_wrap(cmd_string: str) -> str:
    """Wrap a dangerous command in sandbox isolation.

    Automatically picks the best available isolation method:
    1. unshare with namespaces (Linux + CAP_SYS_ADMIN)
    2. ulimit resource limits (fallback)
    """
    if _CAN_USE_UNSHARE:
        return _sandbox_command_unshare(cmd_string)
    else:
        return _sandbox_command_fallback(cmd_string)

# Command danger classification for auto-sandbox decisions
_DESTRUCTIVE_PATTERNS = [
    r'\\brm\\b.*(-rf|-r|-fr|/)',       # rm with recursive/force or absolute paths
    r'\\bshred\\b',
    r'\\bdd\\b.*of=',
    r'\\bmkfs\\b',
    r'\\bformat\\b',
    r'\\bwipefs\\b',
    r'\\bchmod\\b.*777',
    r'\\bchown\\b.*-R',
    r'\\bmv\\b.*/(boot|etc|usr|lib|bin|sbin)',
    r'>\\s*/dev/',
    r'\\bkill\\b.*(-9|-s\\s*9|SIGKILL)',
    r'\\bpkill\\b',
    r'\\bkillall\\b',
]

_NETWORK_PATTERNS = [
    r'\\bcurl\\b.*\\bupload\\b',
    r'\\bwget\\b',
    r'\\bnc\\b',
    r'\\bncat\\b',
    r'\\bnmap\\b',
    r'\\bpython[23]?\\b.*socket\\b',
    r'\\bpython[23]?\\b.*requests\\.(post|put)',
    r'\\bssh\\b.*(@|connect)',
    r'\\bscp\\b',
    r'\\brsync\\b',
    r'\\bnc\\b.*-e',
]

def should_sandbox(command: str) -> bool:
    """Determine if a command should be sandboxed based on pattern matching.

    Returns True for destructive, network-exfiltration, or process-management
    commands that benefit from isolation.
    """
    for pattern in _DESTRUCTIVE_PATTERNS + _NETWORK_PATTERNS:
        if re.search(pattern, command, re.IGNORECASE):
            return True
    return False

'''


def patch_file(filepath: str) -> bool:
    """Patch local.py to add sandbox wrapping for dangerous commands."""
    with open(filepath, "r") as f:
        content = f.read()

    if "sandbox_wrap" in content:
        print(f"  Already patched: {filepath}")
        return True

    applied = False

    # ── Patch: Add sandbox wrapper code after imports ──
    # Find a good insertion point: after the last import
    import_section_end = content.rfind('\n\n')
    if import_section_end > 0 and import_section_end < len(content) // 2:
        insertion_point = import_section_end + 2
    else:
        # Fallback: after the module docstring
        docstring_end = content.find('"""', content.find('"""') + 3)
        if docstring_end > 0:
            insertion_point = docstring_end + 3
        else:
            insertion_point = 0

    content = content[:insertion_point] + "\n" + SANDBOX_WRAPPER_CODE + "\n" + content[insertion_point:]
    applied = True
    print("  [local.py] Added sandbox wrapper code")

    # ── Patch: Hook sandbox into _run_bash method ──
    # Find the _run_bash method and wrap commands that should be sandboxed
    # We look for the line where cmd_string is passed to bash and inject
    # a sandbox check before it.

    # Pattern: in the _run_bash method, right before subprocess.Popen
    # We want to wrap cmd_string if should_sandbox() returns True

    # Find the Popen call in _run_bash
    old_popen = "proc = subprocess.Popen(\n            args,"
    new_popen = (
        "        # Hermes Bot patch: auto-sandbox dangerous commands\n"
        "        if should_sandbox(cmd_string):\n"
        "            original_cmd = cmd_string\n"
        "            cmd_string = sandbox_wrap(cmd_string)\n"
        "            args = [bash, \"-c\", cmd_string]\n"
        "            logger.info(\"Sandbox isolation applied for dangerous command\")\n"
        "        proc = subprocess.Popen(\n"
        "            args,"
    )

    if old_popen in content:
        content = content.replace(old_popen, new_popen, 1)
        applied = True
        print("  [local.py] Hooked sandbox into _run_bash subprocess.Popen")
    else:
        # Try alternative pattern (different indentation)
        alt_popen = "proc = subprocess.Popen(\n                args,"
        if alt_popen in content:
            content = content.replace(
                alt_popen,
                "            # Hermes Bot patch: auto-sandbox dangerous commands\n"
                "            if should_sandbox(cmd_string):\n"
                "                cmd_string = sandbox_wrap(cmd_string)\n"
                "                args = [bash, \"-c\", cmd_string]\n"
                "                logger.info(\"Sandbox isolation applied for dangerous command\")\n"
                "            proc = subprocess.Popen(\n"
                "                args,",
                1,
            )
            applied = True
            print("  [local.py] Hooked sandbox into _run_bash (alt pattern)")

    if applied:
        with open(filepath, "w") as f:
            f.write(content)
        return True
    else:
        print(f"  WARNING: Could not hook sandbox into {filepath}", file=sys.stderr)
        # Still save with the sandbox code added
        with open(filepath, "w") as f:
            f.write(content)
        return False


if __name__ == "__main__":
    candidates = [
        "/app/hermes-agent/tools/environments/local.py",
    ]
    candidates.extend(
        glob.glob("/app/venv/lib/**/tools/environments/local.py", recursive=True)
    )

    filepath = None
    for c in candidates:
        if os.path.isfile(c):
            filepath = c
            break

    if not filepath:
        print("WARNING: local.py not found", file=sys.stderr)
        print(f"Checked: {candidates}", file=sys.stderr)
        sys.exit(0)

    ok = patch_file(filepath)
    if ok:
        print(f"\nSandbox isolation patch applied to {filepath}")
        print(f"  unshare available: {_CAN_USE_UNSHARE if 'sandbox_wrap' in dir() else 'unknown (runtime)'}")
        print(f"  Sandbox tmp: {_SANDBOX_TMP if '_SANDBOX_TMP' in dir() else '/tmp/hermes-sandbox'}")
    else:
        print("Patch partially failed", file=sys.stderr)
        sys.exit(1)