Device check

# py312

# requirements:

# pyeapi==1.0.4

# wmill==1.851.0

import logging

import os

import re

import ssl

import time

import urllib.error

import urllib.request

from typing import Literal, TypedDict

from urllib.parse import urlsplit

import pyeapi

import wmill

class arista_eos(TypedDict):

    host: str

    username: str

    password: str

    port: int

    verify_tls: bool

def main(

    auth: arista_eos,

    phase: Literal["pre", "post"],

    target_version: str = "",

    image_base_url: str = "",

    required_free_mb: int = 1024,

    return_timeout_seconds: int = 900,

    stabilisation_seconds: int = 120,

    poll_interval_seconds: int = 15,

):

    """

    Device check

    Pre- or post-upgrade check of one switch: returns { running_version, status ok|ko, checks_ok, checks_ko[{check, message}] }. Pre: flash space for the target image, MLAG redundancy, and a snapshot of interfaces/LLDP/MLAG kept in the flow run's state. Post: waits for the switch to answer and for the MLAG reload-delay, then compares with the snapshot over a bounded stabilisation window. Throws DeviceUnreachable or AuthFailure. Does not compare the version with the target: the workflow does.

    """

    poll = poll_interval_seconds or 15

    durations = {}

    key = "arista_eos_" + re.sub(r"\W", "_", auth["host"])

    state = load_state(key) or {}

    if phase == "post":

        # Wait for the switch to answer on a boot later than the one each reload of this run was sent on: `reload now`

        # can answer before EOS goes down, and the old boot must not pass for the new one. AuthFailure stops at once.

        markers = [state[k] for k in ("reload", "rollback") if state.get(k)]

        timeout = return_timeout_seconds or 900

        start = time.time()

        while True:

            try:

                [v] = eapi(auth, ["show version"])

                m = next(

                    (m for m in markers if m["boot_timestamp"] == v["bootupTimestamp"]),

                    None,

                )

                if not m:

                    break

                pending = (

                    f"still on the boot the reload sent at {m['sent_at']} was sent on"

                )

            except DeviceUnreachable as e:

                pending = str(e)

            if time.time() - start > timeout:

                raise named(

                    DeviceUnreachable,

                    f"{auth['host']}: not back after {timeout} s ({pending})",

                )

            time.sleep(poll)

        durations["return_timeout_seconds"] = timeout

        durations["returned_after_seconds"] = round(time.time() - start)

    version, mlag = eapi(auth, ["show version", "show mlag"])

    running_version = short_version(version["version"])

    durations["reload_delay_seconds"] = max(

        mlag.get("reloadDelay") or 0, mlag.get("reloadDelayNonMlag") or 0

    )

    ok, ko = [], []

    def record(check, problems):

        if problems:

            ko.append({"check": check, "message": "; ".join(problems)})

        else:

            ok.append(check)

    if phase == "pre":

        record(

            "flash",

            flash_problems(auth, target_version, image_base_url, required_free_mb),

        )

        [sanity] = eapi(auth, ["show mlag config-sanity"])

        record("redundancy", redundancy_problems(mlag, sanity))

        save_state(

            key,

            {

                **state,

                "snapshot": snapshot(auth),

                "snapshot_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),

            },

        )

        ok.append("interface")

    elif not state.get("snapshot"):

        ko.append(

            {

                "check": "interface",

                "message": f"No pre-check snapshot for {auth['host']} in this flow run: run the pre-check first, in the same run",

            }

        )

    else:

        # Ports stay errdisabled until the MLAG reload-delay has elapsed since boot: wait it out.

        wait = durations["reload_delay_seconds"] - version["uptime"]

        if wait > 0:

            time.sleep(wait)

        durations["reload_delay_waited_seconds"] = max(0, round(wait))

        stab = 120 if stabilisation_seconds is None else stabilisation_seconds

        start = time.time()

        problems = diff(state["snapshot"], snapshot(auth))

        while problems and time.time() - start < stab:

            time.sleep(poll)

            problems = diff(state["snapshot"], snapshot(auth))

        durations["stabilisation_seconds"] = stab

        durations["stabilised_after_seconds"] = round(time.time() - start)

        record("interface", problems)

    return {

        "running_version": running_version,

        "status": "ko" if ko else "ok",

        "checks_ok": ok,

        "checks_ko": ko,

        "durations": durations,

    }

def flash_problems(auth, target, base, required_mb):

    [d] = eapi(auth, [{"cmd": "dir flash:/", "revision": 2}])

    flash = d["urls"]["flash:/"]

    image = f"EOS-{target}.swi" if target else ""

    if image and image in flash["entries"]:

        return []

    required = (required_mb or 1024) * 1e6

    if image and base:

        try:

            with urllib.request.urlopen(

                urllib.request.Request(f"{base.rstrip('/')}/{image}", method="HEAD"),

                timeout=30,

            ) as head:

                required = int(head.headers.get("content-length") or 0) or required

        except urllib.error.HTTPError as e:

            return [f"Target image {image} not found on the image server ({e.code})"]

        except OSError as e:

            return [

                f"Image server unreachable from the worker: {getattr(e, 'reason', e)}"

            ]

    free = flash["filesystem"]["freeBytes"]

    if free < required:

        return [

            f"Insufficient flash space ({mb(free)} free, {mb(required)} required{f' for {image}' if image else ''})"

        ]

    return []

def redundancy_problems(mlag, sanity):

    if mlag.get("state") == "disabled":

        return [

            "MLAG is not configured: no peer to carry the traffic during the reload"

        ]

    p = []

    if mlag.get("state") != "active":

        p.append(f"MLAG state {mlag.get('state')} (expected active)")

    if mlag.get("negStatus") != "connected":

        p.append(f"MLAG negotiation {mlag.get('negStatus')} (expected connected)")

    if mlag.get("peerLinkStatus") != "up":

        p.append(f"peer-link {mlag.get('peerLink')} {mlag.get('peerLinkStatus')}")

    if mlag.get("localIntfStatus") != "up":

        p.append(

            f"MLAG local interface {mlag.get('localInterface')} {mlag.get('localIntfStatus')}"

        )

    inconsistent = [

        *(sanity.get("globalConfiguration") or {}),

        *(sanity.get("interfaceConfiguration") or {}),

    ]

    if inconsistent:

        p.append(f"config-sanity inconsistent: {', '.join(inconsistent)}")

    partial = (mlag.get("mlagPorts") or {}).get("Active-partial")

    if partial:

        p.append(f"{partial} MLAG port(s) active-partial")

    if mlag.get("portsErrdisabled"):

        p.append("MLAG ports errdisabled")

    return p

def snapshot(auth):

    status, lldp, mlag, mlag_intf = eapi(

        auth,

        [

            "show interfaces status",

            "show lldp neighbors",

            "show mlag",

            "show mlag interfaces",

        ],

    )

    return {

        "interfaces_up": sorted(

            n

            for n, i in status["interfaceStatuses"].items()

            if i.get("linkStatus") == "connected"

        ),

        "lldp": [

            {

                "port": n["port"],

                "neighbor": f"{n['neighborDevice']} {n['neighborPort']}",

            }

            for n in lldp["lldpNeighbors"]

        ],

        "mlag": {

            "state": mlag.get("state"),

            "negStatus": mlag.get("negStatus"),

            "peerLinkStatus": mlag.get("peerLinkStatus"),

            "ports": {

                f"{i['localInterface']} (mlag {id})": i["status"]

                for id, i in (mlag_intf.get("interfaces") or {}).items()

            },

        },

    }

def diff(before, after):

    p = []

    for i in before["interfaces_up"]:

        if i not in after["interfaces_up"]:

            p.append(f"{i} down after the upgrade (up before)")

    for n in before["lldp"]:

        if not any(

            a["port"] == n["port"] and a["neighbor"] == n["neighbor"]

            for a in after["lldp"]

        ):

            p.append(f"LLDP neighbor {n['neighbor']} no longer seen on {n['port']}")

    for k in ("state", "negStatus", "peerLinkStatus"):

        if before["mlag"][k] != after["mlag"][k]:

            p.append(f"MLAG {k} {after['mlag'][k]} (was {before['mlag'][k]})")

    for port, s in before["mlag"]["ports"].items():

        if after["mlag"]["ports"].get(port) != s:

            p.append(f"{port} {after['mlag']['ports'].get(port, 'missing')} (was {s})")

    return p

# The snapshot lives in the root flow run's user state, keyed by switch: pre- and post-check are different steps

# (often in nested subflows) of the same run. Outside a flow, it falls back to this script's own state.

# The endpoint is called through the client so an unreachable state raises (wmill.get_flow_user_state returns None).

# ponytail: the fallback is one map per script, so concurrent standalone runs can race; flows don't.

def load_state(key):

    if os.environ.get("WM_FLOW_JOB_ID"):

        c = wmill.Windmill()

        return c.get(

            f"/w/{c.workspace}/jobs/flow/user_states/{c.get_root_job_id()}/{key}"

        ).json()

    return (wmill.get_state() or {}).get(key)

def save_state(key, value):

    if os.environ.get("WM_FLOW_JOB_ID"):

        c = wmill.Windmill()

        c.post(

            f"/w/{c.workspace}/jobs/flow/user_states/{c.get_root_job_id()}/{key}",

            json=value,

        )

    else:

        wmill.set_state({**(wmill.get_state() or {}), key: value})

# "4.34.9M-51132493.4349M (engineering build)" on cEOS, "4.34.9M" on hardware.

def short_version(v):

    m = re.match(r"\d+\.\d+\.\d+(\.\d+)?[A-Z]*", v)

    return m.group(0) if m else v

def mb(n):

    return f"{n / 1e9:.1f} GB" if n >= 1e9 else f"{round(n / 1e6)} MB"

class DeviceUnreachable(Exception):

    pass

class AuthFailure(Exception):

    pass

def named(cls, message: str):

    return cls(f"{cls.__name__}: {message}")

# pyeapi logs a traceback for every socket error, and the post-check polls a switch that is rebooting.

logging.getLogger("pyeapi").addHandler(logging.NullHandler())

def eapi(auth: arista_eos, cmds: list, fmt: str = "json", timeout: float = 60):

    eos = node(auth, timeout)

    try:

        return eos.run_commands(cmds, fmt)

    except pyeapi.eapilib.ConnectionError as e:

        raise connection_error(auth, eos, e)

# pyeapi 1.0.4 raises ConnectionError for a socket error (kept in socket_error), an answer that isn't JSON (eAPI

# still starting, 5xx) and a 401, the only one with neither.

def connection_error(auth, eos, e):

    if eos.connection.socket_error is None and e.message != "unable to connect to eAPI":

        return named(

            AuthFailure, f"{auth['host']}: eAPI refused the credential ({e.message})"

        )

    return named(DeviceUnreachable, f"{auth['host']}: {e.message}")

# host is an address or a URL (https://10.0.0.11:8443). pyeapi skips certificate checks unless it is given a context,

# so one is always passed: verifying, or not when the resource unticks verify_tls (EOS ships a self-signed cert).

def node(auth: arista_eos, timeout: float = 60):

    url = urlsplit(auth["host"] if "://" in auth["host"] else f"https://{auth['host']}")

    context = ssl.create_default_context()

    if auth.get("verify_tls") is False:

        context.check_hostname = False

        context.verify_mode = ssl.CERT_NONE

    return pyeapi.connect(

        transport=url.scheme,

        host=url.hostname,

        port=url.port or auth.get("port") or None,

        username=auth["username"],

        password=auth["password"],

        context=context,

        timeout=timeout,

        return_node=True,

    )