#!/usr/bin/python3

import array
import json
import os
import socket
import subprocess
import sys


def load_fds(sock, msglen):
    fds = array.array("i")   # Array of ints
    msg, ancdata, _, addr = sock.recvmsg(msglen, socket.CMSG_LEN(253 * fds.itemsize))
    for cmsg_level, cmsg_type, cmsg_data in ancdata:
        if (cmsg_level == socket.SOL_SOCKET and cmsg_type == socket.SCM_RIGHTS):
            # Append data, ignoring any truncated integers at the end.
            fds.frombytes(cmsg_data[:len(cmsg_data) - (len(cmsg_data) % fds.itemsize)])
    return json.loads(msg), list(fds), addr

def setup_stdio():
    with socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) as sock:
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_PASSCRED, 1)
        sock.connect("/run/osbuild/api/osbuild")
        req = {'method': 'setup-stdio'}
        sock.send(json.dumps(req).encode('utf-8'))
        msg, fds, _ = load_fds(sock, 1024)
        for io in ['stdin', 'stdout', 'stderr']:
            target = getattr(sys, io)
            source = fds[msg[io]]
            os.dup2(source, target.fileno())
            os.close(source)


if __name__ == "__main__":
    setup_stdio()

    r = subprocess.run(sys.argv[1:], check=False)
    sys.exit(r.returncode)