unixshim

#!/usr/bin/python3

import json
import os
import re
import selectors
import socket
import struct
import sys


def ext_id():
    chrome = re.match(r'^chrome-extension://([^/]+)', sys.argv[1])
    if (chrome):
        return chrome.group(1)
    else:  # Firefox
        return sys.argv[2]

def from_ext(stdin):
    header = stdin.buffer.read(HEADER_SIZE)
    if len(header) != HEADER_SIZE:
        sys.exit(1)

    json_content_size = struct.unpack(HEADER_FORMAT, header)[0]
    json_content = stdin.buffer.read(json_content_size)
    if len(json_content) != json_content_size:
        sys.exit(1)

    content = json.loads(json_content.decode('utf-8')).encode('utf-8')

    with socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) as from_ext_socket:
        from_ext_socket.connect(FROM_EXT_SOCKET_PATH)
        from_ext_socket.sendall(content)

def to_ext(server):
    client, dummy_address = server.accept()
    try:
        with client.makefile() as client_stream:
            content = client_stream.buffer.read()
    finally:
        client.close()

    json_content = json.dumps(content.decode('utf8')).encode('utf-8')
    header = struct.pack(HEADER_FORMAT, len(json_content))
    os.writev(sys.stdout.fileno(), (header, json_content))

def cleanup():
    try:
        os.remove(TO_EXT_SOCKET_PATH)
    except FileNotFoundError:
        pass

def run():
    os.makedirs(UNIXSHIM_SUBDIR, exist_ok=True)
    cleanup()

    to_ext_socket = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    to_ext_socket.bind(TO_EXT_SOCKET_PATH)
    to_ext_socket.listen()

    sel = selectors.DefaultSelector()
    sel.register(sys.stdin, selectors.EVENT_READ, from_ext)
    sel.register(to_ext_socket, selectors.EVENT_READ, to_ext)

    while True:
        for key, dummy_mask in sel.select():
            key.data(key.fileobj)


HEADER_SIZE = 4
HEADER_FORMAT = '@I'
UNIXSHIM_SUBDIR = os.path.join(os.environ['UNIXSHIM_DIR'], ext_id())
FROM_EXT_SOCKET_PATH = os.path.join(UNIXSHIM_SUBDIR, 'from-ext')
TO_EXT_SOCKET_PATH = os.path.join(UNIXSHIM_SUBDIR, 'to-ext')

try:
    run()
finally:
    cleanup()