subprocess-util/unix_sock_input.py

import socket
from pathlib import Path


def wait_until_command_received(socket_file: Path, commands: list[bytes]) -> bytes:
    """
    Creates a UNIX socket at socket_file.

    Accepts connections on the UNIX socket,
    until one client sends one of the given commands as first message.

    Closes the UNIX socket.

    :returns: The command that was received.
    """

    # INSPIRATION: https://pymotw.com/3/socket/uds.html

    print(f'Listening on socket {socket_file}')
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.bind(str(socket_file))
    sock.listen(1)

    command = accept_loop_until_command_received(sock, commands)

    print(f'Closing socket {socket_file}')
    sock.close()
    socket_file.unlink(missing_ok=False)

    return command


def accept_loop_until_command_received(sock: socket.socket, commands: list[bytes]) -> bytes:
    """
    Uses an open UNIX socket.

    Accepts connections on the UNIX socket,
    until one client sends one of the given commands as first message.

    Does not close the UNIX socket.

    :returns: The command that was received.
    """

    bufsize = max([len(cmd) for cmd in commands]) + len(b'\n')

    while True:
        connection, client_address = sock.accept()
        try:
            b: bytes = connection.recv(bufsize)
            for cmd in commands:
                if b == cmd:
                    print(f'Received "{cmd}".')
                    return cmd
            print(f'Received unknown message: {b}')
        except Exception as e:
            print(f'Error while reading message: {e}')
        finally:
            connection.close()