prplos/.gitlab/scripts/testbed-device.py

#!/usr/bin/env python3

import argparse
import glob
import os
import sys
import logging
import json
import time
import traceback

from labgrid import Environment, StepReporter
from labgrid.driver import ExecutionError
from labgrid.consoleloggingreporter import ConsoleLoggingReporter


class OpenWrtConsoleShell:
    def __init__(self, args, shell):
        self.args = args
        self.shell = shell

    def run(self, cmd, timeout=30):
        result = None

        try:
            logging.info(f"executing: {cmd}")
            result = self.shell.run_check(cmd, timeout=timeout)
        except ExecutionError as e:
            stdout = " ".join(e.stdout)
            stderr = " ".join(e.stderr)
            logging.error(f"failed: '{e.msg}' stdout='{stdout}' stderr='{stderr}'")
            return None, (e.msg, e.stdout, e.stderr)

        if len(result) > 0:
            logging.info("result:\n" + "\n".join(result) + "\n" + "-" * 80)

        return result

    def ubus_call(self, cmd):
        result = None

        cmd = f"ubus call {cmd}"
        try:
            logging.info(f"executing: {cmd}")
            result = self.shell.run_check(cmd)
        except ExecutionError as e:
            stdout = " ".join(e.stdout).replace("\t", "")
            stderr = " ".join(e.stderr)
            logging.error(f"failed: {e.msg} stdout='{stdout}' stderr='{stderr}'")
            return None, e.stdout[-1], (e.msg, e.stdout, e.stderr)

        result = json.loads(" ".join(result))
        logging.info(f"result: {result}")
        return result


class UbusTR181:
    def __init__(self, args, shell):
        self.args = args
        self.shell = shell

    def api_call(self, method, path, parameters=None):
        cmd = f"{path} {method}"
        if parameters:
            p = json.dumps({"parameters": parameters})
            cmd += f" '{p}'"

        return self.shell.ubus_call(cmd)

    def get(self, path, parameters=None):
        return self.api_call("_get", path, parameters)

    def set(self, path, parameters):
        return self.api_call("_set", path, parameters)

    def add(self, path, parameters):
        return self.api_call("_add", path, parameters)

    def remove(self, path, parameters=None):
        return self.api_call("_del", path, parameters)

    def call(self, path, method, parameters):
        cmd = f"{path} {method}"
        if parameters:
            p = json.dumps(parameters)
            cmd += f" '{p}'"

        return self.shell.ubus_call(cmd)


class TestbedDevice:
    def __init__(self, args):
        self.args = args
        self.shell = None
        self.env = Environment(config_file=self.args.config)
        ConsoleLoggingReporter.start(args.console_logpath)
        self.target = self.env.get_target(args.target)

    def init_shell(self):
        if self.shell:
            return

        self.shell_driver = self.target.get_driver("ShellDriver")
        self.shell = OpenWrtConsoleShell(self.args, self.shell_driver)

    def _init_wan_vlan(self):
        self.shell.run("ubus -t 60 wait_for WANManager.WAN")
        self.shell.run("ubus-cli WANManager.WAN.2.Intf.1.VlanID=101")
        self.shell.run("ubus-cli WANManager.WAN.2.Intf.2.VlanID=100")
        self.shell.run(
            'ubus call WANManager setWANMode \'{ "WANMode": "Ethernet_DHCP" }\''
        )

    def _init_lan_vlan(self):
        bridge_lan_ports = {
            "mxl,osp-tb341-v2": "Bridging.Bridge.1.Port.4",
            "mxl,osp-tb341": "Bridging.Bridge.1.Port.3",
            "prpl,haze": "Bridging.Bridge.1.Port.4",
            "arcadyan,mozart": "Bridging.Bridge.1.Port.2",
            "cznic,turris-omnia": "Bridging.Bridge.2.Port.5",
            "EASY350 ANYWAN (GRX350) Axepoint Asurada model": "Bridging.Bridge.1.Port.5",
        }
        bridge_vlan_port = bridge_lan_ports.get(
            self.board_name,
            bridge_lan_ports.get(self.model, "Bridging.Bridge.1.Port.2"),
        )

        self.shell.run('ubus-cli Bridging.Bridge.1.Standard="802.1Q-2005"')
        self.shell.run(
            "ubus-cli Bridging.Bridge.1.VLAN.+{Alias='vlan201', Name='vlan201', VLANID=201, Enable=1}"
        )
        self.shell.run(
            f'ubus-cli {bridge_vlan_port}.AcceptableFrameTypes="AdmitOnlyVLANTagged"'
        )
        self.shell.run(f'ubus-cli {bridge_vlan_port}.PVID="201"')
        self.shell.run(f'ubus-cli {bridge_vlan_port}.Type="CustomerVLANPort"')
        self.shell.run(f"ubus-cli {bridge_vlan_port}.Enable=1")

        self.shell.run(
            f"ubus-cli Bridging.Bridge.1.VLANPort+{{Alias='LAN', Name='vlan201', Port='Device.{bridge_vlan_port}', VLAN='Device.Bridging.Bridge.1.VLAN.1', Enable=1}}"
        )

    def init_vlans(self):
        self.init_shell()

        system_info = self.shell.ubus_call("system board")
        if not system_info:
            logging.error("Unable to determine running DUT board!")
            return

        self.model = system_info["model"]
        self.board_name = system_info["board_name"]
        logging.info(f"Running on `{self.board_name}` board, model is `{self.model}`.")

        self.ubus_tr181 = UbusTR181(self.args, self.shell)
        self._init_wan_vlan()
        self._init_lan_vlan()

        logging.info(
            "Let the system apply the new configuration, waiting 15 seconds..."
        )
        time.sleep(15)

        self.shell.run("ubus-cli Ethernet.VLANTermination.?")
        self.shell.run("(cat /proc/vlan101 || cat /proc/net/vlan/vlan101) 2> /dev/null")
        self.shell.run("(cat /proc/vlan201 || cat /proc/net/vlan/vlan201) 2> /dev/null")
        self.shell.run("ip address show vlan101")
        self.shell.run("ip address show vlan201")
        self.shell.run("brctl show")
        self.shell.run("ip route show")

    def console_dump_system_state(self):
        self.init_shell()
        self.shell.ubus_call("system board")
        self.shell.run("getDebugInformation -a -o /dev/console", timeout=360)

    def console_recover_ssh_access(self):
        self.init_shell()
        self.shell.run("iptables -P INPUT ACCEPT")
        self.shell.run("iptables -L INPUT")
        self.shell.run("/etc/init.d/dropbear restart; sleep 5")
        self.shell.run("logread | grep dropbear | tail -10")
        self.shell.run("netstat -nlp | grep :22")
        self.shell.run("ping -c1 192.168.1.1; ping -c1 192.168.1.2")

    def boot_into(self):
        strategy = self.target.get_driver("UBootStrategy")
        dest = self.args.destination
        if dest == "shell":
            strategy.transition("shell")
        if dest == "bootloader":
            strategy.transition("uboot")

    def power(self):
        power = self.target.get_driver("PowerProtocol")
        action = self.args.action
        if action == "on":
            power.on()
        if action == "off":
            power.off()
        if action == "cycle":
            power.cycle()

    def check_network(self):
        host = self.args.remote_host
        shell = self.target.get_driver("ShellDriver")
        shell.wait_for("ubus-cli 'IP.Interface.[Alias==\"lan\"].Status?0'", "Up", 60.0)

        shell.wait_for("ping -c1 {} || true".format(host), ", 0% packet loss", 180.0)


def handle_exception_dump_console(args):
    """Handle exception by printing stack trace and dumping console logs"""
    logging.error("An error occurred during execution")
    logging.error("=" * 80)
    logging.error("Stack trace:")
    traceback.print_exc()
    logging.error("=" * 80)
    logging.error("Console log dump:")

    console_log_pattern = os.path.join(args.console_logpath, "console_*")
    console_log_files = glob.glob(console_log_pattern)
    if console_log_files:
        console_log_file = max(console_log_files, key=os.path.getctime)
        with open(console_log_file, "r") as f:
            logging.error(f.read())
    else:
        logging.error("Console log file not found or not available")
    logging.error("=" * 80)
    sys.exit(1)


def main():
    logging.basicConfig(
        level=logging.INFO, format="%(levelname)7s: %(message)s", stream=sys.stderr
    )

    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default=os.environ.get("TB_CONFIG", ".testbed/labgrid/default.yaml"),
        help="config file (default: %(default)s)",
    )
    parser.add_argument(
        "-t",
        "--target",
        type=str,
        default=os.environ.get("TB_TARGET", None),
        help="target device",
    )
    parser.add_argument(
        "-o",
        "--console-logpath",
        type=str,
        default=os.environ.get("TB_CONSOLE_LOGPATH", os.getcwd()),
        help="path for console logfile (default: %(default)s)",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="count",
        default=int(os.environ.get("TB_VERBOSE", 0)),
        help="enable verbose mode",
    )
    parser.add_argument(
        "-d",
        "--debug",
        action="store_true",
        default=os.environ.get("TB_DEBUG"),
        help="enable debug mode",
    )

    subparsers = parser.add_subparsers(dest="command", title="available subcommands")

    subparser = subparsers.add_parser("power", help="control target power")
    subparser.add_argument(
        "action", choices=["on", "off", "cycle"], help="power on/off/cycle target"
    )
    subparser.set_defaults(func=TestbedDevice.power)

    subparser = subparsers.add_parser("boot_into", help="boot target into console")
    subparser.add_argument(
        "destination",
        choices=["shell", "bootloader"],
        help="boot target either into system shell or bootloader console",
    )
    subparser.set_defaults(func=TestbedDevice.boot_into)

    subparser = subparsers.add_parser(
        "check_network", help="ensure that network is usable"
    )
    subparser.add_argument(
        "-r",
        "--remote-host",
        default="192.168.1.2",
        help="remote host used for ping check (default: %(default)s)",
    )
    subparser.add_argument(
        "-n", "--network", default="lan", help="target network (default: %(default)s)"
    )
    subparser.set_defaults(func=TestbedDevice.check_network)

    subparser = subparsers.add_parser(
        "init_vlans", help="initialize VLAN configuration"
    )
    subparser.set_defaults(func=TestbedDevice.init_vlans)

    subparser = subparsers.add_parser(
        "console_dump_system_state", help="dump system state using serial console"
    )
    subparser.set_defaults(func=TestbedDevice.console_dump_system_state)

    subparser = subparsers.add_parser(
        "console_recover_ssh_access",
        help="try to recover SSH access using serial console",
    )
    subparser.set_defaults(func=TestbedDevice.console_recover_ssh_access)

    args = parser.parse_args()
    if args.verbose >= 1:
        StepReporter.start()

    if args.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    if not args.target:
        print("target device name is mandatory")
        exit(1)

    if not args.command:
        print("command is missing")
        exit(1)

    try:
        device = TestbedDevice(args)
        args.func(device)
    except Exception as e:
        handle_exception_dump_console(args)


if __name__ == "__main__":
    main()