reboot-ping-pong/rpp

#!/usr/bin/python3

import sys
import time
from collections import namedtuple

import pexpect

Endpoint = namedtuple("Endpoint", "host, port")
Credentials = namedtuple("Credentials", "username, password")
Target = namedtuple("Target", "credentials, endpoint")
Entity = namedtuple("Entity", "console, ssh")
Machine = namedtuple("Machine", "bmc, host")


class Obmcutil(object):
    BMC_READY = "xyz.openbmc_project.State.BMC.BMCState.Ready"
    BMC_NOT_READY = "xyz.openbmc_project.State.BMC.BMCState.NotReady"

    HOST_OFF = "xyz.openbmc_project.State.Host.HostState.Off"
    HOST_ON = "xyz.openbmc_project.State.Host.HostState.Running"
    HOST_QUIESCED = "xyz.openbmc_project.State.Host.HostState.Quiesced"

    def __init__(self, session, prompt):
        self.session = session
        self.prompt = prompt

    def _clear(self):
        self.session.expect([".+".encode(), pexpect.TIMEOUT], timeout=5)

    def _state(self, cmd, needle):
        self.session.sendline()
        self._clear()
        self.session.sendline("obmcutil -w {}".format(cmd).encode())
        self.session.expect(needle, timeout=None)
        rc = self.session.after.decode()
        return rc

    def hoststate(self):
        return self._state(
            "hoststate",
            "xyz\\.openbmc_project\\.State\\.Host\\.HostState\\."
            + "(Off|Running|Quiesced)",
        )

    def bmcstate(self):
        return self._state(
            "bmcstate",
            "xyz\\.openbmc_project\\.State\\.BMC\\.BMCState\\.(Not)?Ready",
        )

    def poweron(self):
        self.session.sendline("obmcutil -w poweron")
        self.session.expect(self.prompt)

    def chassisoff(self):
        self.session.sendline("obmcutil -w chassisoff")
        self.session.expect(self.prompt)


class PexpectLogger(object):
    def write(self, bstring):
        try:
            sys.stdout.write(bstring.decode())
        except UnicodeDecodeError:
            print("Dropping broken unicode line")

    def flush(self):
        sys.stdout.flush()


class Bmc(object):
    def __init__(self, entity):
        self.getty = "login: ".encode()
        self.shell = "# ".encode()
        self.entity = entity
        fargs = (entity.console.endpoint.host, entity.console.endpoint.port)
        self.session = pexpect.spawn("telnet {} {}".format(*fargs))
        self.session.logfile = PexpectLogger()
        self.obmcutil = Obmcutil(self.session, self.shell)
        self.session.sendline()
        rc = self.session.expect([self.getty, self.shell])
        if rc == 0:
            self.login()

    def login(self):
        self.session.sendline(
            self.entity.console.credentials.username.encode()
        )
        self.session.expect("Password: ".encode())
        self.session.sendline(
            self.entity.console.credentials.password.encode()
        )
        self.session.expect(self.shell)

    def reboot(self):
        self.session.sendline("reboot")
        self.session.expect(
            "Hit any key to stop autoboot:".encode(), timeout=None
        )
        self.session.expect(self.getty, timeout=None)
        self.login()
        state = self.obmcutil.bmcstate()
        while state != self.obmcutil.BMC_READY:
            print(
                "Wanted state '{}', got state '{}'".format(
                    self.obmcutil.BMC_READY, state
                )
            )
            time.sleep(5)
            state = self.obmcutil.bmcstate()

    def chassisoff(self):
        self.obmcutil.chassisoff()

    def poweron(self):
        hs = self.obmcutil.hoststate()
        print("Host state is: {}".format(hs))
        if hs in (self.obmcutil.HOST_ON, self.obmcutil.HOST_QUIESCED):
            self.obmcutil.chassisoff()
        self.obmcutil.poweron()


class Host(object):
    def __init__(self, entity, bmc):
        self.shell = "/? *#".encode()
        self.petitboot = "Petitboot".encode()
        self.session = None
        self.entity = entity
        self.bmc = bmc
        self.connect()

    def connect(self):
        fargs = (
            self.entity.console.endpoint.port,
            self.entity.console.credentials.username,
            self.entity.console.endpoint.host,
        )
        self.session = pexpect.spawn("ssh -p{} {}@{}".format(*fargs))
        self.session.logfile = PexpectLogger()
        self.session.expect("password:".encode())
        self.session.sendline(
            self.entity.console.credentials.password.encode()
        )

    def poweron(self):
        self.bmc.chassisoff()
        self.bmc.poweron()
        self.session.send("\f")
        rc = self.session.expect([self.petitboot, self.shell], timeout=None)
        if rc == 0:
            self.session.sendline()
            self.session.expect(self.shell)

    def reboot(self):
        self.session.send("\f")
        rc = self.session.expect([self.petitboot, self.shell], timeout=None)
        if rc == 0:
            self.session.sendline()
            self.session.expect(self.shell)
        self.session.sendline("reboot".encode())
        self.session.expect(
            "INIT: Waiting for kernel...".encode(), timeout=None
        )
        self.session.expect("Petitboot".encode(), timeout=None)
        self.session.sendline()
        self.session.expect(self.shell)


def rpp(machine):
    bmc = Bmc(machine.bmc)
    host = Host(machine.host, bmc)
    host.poweron()
    while True:
        bmc.reboot()
        host.connect()
        host.reboot()


def main():
    bmccreds = Credentials("root", "0penBmc")
    b = Entity(
        Target(bmccreds, Endpoint("serial.concentrator.somewhere.com", 1234)),
        Target(bmccreds, Endpoint("bmc.somewhere.com", 22)),
    )
    h = Entity(
        Target(bmccreds, Endpoint("bmc.somewhere.com", 2200)),
        Target(
            Credentials("user", "password"), Endpoint("host.somewhere.com", 22)
        ),
    )
    m = Machine(b, h)
    return rpp(m)


if __name__ == "__main__":
    main()