Edits history of script submission #23902 for ' Device upgrade (arista_eos)'

• python3

  Ap­pro­ved version

  Edit in Windmill

  Embed

  # 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
  
  
  class ActivateFailed(Exception):
      pass
  
  
  class RollbackFailed(Exception):
      pass
  
  
  def main(
      auth: arista_eos,
      action: Literal["activate", "reload", "commit", "rollback"],
      target_version: str = "",
      image_base_url: str = "",
      vrf: str = "",
      min_uptime_seconds: int = 300,
      delete_old_images: bool = False,
  ):
      """
      Device upgrade
      One upgrade step on one switch: activate (copy EOS-<target>.swi from the image server, verify its sha512, boot system it, no reload), reload (exactly once, guarded), commit (write memory, keep the previous image), rollback (boot the previous image and reload, same guard). Returns { status: "done" } or throws ActivateFailed, RollbackFailed, DeviceUnreachable or AuthFailure as the first token of the error.
      """
      key = "arista_eos_" + re.sub(r"\W", "_", auth["host"])
      state = load_state(key) or {}
      min_uptime = 300 if min_uptime_seconds is None else min_uptime_seconds
  
      if action == "activate":
          if not target_version or not image_base_url:
              raise named(
                  ActivateFailed, "target_version and image_base_url are required"
              )
          image = f"EOS-{target_version}.swi"
          url = f"{image_base_url.rstrip('/')}/{image}"
          steps = []
          previous = state.get("previous_image") or ""
          try:
              try:
                  with urllib.request.urlopen(f"{url}.sha512", timeout=60) as r:
                      expected = sha512(r.read().decode(errors="replace"))
              except urllib.error.HTTPError as e:
                  raise Exception(f"cannot read {url}.sha512 ({e.code})")
              except OSError as e:
                  raise Exception(f"cannot read {url}.sha512 ({getattr(e, 'reason', e)})")
              if not expected:
                  raise Exception(f"no sha512 hash in {url}.sha512")
  
              d, boot, version = eapi(
                  auth,
                  [
                      {"cmd": "dir flash:/", "revision": 2},
                      "show boot-config",
                      "show version",
                  ],
              )
              files = list(d["urls"]["flash:/"]["entries"])
              running = short_version(version["version"])
              booted = boot["softwareImage"].replace("flash:/", "")
              # boot-config is the running image (first run) or the target (an earlier run got past boot system).
              # Anything else is a stale boot-config this step won't guess around.
              if image_version(booted) != running and booted != image:
                  raise Exception(
                      f"boot-config points to {boot['softwareImage']}, neither the running version {running} nor {image}: fix boot-config first"
                  )
              # The rollback target is the image of the running version, recorded once per run before anything changes.
              if not previous and running != target_version:
                  candidates = (
                      [booted]
                      if booted != image
                      else [
                          f
                          for f in files
                          if f.endswith(".swi") and image_version(f) == running
                      ]
                  )
                  if len(candidates) != 1 or candidates[0] not in files:
                      raise Exception(
                          f"cannot tell which image on flash runs {running} (candidates: {', '.join(candidates) or 'none'}): no rollback target"
                      )
                  previous = f"flash:/{candidates[0]}"
                  save_state(key, {**state, "previous_image": previous})
  
              def copy():
                  # The copy runs on the switch and can take minutes for a ~1 GB image.
                  eapi(
                      auth,
                      [*([f"cli vrf {vrf}"] if vrf else []), f"copy {url} flash:{image}"],
                      "json",
                      3600,
                  )
                  steps.append(f"copied {url}")
  
              def verify():
                  [v] = eapi(auth, [f"verify /sha512 flash:{image}"])
                  return sha512((v.get("messages") or [""])[0])
  
              copied = image not in files
              if copied:
                  copy()
              actual = verify()
              if actual != expected and not copied:
                  # Most likely a partial copy left by a killed run: copy it once more.
                  eapi(auth, [f"delete flash:{image}"])
                  copy()
                  actual = verify()
              if actual != expected:
                  # An earlier run may have booted it already: never leave boot-config on a deleted image.
                  restored = booted == image and previous
                  if restored:
                      boot_system(auth, previous.replace("flash:/", ""))
                  eapi(auth, [f"delete flash:{image}"])
                  raise Exception(
                      f"sha512 mismatch for {image} (switch {actual or 'none'}, image server {expected}), deleted it"
                      + (f", boot-config back on {previous}" if restored else "")
                  )
              steps.append("sha512 verified")
              steps.extend(boot_system(auth, image))
          except (DeviceUnreachable, AuthFailure):
              raise
          except Exception as e:
              raise named(ActivateFailed, str(e))
          return {"status": "done", "steps": steps, "previous_image": previous}
  
      if action == "commit":
          _, boot, d = eapi(
              auth,
              ["write memory", "show boot-config", {"cmd": "dir flash:/", "revision": 2}],
          )
          deleted = []
          # Rollback needs the previous image: only images that are neither booted nor the recorded previous one go.
          if delete_old_images and state.get("previous_image"):
              keep = [
                  i.replace("flash:/", "")
                  for i in (boot["softwareImage"], state["previous_image"])
              ]
              for f in list(d["urls"]["flash:/"]["entries"]):
                  if re.match(r"^EOS.*\.swi$", f) and f not in keep:
                      eapi(auth, [f"delete flash:{f}"])
                      deleted.append(f)
          return {"status": "done", "saved": True, "deleted": deleted}
  
      if action == "reload":
          if not target_version:
              raise Exception("target_version is required for the reload guard")
          sw = read_switch(auth, state, min_uptime)
          sent = sw["marker"]("reload")
          if sw.get("unreachable"):
              if sent:
                  return skip(
                      f"reload sent at {sent['sent_at']}; the switch no longer answers: it is reloading"
                  )
              raise named(DeviceUnreachable, sw["unreachable"])
          if sent and (sent["confirmed"] or sent["boot_timestamp"] != sw["boot"]):
              return skip(guard_reason("reload", sent, sw))
          if sw["running"] == target_version:
              return skip(f"already running {target_version}")
          if sw["image"] != f"flash:/EOS-{target_version}.swi":
              raise Exception(
                  f"boot-config points to {sw['image']}, not EOS-{target_version}.swi: run activate first"
              )
          if sw["uptime"] < min_uptime:
              return skip(f"rebooted {round(sw['uptime'])} s ago")
          fire_reload(auth, key, "reload", sw["boot"])
          return {"status": "done", "reloaded": True}
  
      # rollback: boot the previous image (still on flash) and reload, with the same once-only guard.
      try:
          sw = read_switch(auth, state, min_uptime)
          sent = sw["marker"]("rollback")
          if sw.get("unreachable"):
              if sent:
                  return skip(
                      f"rollback sent at {sent['sent_at']}; the switch no longer answers: it is reloading"
                  )
              raise Exception(sw["unreachable"])
          if sent and (sent["confirmed"] or sent["boot_timestamp"] != sw["boot"]):
              return skip(guard_reason("rollback", sent, sw))
          d, boot = eapi(
              auth, [{"cmd": "dir flash:/", "revision": 2}, "show boot-config"]
          )
          files = list(d["urls"]["flash:/"]["entries"])
          current = boot["softwareImage"].replace("flash:/", "")
          # Without a recorded previous image (new run), the only other EOS image on flash is unambiguous.
          others = [f for f in files if re.match(r"^EOS.*\.swi$", f) and f != current]
          previous = (state.get("previous_image") or "").replace("flash:/", "") or (
              others[0] if len(others) == 1 else ""
          )
          if not previous:
              raise Exception(
                  f"previous image unknown (EOS images on flash: {', '.join(others) or 'none'})"
              )
          if previous not in files:
              raise Exception(f"previous image {previous} is no longer on flash")
          reload = sw["marker"]("reload")
          if reload and reload["confirmed"] and reload["boot_timestamp"] == sw["boot"]:
              raise Exception(
                  f"the reload sent at {reload['sent_at']} has not happened yet (same boot as when it was sent): rolling back now would race it, check the switch"
              )
          steps = boot_system(auth, previous)
          if sw["running"] == image_version(previous):
              return skip(
                  f"already running the previous image {previous}; boot-config set back to it"
              )
          # A recent boot that this run's reload caused doesn't block the rollback: the post-check that ran after it
          # can end well inside min_uptime_seconds.
          ours = reload and reload["boot_timestamp"] != sw["boot"]
          if not ours and sw["uptime"] < min_uptime:
              return skip(f"rebooted {round(sw['uptime'])} s ago, not by this run")
          fire_reload(auth, key, "rollback", sw["boot"])
          return {
              "status": "done",
              "reloaded": True,
              "boot_image": f"flash:/{previous}",
              "steps": steps,
          }
      except AuthFailure:
          raise
      except Exception as e:
          raise named(
              RollbackFailed, f"{auth['host']}: {re.sub(r'^[A-Za-z]+: ', '', str(e))}"
          )
  
  
  def boot_system(auth, image):
      _, set_, _, after = eapi(
          auth, ["configure", f"boot system flash:{image}", "end", "show boot-config"]
      )
      if after["softwareImage"] != f"flash:/{image}":
          raise Exception(f"boot-config still points to {after['softwareImage']}")
      return [f"boot system flash:{image}", *(set_.get("warnings") or [])]
  
  
  def skip(reason):
      return {"status": "done", "reloaded": False, "reason": reason}
  
  
  def guard_reason(kind, m, sw):
      if m["boot_timestamp"] == sw["boot"]:
          return f"{kind} sent at {m['sent_at']}; the switch is going down"
      return f"{kind} sent at {m['sent_at']} already rebooted the switch"
  
  
  # A reload is sent at most once per switch and flow run. Each one leaves a marker { sent_at, boot_timestamp,
  # confirmed }: boot_timestamp is the boot it was sent on, so a later boot means it happened, and confirmed means
  # the command reached the switch. An unconfirmed marker on the same boot (the worker died first) is sent again.
  # Outside a flow, a marker is dropped once the reboot it caused is min_uptime_seconds old.
  def read_switch(auth, state, min_uptime):
      try:
          version, boot = eapi(auth, ["show version", "show boot-config"])
      except DeviceUnreachable as e:
          return {"unreachable": str(e), "marker": lambda k: state.get(k)}
      sw = {
          "running": short_version(version["version"]),
          "image": boot["softwareImage"],
          "boot": version["bootupTimestamp"],
          "uptime": version["uptime"],
      }
      in_flow = bool(os.environ.get("WM_FLOW_JOB_ID"))
  
      def marker(k):
          m = state.get(k)
          return (
              m
              if m
              and (
                  in_flow
                  or m["boot_timestamp"] == sw["boot"]
                  or sw["uptime"] < min_uptime
              )
              else None
          )
  
      return {**sw, "marker": marker}
  
  
  # Fire and forget: the switch drops the connection while going down, which is not an error. A connection that
  # never opened, or a refusal from the switch, means it did not reload: the marker goes and the error is raised.
  def fire_reload(auth, key, kind, boot_ts):
      def mark(m):
          rest = {k: v for k, v in (load_state(key) or {}).items() if k != kind}
          save_state(key, {**rest, kind: m} if m else rest)
  
      m = {
          "sent_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
          "boot_timestamp": boot_ts,
          "confirmed": False,
      }
      mark(m)
      eos = node(auth, 15)
      try:
          # Opened before sending, so a failure here means the command never left.
          eos.connection.transport.connect()
      except OSError as e:
          mark(None)
          raise named(DeviceUnreachable, f"{auth['host']}: {e}")
      try:
          eos.run_commands(["reload now"])
      except pyeapi.eapilib.ConnectionError as e:
          # Sent, then the connection dropped or timed out as the switch went down. A 401 is a refusal.
          err = connection_error(auth, eos, e)
          if isinstance(err, AuthFailure):
              mark(None)
              raise err
      except Exception:
          mark(None)
          raise
      mark({**m, "confirmed": True})
  
  
  # Shared with Device check: the flow run's user state, keyed by switch, holds the pre-check snapshot, the image
  # activate replaced (rollback target) and the reload markers. 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
  
  
  # "EOS64-4.34.9M.swi" -> "4.34.9M"
  def image_version(file):
      m = re.search(r"-(\d+\.\d+\.\d+(\.\d+)?[A-Z]*)\.swi$", file)
      return m.group(1) if m else None
  
  
  # The 128-hex digest in a sha512sum line, a BSD "SHA512 (f) = <hash>" line or verify /sha512 output.
  def sha512(text):
      m = re.search(r"\b[0-9a-fA-F]{128}\b", text)
      return m.group(0).lower() if m else ""
  
  
  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,
      )
  

  Submitted by hugo989 3 hours ago