status-indicators/push_page.py

#!/usr/bin/env python3
"""Import status-page.companionconfig into Companion's page 19 without clicking through the UI.

Does exactly what Import/Export -> Import -> "import into page 19" does, over the
same internal API the Companion web UI uses (tRPC over a WebSocket at /trpc):
upload the file in chunks, check its SHA-1, then import it as a single page.

Safety: it refuses unless the target page is empty or already the status page
(its name contains "status"), so it cannot overwrite one of your other pages.

  python3 push_page.py                         # na-ingest.local:8000, page and file from config.json
  python3 push_page.py --host 10.6.3.222 --page 19 --file status-page.companionconfig

Stdlib only (a minimal WebSocket client is included). Written against Companion 5.0.7.
"""
import argparse
import base64
import hashlib
import json
import os
import socket
import struct
import sys
import time
import urllib.parse
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
CHUNK = 256 * 1024


class WebSocket:
    def __init__(self, host, port, path, timeout=15):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        key = base64.b64encode(os.urandom(16)).decode()
        self.sock.sendall(("GET %s HTTP/1.1\r\nHost: %s:%d\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n"
                           "Sec-WebSocket-Key: %s\r\nSec-WebSocket-Version: 13\r\n\r\n"
                           % (path, host, port, key)).encode())
        buf = b""
        while b"\r\n\r\n" not in buf:
            chunk = self.sock.recv(4096)
            if not chunk:
                raise ConnectionError("connection closed during WebSocket handshake")
            buf += chunk
        head, self.buf = buf.split(b"\r\n\r\n", 1)
        if b" 101 " not in head.split(b"\r\n", 1)[0]:
            raise ConnectionError("WebSocket handshake refused: %s" % head.split(b"\r\n", 1)[0].decode())

    def _read(self, n):
        while len(self.buf) < n:
            chunk = self.sock.recv(65536)
            if not chunk:
                raise ConnectionError("connection closed")
            self.buf += chunk
        data, self.buf = self.buf[:n], self.buf[n:]
        return data

    def _send(self, opcode, payload):
        n = len(payload)
        header = bytes([0x80 | opcode])
        if n < 126:
            header += bytes([0x80 | n])
        elif n < 65536:
            header += bytes([0x80 | 126]) + struct.pack(">H", n)
        else:
            header += bytes([0x80 | 127]) + struct.pack(">Q", n)
        mask = os.urandom(4)
        masked = bytes(b ^ mask[i & 3] for i, b in enumerate(payload))
        self.sock.sendall(header + mask + masked)

    def send(self, text):
        self._send(0x1, text.encode())

    def recv(self):
        message = b""
        while True:
            b1, b2 = self._read(2)
            fin, opcode, n = b1 & 0x80, b1 & 0x0F, b2 & 0x7F
            if n == 126:
                n = struct.unpack(">H", self._read(2))[0]
            elif n == 127:
                n = struct.unpack(">Q", self._read(8))[0]
            payload = self._read(n)
            if opcode == 0x9:                       # ping
                self._send(0xA, payload)
                continue
            if opcode == 0x8:
                raise ConnectionError("server closed the WebSocket")
            message += payload
            if fin:
                return message.decode()

    def close(self):
        try:
            self._send(0x8, b"")
        finally:
            self.sock.close()


class Trpc:
    def __init__(self, ws):
        self.ws, self.next_id = ws, 1

    def mutation(self, path, data):
        rid = self.next_id
        self.next_id += 1
        self.ws.send(json.dumps({"id": rid, "jsonrpc": "2.0", "method": "mutation",
                                 "params": {"path": path, "input": data}}))
        while True:
            msg = json.loads(self.ws.recv())
            if msg.get("id") != rid:
                continue
            if "error" in msg:
                err = msg["error"]
                raise RuntimeError("%s failed: %s" % (path, err.get("message", err)))
            return msg.get("result", {}).get("data")


def page_name(host, port, page):
    url = "http://%s:%d/api/variable/internal/page_number_%d_name/value" % (host, port, page)
    with urllib.request.urlopen(url, timeout=5) as r:
        return r.read().decode()


def main():
    cfg = json.load(open(os.path.join(HERE, "config.json")))
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("--host", default="na-ingest.local")
    ap.add_argument("--port", type=int, default=8000)
    ap.add_argument("--page", type=int, default=cfg["page"])
    ap.add_argument("--file", default=os.path.join(HERE, "status-page.companionconfig"))
    ap.add_argument("--force", action="store_true", help="overwrite the page whatever its name")
    args = ap.parse_args()

    name = page_name(args.host, args.port, args.page)
    plain = name.replace("\\n", " ").strip()
    empty = plain.upper() in ("", "PAGE", "%d PAGE" % args.page)
    if not (empty or "status" in plain.lower() or args.force):
        sys.exit("page %d is called %r - not empty and not the status page; refusing (use --force)"
                 % (args.page, plain))

    data = open(args.file, "rb").read()
    ws = WebSocket(args.host, args.port, "/trpc")
    try:
        api = Trpc(ws)
        # Only one upload at a time; an abandoned one expires after a short idle time.
        for attempt in range(40):
            try:
                session = api.mutation("importExport.prepareImport.start",
                                       {"name": os.path.basename(args.file), "size": len(data)})
                break
            except RuntimeError as e:
                if "already exists" not in str(e) or attempt == 39:
                    raise
                if attempt == 0:
                    print("another upload is still open in Companion; waiting for it to expire...")
                time.sleep(3)
        for offset in range(0, len(data), CHUNK):
            api.mutation("importExport.prepareImport.uploadChunk",
                         {"sessionId": session, "offset": offset,
                          "data": base64.b64encode(data[offset:offset + CHUNK]).decode()})
        summary = api.mutation("importExport.prepareImport.complete",
                               {"sessionId": session, "expectedChecksum": hashlib.sha1(data).hexdigest(),
                                "userData": None})
        if isinstance(summary, list) and summary and summary[0]:
            sys.exit("Companion could not read the file: %s" % summary[0])
        api.mutation("importExport.importSinglePage",
                     {"targetPage": args.page, "sourcePage": 1, "connectionIdRemapping": {}})
    finally:
        ws.close()
    print("imported %s into page %d (was %r, now %r)" % (os.path.basename(args.file), args.page, plain,
                                                         page_name(args.host, args.port, args.page)))


if __name__ == "__main__":
    main()