frc-openmesh-config.py

#!/usr/bin/env python3
"""
Utility to change settings on the OpenMesh OM5P-AN and OM5P-AC running FRC OpenWrt.

One can flash the radio firmware using ap51-flash <https://github.com/ap51-flash/ap51-flash>.
This is also what the official utility uses to flash the radio firmware.

You'll need to download the official utility to grab a copy of the firmware.
"""

import enum
import socket
import time

from typing import Optional

MIN_ALLOWED_FIRMWARE = (19, 0, 0)
WRT_DEFAULT_IP = "192.168.1.1"


class Mode(enum.Enum):
    BRIDGE = ("Bridge", "BRIDGE", "B5")
    AP24 = ("2.4GHz Access Point", "AP 2.4GHz", "AP24")
    AP5 = ("5GHz Access Point", "AP 5GHz", "AP5")
    #: AP24_5 seems to have disappeared from the official radio config tool,
    #: but it looks like it still works.
    AP24_5 = ("2.4GHz + 5GHz AP", "AP 2.4+5GHz", "APDUO")
    BRIDGE24 = ("2.4GHz Bridge", "BRIDGE24", "B24")

    def __init__(self, long_str: str, short_str: str, chars: str) -> None:
        self.long_str = long_str
        self.short_str = short_str
        self.chars = chars

    def __str__(self):
        return self.chars

    @classmethod
    def interactively_choose(cls):
        modes = list(cls)
        print('Available modes:')
        for i, mode in enumerate(modes):
            print(i, mode.long_str, sep='. ')
        return modes[int(input('Choose mode: '))]

    def is_ap(self) -> bool:
        return self is not self.BRIDGE  # ???


def get_yn(prompt, default: Optional[bool] = None) -> bool:
    choices = {True: 'Yn', False: 'yN', None: 'yn'}[default]
    prompt = '%s [%s] ' % (prompt, choices)
    while True:
        inp = input(prompt).upper()
        if inp.startswith('Y'):
            return True
        if inp.startswith('N'):
            return False
        if default is not None:
            return default
        print('Invalid choice! Please choose yes or no.')


def main():
    mode = Mode.interactively_choose()
    team = int(input('Team number: '))
    ssid = input('Custom SSID [%d]: ' % team) or team
    wpa_key = input('WPA key (optional): ')
    firewall = get_yn('Firewall?', default=False)
    bw_limit = get_yn('Bandwidth limit?', default=False)
    target = input('Target to reconfigure [%s]: ' % WRT_DEFAULT_IP) or WRT_DEFAULT_IP

    print()
    program(
        mode,
        team,
        ssid,
        wpa_key,
        firewall,
        bw_limit=4000 if bw_limit else None,
        target=target,
    )


def program(
    mode: Mode,
    team: int,
    ssid: str,
    wpa_key: str,
    firewall: bool = False,
    *,
    bw_limit: Optional[int] = 4000,
    dhcp: bool = True,
    chan_24: int = 0,
    chan_5: int = 0,
    event: str = "",
    date: int = 0,
    target: str = WRT_DEFAULT_IP
):
    if dhcp is None:
        dhcp = mode not in (Mode.BRIDGE, Mode.BRIDGE24)

    prog = (
        mode,
        team,
        ssid,
        wpa_key,
        firewall,
        bw_limit,
        dhcp,
        chan_24,
        chan_5,
        event,
        date,
    )
    prog = ",".join(
        ["NY"[x] if isinstance(x, bool) else "" if x is None else str(x) for x in prog]
    ).encode()

    # 2017 /bin/programmer sends 'Model: <?>\nVersion: <version>\n\nEvent: <event>'(?) on connect
    # 2018 /bin/programmer sends '1-0:<date>:<version>\n:<event>' or something.
    # I suspect the extra newline is so the 2017 tool gives a sensible error.

    # The 2017 utility stopped reading after it received a line with "Event:",
    # then split everything on /:|\n/g, reading the version from index 3.
    # This changed in 2018 to stop reading lines after 3 ":" are received,
    # and reads the version from index 2 instead (for obvious reasons).

    with socket.create_connection((target, 8888)) as sock, sock.makefile() as f:
        # I feel like doing this without regex, so I'm making assumptions about the protocol here.
        line = f.readline()
        if not line:
            print('ERROR: programmer disconnected before sending firmware version?')
            return False
        print(line, end='')
        try:
            _, curr_date, firmware_version = line.rstrip().split(':')
        except ValueError:
            print('ERROR: firmware *way* too old, update to the 2018 firmware')
            return False

        firmware_version_split = tuple(map(int, firmware_version.split('.')))

        if firmware_version_split < MIN_ALLOWED_FIRMWARE:
            print('ERROR: firmware too old (%s)' % firmware_version)
            return False

        line = f.readline()
        print(line, end='')
        curr_event = line.strip(':\n')
        print()

        if time.time() < int(curr_date) / 1000:
            print(
                "NOTICE: This radio is programmed for an event in progress (%s)."
                % curr_event
            )
            if not get_yn(
                "Are you sure you wish to erase the field programming?", default=False
            ):
                print('Aborting.')
                return False

        sock.send(prog + b'\n')

        for line in f:
            print(line, end='')

    return True


if __name__ == '__main__':
    main()