#!/usr/bin/env python3
"""Sandbox escape hatch for execute_code elevated steps.

Connects to the socket at OLLIE_ELEVATE_SOCKET, sends the command and cwd,
and streams output back in real-time. The adapter listening on that socket
is responsible for running the command outside the Landlock sandbox.

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

Usage: elevate -- <cmd>
"""
import sys
sys.path = sys.path[1:]  # strip script dir — may be a non-navigable fs (e.g. 9P)
import json
import os
import socket
import struct


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 recv_exact(sock, n):
    """Read exactly n bytes from socket."""
    buf = b""
    while len(buf) < n:
        chunk = sock.recv(n - len(buf))
        if not chunk:
            return None
        buf += chunk
    return buf


def main():
    sock_path = os.environ.get("OLLIE_ELEVATE_SOCKET") or default_socket_path()

    if len(sys.argv) < 3:
        print("usage: elevate -- <cmd>", file=sys.stderr)
        sys.exit(1)

    req = json.dumps({"cmd": sys.argv[2], "cwd": os.getcwd()}).encode() + b"\n"

    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    try:
        sock.connect(sock_path)
    except OSError as e:
        print(f"elevate: connect {sock_path}: {e}", file=sys.stderr)
        sys.exit(1)

    with sock:
        sock.sendall(req)
        sock.shutdown(socket.SHUT_WR)

        exit_code = 1
        while True:
            header = recv_exact(sock, 5)
            if header is None:
                break
            frame_type = chr(header[0])
            length = struct.unpack(">I", header[1:5])[0]
            payload = recv_exact(sock, length) if length > 0 else b""
            if payload is None:
                break

            if frame_type == "d":
                sys.stdout.buffer.write(payload)
                sys.stdout.buffer.flush()
            elif frame_type == "x":
                exit_code = int(payload.decode())
                break

    sys.exit(exit_code)


if __name__ == "__main__":
    main()