added potion-polyglot worker folder w/o repos

worker-toolkit-potion-polyglot/scripts/dnsjail.py

@@ -0,0 +1,98 @@
+"""Apply the DNS jail to a trial container: the model endpoint resolves, nothing else does.
+
+Opt-in with RACCOON_DNS_JAIL=1. Runs from the agent's own turn rather than from a compose
+overlay — the allowlist comes from the proxy URL this process already holds (plus any hosts
+RACCOON_DNS_JAIL_ALLOW adds), so nothing has to be injected into the container, and the jail works on every harbor backend. Deliberately
+after agent-setup: a harness that downloads its CLI there still reaches the network to do it.
+"""
+
+import logging
+import os
+import shlex
+from typing import Any
+
+JAIL = "/usr/local/bin/raccoon-dns-jail"
+_NO_SCRIPT = "raccoon-dns-jail: not in this image"
+
+_URL_VARS = (
+    "ANTHROPIC_BASE_URL", "OPENAI_BASE_URL", "GOOGLE_GEMINI_BASE_URL",
+    "HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy",
+)
+
+_EXTRA_VAR = "RACCOON_DNS_JAIL_ALLOW"
+
+_log = logging.getLogger(__name__)
+
+
+def _host(url: str) -> str:
+    """Hostname out of a URL, or "" when it is not a plain hostname we can allow."""
+    h = url.split("://", 1)[-1].split("/", 1)[0].rsplit("@", 1)[-1].split(":", 1)[0]
+    if not h or h.startswith((".", "-")) or h.endswith(".") or not all(
+        c.isascii() and (c.isalnum() or c in ".-") for c in h
+    ):
+        return ""
+    # An IP-literal endpoint (a loopback proxy shim, say) needs no DNS at all, and a
+    # --server rule for it would only be checked by a PTR query the catch-all answers.
+    if all(part.isdigit() for part in h.split(".")):
+        return ""
+    return h
+
+
+def dns_jail_allowlist() -> tuple[list[str], list[str]]:
+    """(required, advisory).
+
+    Required = the hosts this process's own env says the agent will dial; every one must
+    resolve through the jail or no jail is applied, because a host the agent needs and
+    cannot resolve is a dead trial. Advisory = whatever RACCOON_DNS_JAIL_ALLOW adds, which
+    only warns: an added host that CNAMEs outside the allowlist cannot resolve through the
+    catch-all, and must not take the whole jail down with it.
+    """
+    required: list[str] = []
+    for var in _URL_VARS:
+        h = _host(os.environ.get(var) or "")
+        if h and h not in required:
+            required.append(h)
+    advisory: list[str] = []
+    for entry in (os.environ.get(_EXTRA_VAR) or "").replace(",", " ").split():
+        # Bare hostnames only: a URL silently truncated to its first path segment would
+        # allow a name nobody asked for and block the one they meant.
+        h = "" if ("/" in entry or ":" in entry) else _host(entry)
+        if not h:
+            _log.warning("DNS jail: ignoring unusable %s entry %r", _EXTRA_VAR, entry)
+        elif h not in required and h not in advisory:
+            advisory.append(h)
+    return required, advisory
+
+
+def dns_jail_enabled() -> bool:
+    return os.environ.get("RACCOON_DNS_JAIL") == "1"
+
+
+async def apply_dns_jail(agent: Any, environment: Any) -> None:
+    """No-op unless enabled; leaves the container's DNS untouched on any doubt."""
+    if not dns_jail_enabled():
+        return
+    required, advisory = dns_jail_allowlist()
+    allow = " ".join(required)
+    # A blank allowlist means no model endpoint was found: jailing would strand the agent.
+    if not allow:
+        _log.warning("DNS jail: no usable model endpoint — the trial keeps normal network access")
+        return
+    try:
+        result = await agent.exec_as_root(
+            environment,
+            command=(
+                f"if [ -x {JAIL} ]; then DNSJAIL_ALLOW={shlex.quote(allow)} "
+                f"DNSJAIL_ALLOW_EXTRA={shlex.quote(' '.join(advisory))} {JAIL}; "
+                f'else echo "{_NO_SCRIPT}"; fi'
+            ),
+        )
+    except Exception as exc:  # a jail that cannot be applied must not fail the trial
+        _log.warning("DNS jail: could not apply (%s) — the trial keeps normal network access", exc)
+        return
+    # An image frozen before this feature has nothing to invoke. Say so: a launcher that
+    # believes the network is restricted when it is not is worse than no jail at all.
+    if _NO_SCRIPT in (getattr(result, "stdout", "") or ""):
+        _log.warning(
+            "DNS jail: this task's image ships no resolver — the trial keeps normal network access"
+        )