ollie/scripts/x/elevate-adapter-superpowerd

#!/usr/bin/env python3
"""Elevation adapter: runs commands outside the Landlock sandbox with streaming I/O.

Creates a Unix socket at OLLIE_ELEVATE_SOCKET, accepts {cmd, cwd} requests
from x/elevate, sends them to superpowerd via ExecuteCommandStreaming with
fd-passing, and streams the output back to the client in real-time.

Protocol (client-facing, streaming):
  Client sends: {"cmd": ..., "cwd": ...}\n then shutdown(WR)
  Server sends frames: <type:1><length:4 big-endian><payload>
    type 'd' = output data (combined stdout+stderr)
    type 'x' = exit code (payload is ASCII integer)

Usage:
    superpowers run-session -- python3 elevate-adapter-superpowerd

Environment:
    OLLIE_ELEVATE_SOCKET       path for the listening socket
                               (default: $XDG_RUNTIME_DIR/ollie/elevate.sock)
    SUPERPOWERD_SESSION_TOKEN  set automatically by superpowers run-session
    SUPERPOWERD_SOCKET         override superpowerd socket path
    SUPERPOWERD_SOCKET_DIR     directory containing superpowerd.sock
"""
import sys
sys.path = sys.path[1:]  # strip script dir — may be a non-navigable fs (e.g. 9P)
import array
import json
import os
import signal
import socket
import struct
import subprocess
import threading

_token_lock = threading.Lock()


def refresh_token():
    """Obtain a fresh session token from superpowerd via run-session.

    Returns the new token string, or None on failure.
    Updates os.environ so subsequent calls use the fresh token.
    """
    try:
        result = subprocess.run(
            ["superpowers", "run-session", "--child-session", "--",
             "sh", "-c", "echo $SUPERPOWERD_SESSION_TOKEN"],
            capture_output=True, text=True, timeout=10,
        )
        token = result.stdout.strip()
        if token and result.returncode == 0:
            os.environ["SUPERPOWERD_SESSION_TOKEN"] = token
            print(f"elevate-adapter-superpowerd: refreshed session token", file=sys.stderr)
            return token
    except (subprocess.TimeoutExpired, OSError) as e:
        print(f"elevate-adapter-superpowerd: token refresh failed: {e}", file=sys.stderr)
    return None


def default_socket_path():
    xdg = os.environ.get("XDG_RUNTIME_DIR") or f"/run/user/{os.getuid()}"
    return os.path.join(xdg, "ollie", "elevate.sock")


def sp_socket_path():
    if p := os.environ.get("SUPERPOWERD_SOCKET"):
        return p
    if d := os.environ.get("SUPERPOWERD_SOCKET_DIR"):
        return d + "/superpowerd.sock"
    if x := os.environ.get("XDG_RUNTIME_DIR"):
        return x + "/superpowerd/superpowerd.sock"
    return f"/run/user/{os.getuid()}/superpowerd/superpowerd.sock"


def send_frame(conn, frame_type, payload):
    """Send a framed message: <type:1><length:4 big-endian><payload>"""
    data = payload if isinstance(payload, bytes) else payload.encode()
    header = frame_type.encode() + struct.pack(">I", len(data))
    conn.sendall(header + data)


def send_error(conn, msg, exit_code=1):
    """Send an error message and exit frame."""
    send_frame(conn, "d", msg + "\n")
    send_frame(conn, "x", str(exit_code))


def _is_auth_error(sp_data):
    """Check if superpowerd response is an authentication/session error."""
    if not sp_data:
        return False
    try:
        resp = json.loads(sp_data)
        err = resp.get("error")
        if isinstance(err, dict):
            msg = err.get("message", "").lower()
            return "invalid session token" in msg or "session not found" in msg
    except (json.JSONDecodeError, ValueError):
        pass
    return False


def _execute_on_superpowerd(conn, cmd, cwd, token):
    """Execute a command via superpowerd. Returns (sp_data, True) on success/normal
    error, or (sp_data, False) if the caller should retry (auth error, no output streamed)."""
    sp_path = sp_socket_path()

    # Connect to superpowerd via SEQPACKET
    sp = socket.socket(socket.AF_UNIX, socket.SOCK_SEQPACKET)
    try:
        sp.connect(sp_path)
    except OSError as e:
        send_error(conn, f"elevate-adapter: connect {sp_path}: {e}")
        return None, True

    # Create pipes for stdout+stderr (merged) and stdin (/dev/null)
    stdin_r = os.open("/dev/null", os.O_RDONLY)
    stdout_r, stdout_w = os.pipe()

    # Build JSON-RPC request
    sp_req = json.dumps({
        "jsonrpc": "2.0",
        "method": "ExecuteCommandStreaming",
        "id": 1,
        "params": {
            "auth": {"session_token": token},
            "command": cmd,
            "working_dir": cwd,
            "approval_timeout": 7200,
        },
    }).encode()

    # Send request + 3 fds (stdin, stdout, stderr) via SCM_RIGHTS
    fds = array.array("i", [stdin_r, stdout_w, stdout_w])
    ancdata = [(socket.SOL_SOCKET, socket.SCM_RIGHTS, fds)]
    try:
        sp.sendmsg([sp_req], ancdata)
    except OSError as e:
        os.close(stdin_r)
        os.close(stdout_r)
        os.close(stdout_w)
        sp.close()
        send_error(conn, f"elevate-adapter: sendmsg: {e}")
        return None, True

    # Close write-ends in this process (superpowerd has them now)
    os.close(stdin_r)
    os.close(stdout_w)

    # Stream output from the pipe to the elevate client
    stream_done = threading.Event()
    client_broken = threading.Event()
    output_received = threading.Event()

    def stream_output():
        try:
            while True:
                chunk = os.read(stdout_r, 4096)
                if not chunk:
                    break
                output_received.set()
                send_frame(conn, "d", chunk)
        except (BrokenPipeError, OSError):
            client_broken.set()
        finally:
            os.close(stdout_r)
            stream_done.set()

    t = threading.Thread(target=stream_output, daemon=True)
    t.start()

    # Wait for superpowerd's response (comes after command exits)
    try:
        sp_data = sp.recv(1024 * 1024)
    except OSError:
        sp_data = b""
    sp.close()

    # Wait for streaming thread to finish
    stream_done.wait()

    if client_broken.is_set():
        return None, True

    # If auth error and no output was streamed, signal retry
    if _is_auth_error(sp_data) and not output_received.is_set():
        return sp_data, False

    return sp_data, True


def handle(conn):
    with conn:
        try:
            _handle_inner(conn)
        except Exception as e:
            try:
                send_error(conn, f"elevate-adapter: internal error: {e}")
            except OSError:
                pass


def _handle_inner(conn):
    data = b""
    while chunk := conn.recv(65536):
        data += chunk
        if b"\n" in data:
            break

    try:
        req = json.loads(data)
    except json.JSONDecodeError as e:
        send_error(conn, f"elevate-adapter: bad request: {e}")
        return

    cmd = req.get("cmd", "")
    cwd = req.get("cwd", os.getcwd())

    token = os.environ.get("SUPERPOWERD_SESSION_TOKEN", "")

    # First attempt
    sp_data, done = _execute_on_superpowerd(conn, cmd, cwd, token)
    if done:
        _finish_response(conn, sp_data)
        return

    # Auth error — refresh token and retry once
    with _token_lock:
        # Re-read in case another thread already refreshed
        current = os.environ.get("SUPERPOWERD_SESSION_TOKEN", "")
        if current == token:
            new_token = refresh_token()
        else:
            new_token = current

    if not new_token:
        # Refresh failed — report the original auth error
        _finish_response(conn, sp_data)
        return

    sp_data, _ = _execute_on_superpowerd(conn, cmd, cwd, new_token)
    _finish_response(conn, sp_data)


def _finish_response(conn, sp_data):
    """Parse superpowerd response and send exit code frame to client."""
    if sp_data is None:
        return

    exit_code = 1
    try:
        sp_resp = json.loads(sp_data)
        rpc_err = sp_resp.get("error")
        if rpc_err:
            msg = rpc_err.get("message", str(rpc_err)) if isinstance(rpc_err, dict) else str(rpc_err)
            send_frame(conn, "d", f"elevate-adapter: {msg}\n".encode())
        else:
            result = sp_resp.get("result", {})
            exit_code = result.get("exit_code", 0)
    except (json.JSONDecodeError, ValueError):
        pass

    send_frame(conn, "x", str(exit_code))


def main():
    if "SUPERPOWERD_SESSION_TOKEN" not in os.environ:
        os.execvp("superpowers", ["superpowers", "run-session", "--",
                                   sys.executable, __file__] + sys.argv[1:])

    sock_path = os.environ.get("OLLIE_ELEVATE_SOCKET") or default_socket_path()
    sock_path = os.path.expanduser(sock_path)
    os.makedirs(os.path.dirname(sock_path), exist_ok=True)

    if os.path.exists(sock_path):
        os.unlink(sock_path)

    srv = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    srv.bind(sock_path)
    srv.listen()

    def shutdown(sig, frame):
        srv.close()
        if os.path.exists(sock_path):
            os.unlink(sock_path)
        sys.exit(0)

    signal.signal(signal.SIGTERM, shutdown)
    signal.signal(signal.SIGINT, shutdown)

    print(f"elevate-adapter-superpowerd: listening on {sock_path}", file=sys.stderr)

    while True:
        try:
            conn, _ = srv.accept()
        except OSError:
            break
        threading.Thread(target=handle, args=(conn,), daemon=True).start()


if __name__ == "__main__":
    main()