roro9stack/scripts/rdbg.py

#!/usr/bin/env python3
"""The Debug Console, over Wi-Fi (TCP 2323, see ADR 0010). Switch it on first: Settings > Debug Console.

Usage: scripts/rdbg.py [-H host] [-t token] [-b] [command ...]
  no command   interactive: type commands, see every console line live; Ctrl-D or `quit` leaves
  command      runs it and prints what follows, until the console has been quiet for a moment
  -H host      the device's IP (Settings > Debug Console), default $RORO_OTA_HOST
  -t token     the device's token (also --token), as Settings > Debug Console shows it: dashes and
               case don't matter. Otherwise $RORO_DEBUG_TOKEN, otherwise ~/.config/roro9stack/debug-token
  -b           also print the backlog the device sends on connecting (boot messages and so on)

Commands handled here as well as on the device:
  crash                    the last crash, with its backtrace decoded (scripts/decode_backtrace.sh)
  coredump [file]          fetch the core dump (default core-<date>.bin) and decode it with esp-coredump
  get <card path> [file]   copy a file from the SD card
  put <file> [card path]   copy a file to the SD card (default /updates/<name>), checked with SHA-256
  screenshot [file.png]    what the screen shows (default screen-<date>.png), at 2x
  reset                    restart at once, even if the main loop is stuck
The token never crosses the network: the device sends a challenge, and this answers with its HMAC.
"""
import gzip
import hashlib
import hmac
import os
import re
import select
import struct
import subprocess
import zlib
import socket
import sys
import time
import urllib.request

PORT = 2323
TOKEN_FILE = os.path.expanduser("~/.config/roro9stack/debug-token")
RELEASES = "https://git.twis.la/twisla/roro9stack/releases/download"


def tidy_token(typed):
    """As the device stores it (lib/debug/src/debug_auth.cpp): no dashes or spaces, in capitals."""
    tidy = "".join(c for c in typed if c not in "- \t\r\n").upper()
    return tidy.replace("O", "0").replace("I", "1").replace("L", "1")  # read as the digits they look like


def answer_for(token, nonce):
    """What the device expects back for a challenge: HMAC-SHA256 of the nonce, keyed by the token."""
    return hmac.new(token.encode(), nonce, hashlib.sha256).hexdigest()


def find_token(given):
    token = given or os.environ.get("RORO_DEBUG_TOKEN")
    if not token and os.path.exists(TOKEN_FILE):
        token = open(TOKEN_FILE).read()
    token = tidy_token(token or "")
    if not token:
        sys.exit("No token: pass -t <token>, set RORO_DEBUG_TOKEN, or put it in " + TOKEN_FILE +
                 "\n(the device shows it in Settings > Debug Console)")
    return token


def log_in(sock, token):
    """Answers the device's challenge; returns the banner line, or exits saying what went wrong."""
    first = read_until(sock, b"\n", 10)
    if first is None:
        sys.exit("device: nothing came back. Is the console switched on (Settings > Debug Console)?")
    line = first.decode(errors="replace").strip()
    if line.startswith("locked"):
        sys.exit("device: closed for a minute after too many wrong tokens")
    challenge = re.search(r"challenge ([0-9a-f]{32})$", line)
    if not challenge:
        sys.exit("device: no challenge (an older firmware?): " + line[:60])
    sock.sendall((answer_for(token, bytes.fromhex(challenge.group(1))) + "\n").encode())
    banner = read_until(sock, b"\n", 15)
    if banner is None:
        sys.exit("device: no answer")
    if banner.startswith(b"denied"):
        sys.exit("device: wrong token (Settings > Debug Console shows the right one)")
    return banner


def read_until(sock, marker, timeout):
    """Everything up to and including `marker`, or None on a timeout or a closed connection."""
    sock.settimeout(timeout)
    data = b""
    while marker not in data:
        try:
            chunk = sock.recv(4096)
        except socket.timeout:
            return None
        if not chunk:
            return None
        data += chunk
    return data


def read_until_quiet(sock, quiet, out):
    """Copies everything the device sends to `out` until nothing has arrived for `quiet` seconds."""
    sock.settimeout(quiet)
    while True:
        try:
            data = sock.recv(4096)
        except socket.timeout:
            return True
        if not data:
            return False
        if out:
            out.write(data.decode(errors="replace"))
            out.flush()


SCRIPTS = os.path.dirname(os.path.abspath(__file__))


def run(sock, command, out=sys.stdout):
    """Sends one command and returns its reply (also copied to `out`)."""
    sock.sendall((command + "\n").encode())
    # The main loop echoes "> command" when it runs it; the reply follows.
    echoed = read_until(sock, f"> {command}\n".encode(), 15)
    if echoed is None:
        sys.exit("device: the command never ran")
    reply = echoed.decode(errors="replace").split(f"> {command}\n", 1)[1]

    class Tee:
        def write(self, text):
            nonlocal reply
            reply += text
            if out:
                out.write(text)

        def flush(self):
            if out:
                out.flush()

    if out:
        out.write(reply)
    try:
        read_until_quiet(sock, 1.5, Tee())
    except BrokenPipeError:  # e.g. piped into head
        pass
    return reply


def crash_firmware(reply):
    """The archived-ELF key for the crashed firmware: its ELF digest, or else its version."""
    sha = re.search(r"elf sha256 ([0-9a-f]{8,})", reply)
    if sha:
        return sha.group(1)
    version = re.search(r"last one in (\S+)", reply)
    return version.group(1) if version else None


def have_elf(reply):
    """Makes sure .pio/elves/ holds the ELF of the firmware that crashed: a release's is on Gitea."""
    folder = os.path.join(os.path.dirname(SCRIPTS), ".pio", "elves")
    key = crash_firmware(reply)
    version = re.search(r"last one in (\S+)", reply)
    version = version.group(1) if version else None
    names = os.listdir(folder) if os.path.isdir(folder) else []
    if not key or any(key in n for n in names) or not version or not re.fullmatch(r"v\d+\.\d+\.\d+", version):
        return  # there already, or not a release: only tags are published
    url = f"{RELEASES}/{version}/roro9stack-{version}.elf.gz"
    try:
        elf = gzip.decompress(urllib.request.urlopen(url, timeout=60).read())
    except Exception as e:
        return print(f"(no ELF for {version} here, and none fetched from {url}: {e})")
    digest = hashlib.sha256(elf).hexdigest()[:16]
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, f"{version}.{digest}.elf")  # as scripts/version.py names them
    open(path, "wb").write(elf)
    print(f"(fetched the ELF of {version} from its release: {os.path.relpath(path)})")


def crash(sock):
    reply = run(sock, "crash")
    trace = re.search(r"backtrace((?: 0x[0-9a-f]+)+)", reply)
    key = crash_firmware(reply)
    if trace and key:
        have_elf(reply)
        print()
        subprocess.call([os.path.join(SCRIPTS, "decode_backtrace.sh"), key] + trace.group(1).split())


def coredump(sock, path):
    info = run(sock, "crash", out=None)
    have_elf(info)
    sock.sendall(b"coredump get\n")
    # One buffer throughout: the header, the size and the first bytes often share a packet.
    buf = b""
    sock.settimeout(15)
    while b"coredump: none" not in buf and not re.search(rb"coredump: data (\d+)\n", buf):
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit("device: closed the connection")
        buf += chunk
    header = re.search(rb"coredump: data (\d+)\n", buf)
    if not header:
        sys.exit("device: no core dump in flash")
    size = int(header.group(1))
    data = buf[header.end():]
    while len(data) < size:
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit(f"device: the connection closed after {len(data)} of {size} bytes")
        data += chunk
    data = data[:size]
    path = path or time.strftime("core-%Y%m%d-%H%M%S.bin")
    open(path, "wb").write(data)
    print(f"saved {path} ({size} bytes)")
    key = crash_firmware(info)
    if key:
        subprocess.call([os.path.join(SCRIPTS, "decode_coredump.sh"), path, key])


def binary(sock, command, tag):
    """Sends a binary command; returns (header words, payload) or exits with the device's error."""
    sock.sendall((command + "\n").encode())
    buf = b""
    sock.settimeout(30)
    pattern = re.compile(tag.encode() + rb": (data|ready|rgb332|error)([^\n]*)\n")
    while not (m := pattern.search(buf)):
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit("device: closed the connection")
        buf += chunk
    if m.group(1) == b"error":
        sys.exit(f"device: {tag}: error{m.group(2).decode()}")
    return m.group(1).decode(), m.group(2).decode().split(), buf[m.end():]


def receive(sock, have, size):
    while len(have) < size:
        chunk = sock.recv(65536)
        if not chunk:
            sys.exit(f"device: the connection closed after {len(have)} of {size} bytes")
        have += chunk
    return have[:size]


def get(sock, remote, local):
    _, words, rest = binary(sock, f"get {remote}", "get")
    start, size = time.time(), int(words[0])
    data = receive(sock, rest, size)
    local = local or os.path.basename(remote)
    open(local, "wb").write(data)
    print(f"saved {local} ({size} bytes, {size / 1024 / max(time.time() - start, 0.001):.0f} KB/s)")


def put(sock, local, remote):
    data = open(local, "rb").read()
    remote = remote or "/updates/" + os.path.basename(local)
    digest = hashlib.sha256(data).hexdigest()
    binary(sock, f"put {remote} {len(data)} {digest}", "put")
    start = time.time()
    sock.sendall(data)
    answer = read_until(sock, b"\n", 30) or b"put: error no answer"
    answer = answer.decode(errors="replace").strip().splitlines()[-1]
    print(f"device: {answer} ({len(data) / 1024 / max(time.time() - start, 0.001):.0f} KB/s)")
    if " error " in answer:
        sys.exit(1)


def png(path, width, height, rgb, scale):
    rows = b""
    for y in range(height):
        row = b"".join(rgb[(y * width + x) * 3:(y * width + x) * 3 + 3] * scale for x in range(width))
        rows += (b"\x00" + row) * scale
    def chunk(kind, body):
        return struct.pack(">I", len(body)) + kind + body + struct.pack(">I", zlib.crc32(kind + body))
    header = struct.pack(">IIBBBBB", width * scale, height * scale, 8, 2, 0, 0, 0)
    with open(path, "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n" + chunk(b"IHDR", header) + chunk(b"IDAT", zlib.compress(rows)) + chunk(b"IEND", b""))


def screenshot(sock, path):
    _, words, rest = binary(sock, "screenshot", "screenshot")
    width, height = int(words[0]), int(words[1])
    pixels = receive(sock, rest, width * height)
    # RGB332, as M5GFX stores an 8-bit sprite: RRRGGGBB.
    rgb = b"".join(bytes(((v >> 5) * 255 // 7, ((v >> 2) & 7) * 255 // 7, (v & 3) * 255 // 3)) for v in pixels)
    path = path or time.strftime("screen-%Y%m%d-%H%M%S.png")
    png(path, width, height, rgb, 2)
    print(f"saved {path} ({width * 2}x{height * 2})")


def interactive(sock):
    sock.setblocking(False)
    while True:
        ready, _, _ = select.select([sock, sys.stdin], [], [])
        if sock in ready:
            data = sock.recv(4096)
            if not data:
                print("\n(connection closed)")
                return
            sys.stdout.write(data.decode(errors="replace"))
            sys.stdout.flush()
        if sys.stdin in ready:
            # Straight from the descriptor: readline() takes every waiting line into Python's own
            # buffer and hands over one, and select() then sees nothing more to read. Piped input
            # written while we were still connecting got stuck until the next line came.
            data = os.read(sys.stdin.fileno(), 4096)
            if not data:
                return
            sock.setblocking(True)
            sock.sendall(data)
            sock.setblocking(False)


def main():
    args = sys.argv[1:]
    host, backlog, token = os.environ.get("RORO_OTA_HOST"), False, None
    while args and args[0].startswith("-"):
        flag = args.pop(0)
        if flag == "-H" and args:
            host = args.pop(0)
        elif flag in ("-t", "--token") and args:
            token = args.pop(0)
        elif flag == "-b":
            backlog = True
        else:
            sys.exit(__doc__)
    if not host:
        sys.exit("No device: pass -H <ip> or set RORO_OTA_HOST\n\n" + __doc__)
    token = find_token(token)

    try:
        sock = socket.create_connection((host, PORT), timeout=10)
    except OSError as e:
        sys.exit(f"device: can't connect to {host}:{PORT} ({e}). Is the console switched on (Settings > Debug Console)?")
    with sock:
        banner = log_in(sock, token)
        show = sys.stdout if backlog or not args else None
        if show:
            show.write(banner.decode(errors="replace"))
        read_until_quiet(sock, 0.5, show)  # the backlog
        if not args:
            return interactive(sock)
        if args == ["crash"]:
            return crash(sock)
        if args[0] == "get" and len(args) >= 2:
            return get(sock, args[1], args[2] if len(args) > 2 else None)
        if args[0] == "put" and len(args) >= 2:
            return put(sock, args[1], args[2] if len(args) > 2 else None)
        if args[0] == "screenshot":
            return screenshot(sock, args[1] if len(args) > 1 else None)
        if args == ["reset"]:  # answered by the console's own task, not the main loop
            sock.sendall(b"reset\n")
            print((read_until(sock, b"restarting now\n", 10) or b"device: no answer").decode().strip().splitlines()[-1])
            return
        if args[0] == "coredump" and args[1:2] != ["erase"]:
            return coredump(sock, args[1] if len(args) > 1 else None)
        run(sock, " ".join(args))


if __name__ == "__main__":
    try:
        main()
    except (ConnectionResetError, BrokenPipeError):
        sys.exit("device: the connection dropped (restarting?)")