0

Device upgrade

by
Published today

One upgrade step on one switch: activate, reload (once-only guard), commit or rollback. Returns { status: done } or throws ActivateFailed, RollbackFailed, DeviceUnreachable or AuthFailure.

Script arista_eos Verified

The script

Submitted by hugo989 Python3
Verified 3 hours ago
1
# py312
2
# requirements:
3
# pyeapi==1.0.4
4
# wmill==1.851.0
5

6
import logging
7
import os
8
import re
9
import ssl
10
import time
11
import urllib.error
12
import urllib.request
13
from typing import Literal, TypedDict
14
from urllib.parse import urlsplit
15

16
import pyeapi
17
import wmill
18

19

20
class arista_eos(TypedDict):
21
    host: str
22
    username: str
23
    password: str
24
    port: int
25
    verify_tls: bool
26

27

28
class ActivateFailed(Exception):
29
    pass
30

31

32
class RollbackFailed(Exception):
33
    pass
34

35

36
def main(
37
    auth: arista_eos,
38
    action: Literal["activate", "reload", "commit", "rollback"],
39
    target_version: str = "",
40
    image_base_url: str = "",
41
    vrf: str = "",
42
    min_uptime_seconds: int = 300,
43
    delete_old_images: bool = False,
44
):
45
    """
46
    Device upgrade
47
    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.
48
    """
49
    key = "arista_eos_" + re.sub(r"\W", "_", auth["host"])
50
    state = load_state(key) or {}
51
    min_uptime = 300 if min_uptime_seconds is None else min_uptime_seconds
52

53
    if action == "activate":
54
        if not target_version or not image_base_url:
55
            raise named(
56
                ActivateFailed, "target_version and image_base_url are required"
57
            )
58
        image = f"EOS-{target_version}.swi"
59
        url = f"{image_base_url.rstrip('/')}/{image}"
60
        steps = []
61
        previous = state.get("previous_image") or ""
62
        try:
63
            try:
64
                with urllib.request.urlopen(f"{url}.sha512", timeout=60) as r:
65
                    expected = sha512(r.read().decode(errors="replace"))
66
            except urllib.error.HTTPError as e:
67
                raise Exception(f"cannot read {url}.sha512 ({e.code})")
68
            except OSError as e:
69
                raise Exception(f"cannot read {url}.sha512 ({getattr(e, 'reason', e)})")
70
            if not expected:
71
                raise Exception(f"no sha512 hash in {url}.sha512")
72

73
            d, boot, version = eapi(
74
                auth,
75
                [
76
                    {"cmd": "dir flash:/", "revision": 2},
77
                    "show boot-config",
78
                    "show version",
79
                ],
80
            )
81
            files = list(d["urls"]["flash:/"]["entries"])
82
            running = short_version(version["version"])
83
            booted = boot["softwareImage"].replace("flash:/", "")
84
            # boot-config is the running image (first run) or the target (an earlier run got past boot system).
85
            # Anything else is a stale boot-config this step won't guess around.
86
            if image_version(booted) != running and booted != image:
87
                raise Exception(
88
                    f"boot-config points to {boot['softwareImage']}, neither the running version {running} nor {image}: fix boot-config first"
89
                )
90
            # The rollback target is the image of the running version, recorded once per run before anything changes.
91
            if not previous and running != target_version:
92
                candidates = (
93
                    [booted]
94
                    if booted != image
95
                    else [
96
                        f
97
                        for f in files
98
                        if f.endswith(".swi") and image_version(f) == running
99
                    ]
100
                )
101
                if len(candidates) != 1 or candidates[0] not in files:
102
                    raise Exception(
103
                        f"cannot tell which image on flash runs {running} (candidates: {', '.join(candidates) or 'none'}): no rollback target"
104
                    )
105
                previous = f"flash:/{candidates[0]}"
106
                save_state(key, {**state, "previous_image": previous})
107

108
            def copy():
109
                # The copy runs on the switch and can take minutes for a ~1 GB image.
110
                eapi(
111
                    auth,
112
                    [*([f"cli vrf {vrf}"] if vrf else []), f"copy {url} flash:{image}"],
113
                    "json",
114
                    3600,
115
                )
116
                steps.append(f"copied {url}")
117

118
            def verify():
119
                [v] = eapi(auth, [f"verify /sha512 flash:{image}"])
120
                return sha512((v.get("messages") or [""])[0])
121

122
            copied = image not in files
123
            if copied:
124
                copy()
125
            actual = verify()
126
            if actual != expected and not copied:
127
                # Most likely a partial copy left by a killed run: copy it once more.
128
                eapi(auth, [f"delete flash:{image}"])
129
                copy()
130
                actual = verify()
131
            if actual != expected:
132
                # An earlier run may have booted it already: never leave boot-config on a deleted image.
133
                restored = booted == image and previous
134
                if restored:
135
                    boot_system(auth, previous.replace("flash:/", ""))
136
                eapi(auth, [f"delete flash:{image}"])
137
                raise Exception(
138
                    f"sha512 mismatch for {image} (switch {actual or 'none'}, image server {expected}), deleted it"
139
                    + (f", boot-config back on {previous}" if restored else "")
140
                )
141
            steps.append("sha512 verified")
142
            steps.extend(boot_system(auth, image))
143
        except (DeviceUnreachable, AuthFailure):
144
            raise
145
        except Exception as e:
146
            raise named(ActivateFailed, str(e))
147
        return {"status": "done", "steps": steps, "previous_image": previous}
148

149
    if action == "commit":
150
        _, boot, d = eapi(
151
            auth,
152
            ["write memory", "show boot-config", {"cmd": "dir flash:/", "revision": 2}],
153
        )
154
        deleted = []
155
        # Rollback needs the previous image: only images that are neither booted nor the recorded previous one go.
156
        if delete_old_images and state.get("previous_image"):
157
            keep = [
158
                i.replace("flash:/", "")
159
                for i in (boot["softwareImage"], state["previous_image"])
160
            ]
161
            for f in list(d["urls"]["flash:/"]["entries"]):
162
                if re.match(r"^EOS.*\.swi$", f) and f not in keep:
163
                    eapi(auth, [f"delete flash:{f}"])
164
                    deleted.append(f)
165
        return {"status": "done", "saved": True, "deleted": deleted}
166

167
    if action == "reload":
168
        if not target_version:
169
            raise Exception("target_version is required for the reload guard")
170
        sw = read_switch(auth, state, min_uptime)
171
        sent = sw["marker"]("reload")
172
        if sw.get("unreachable"):
173
            if sent:
174
                return skip(
175
                    f"reload sent at {sent['sent_at']}; the switch no longer answers: it is reloading"
176
                )
177
            raise named(DeviceUnreachable, sw["unreachable"])
178
        if sent and (sent["confirmed"] or sent["boot_timestamp"] != sw["boot"]):
179
            return skip(guard_reason("reload", sent, sw))
180
        if sw["running"] == target_version:
181
            return skip(f"already running {target_version}")
182
        if sw["image"] != f"flash:/EOS-{target_version}.swi":
183
            raise Exception(
184
                f"boot-config points to {sw['image']}, not EOS-{target_version}.swi: run activate first"
185
            )
186
        if sw["uptime"] < min_uptime:
187
            return skip(f"rebooted {round(sw['uptime'])} s ago")
188
        fire_reload(auth, key, "reload", sw["boot"])
189
        return {"status": "done", "reloaded": True}
190

191
    # rollback: boot the previous image (still on flash) and reload, with the same once-only guard.
192
    try:
193
        sw = read_switch(auth, state, min_uptime)
194
        sent = sw["marker"]("rollback")
195
        if sw.get("unreachable"):
196
            if sent:
197
                return skip(
198
                    f"rollback sent at {sent['sent_at']}; the switch no longer answers: it is reloading"
199
                )
200
            raise Exception(sw["unreachable"])
201
        if sent and (sent["confirmed"] or sent["boot_timestamp"] != sw["boot"]):
202
            return skip(guard_reason("rollback", sent, sw))
203
        d, boot = eapi(
204
            auth, [{"cmd": "dir flash:/", "revision": 2}, "show boot-config"]
205
        )
206
        files = list(d["urls"]["flash:/"]["entries"])
207
        current = boot["softwareImage"].replace("flash:/", "")
208
        # Without a recorded previous image (new run), the only other EOS image on flash is unambiguous.
209
        others = [f for f in files if re.match(r"^EOS.*\.swi$", f) and f != current]
210
        previous = (state.get("previous_image") or "").replace("flash:/", "") or (
211
            others[0] if len(others) == 1 else ""
212
        )
213
        if not previous:
214
            raise Exception(
215
                f"previous image unknown (EOS images on flash: {', '.join(others) or 'none'})"
216
            )
217
        if previous not in files:
218
            raise Exception(f"previous image {previous} is no longer on flash")
219
        reload = sw["marker"]("reload")
220
        if reload and reload["confirmed"] and reload["boot_timestamp"] == sw["boot"]:
221
            raise Exception(
222
                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"
223
            )
224
        steps = boot_system(auth, previous)
225
        if sw["running"] == image_version(previous):
226
            return skip(
227
                f"already running the previous image {previous}; boot-config set back to it"
228
            )
229
        # A recent boot that this run's reload caused doesn't block the rollback: the post-check that ran after it
230
        # can end well inside min_uptime_seconds.
231
        ours = reload and reload["boot_timestamp"] != sw["boot"]
232
        if not ours and sw["uptime"] < min_uptime:
233
            return skip(f"rebooted {round(sw['uptime'])} s ago, not by this run")
234
        fire_reload(auth, key, "rollback", sw["boot"])
235
        return {
236
            "status": "done",
237
            "reloaded": True,
238
            "boot_image": f"flash:/{previous}",
239
            "steps": steps,
240
        }
241
    except AuthFailure:
242
        raise
243
    except Exception as e:
244
        raise named(
245
            RollbackFailed, f"{auth['host']}: {re.sub(r'^[A-Za-z]+: ', '', str(e))}"
246
        )
247

248

249
def boot_system(auth, image):
250
    _, set_, _, after = eapi(
251
        auth, ["configure", f"boot system flash:{image}", "end", "show boot-config"]
252
    )
253
    if after["softwareImage"] != f"flash:/{image}":
254
        raise Exception(f"boot-config still points to {after['softwareImage']}")
255
    return [f"boot system flash:{image}", *(set_.get("warnings") or [])]
256

257

258
def skip(reason):
259
    return {"status": "done", "reloaded": False, "reason": reason}
260

261

262
def guard_reason(kind, m, sw):
263
    if m["boot_timestamp"] == sw["boot"]:
264
        return f"{kind} sent at {m['sent_at']}; the switch is going down"
265
    return f"{kind} sent at {m['sent_at']} already rebooted the switch"
266

267

268
# A reload is sent at most once per switch and flow run. Each one leaves a marker { sent_at, boot_timestamp,
269
# confirmed }: boot_timestamp is the boot it was sent on, so a later boot means it happened, and confirmed means
270
# the command reached the switch. An unconfirmed marker on the same boot (the worker died first) is sent again.
271
# Outside a flow, a marker is dropped once the reboot it caused is min_uptime_seconds old.
272
def read_switch(auth, state, min_uptime):
273
    try:
274
        version, boot = eapi(auth, ["show version", "show boot-config"])
275
    except DeviceUnreachable as e:
276
        return {"unreachable": str(e), "marker": lambda k: state.get(k)}
277
    sw = {
278
        "running": short_version(version["version"]),
279
        "image": boot["softwareImage"],
280
        "boot": version["bootupTimestamp"],
281
        "uptime": version["uptime"],
282
    }
283
    in_flow = bool(os.environ.get("WM_FLOW_JOB_ID"))
284

285
    def marker(k):
286
        m = state.get(k)
287
        return (
288
            m
289
            if m
290
            and (
291
                in_flow
292
                or m["boot_timestamp"] == sw["boot"]
293
                or sw["uptime"] < min_uptime
294
            )
295
            else None
296
        )
297

298
    return {**sw, "marker": marker}
299

300

301
# Fire and forget: the switch drops the connection while going down, which is not an error. A connection that
302
# never opened, or a refusal from the switch, means it did not reload: the marker goes and the error is raised.
303
def fire_reload(auth, key, kind, boot_ts):
304
    def mark(m):
305
        rest = {k: v for k, v in (load_state(key) or {}).items() if k != kind}
306
        save_state(key, {**rest, kind: m} if m else rest)
307

308
    m = {
309
        "sent_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
310
        "boot_timestamp": boot_ts,
311
        "confirmed": False,
312
    }
313
    mark(m)
314
    eos = node(auth, 15)
315
    try:
316
        # Opened before sending, so a failure here means the command never left.
317
        eos.connection.transport.connect()
318
    except OSError as e:
319
        mark(None)
320
        raise named(DeviceUnreachable, f"{auth['host']}: {e}")
321
    try:
322
        eos.run_commands(["reload now"])
323
    except pyeapi.eapilib.ConnectionError as e:
324
        # Sent, then the connection dropped or timed out as the switch went down. A 401 is a refusal.
325
        err = connection_error(auth, eos, e)
326
        if isinstance(err, AuthFailure):
327
            mark(None)
328
            raise err
329
    except Exception:
330
        mark(None)
331
        raise
332
    mark({**m, "confirmed": True})
333

334

335
# Shared with Device check: the flow run's user state, keyed by switch, holds the pre-check snapshot, the image
336
# activate replaced (rollback target) and the reload markers. Outside a flow, it falls back to this script's own state.
337
# The endpoint is called through the client so an unreachable state raises (wmill.get_flow_user_state returns None).
338
# ponytail: the fallback is one map per script, so concurrent standalone runs can race; flows don't.
339
def load_state(key):
340
    if os.environ.get("WM_FLOW_JOB_ID"):
341
        c = wmill.Windmill()
342
        return c.get(
343
            f"/w/{c.workspace}/jobs/flow/user_states/{c.get_root_job_id()}/{key}"
344
        ).json()
345
    return (wmill.get_state() or {}).get(key)
346

347

348
def save_state(key, value):
349
    if os.environ.get("WM_FLOW_JOB_ID"):
350
        c = wmill.Windmill()
351
        c.post(
352
            f"/w/{c.workspace}/jobs/flow/user_states/{c.get_root_job_id()}/{key}",
353
            json=value,
354
        )
355
    else:
356
        wmill.set_state({**(wmill.get_state() or {}), key: value})
357

358

359
# "4.34.9M-51132493.4349M (engineering build)" on cEOS, "4.34.9M" on hardware.
360
def short_version(v):
361
    m = re.match(r"\d+\.\d+\.\d+(\.\d+)?[A-Z]*", v)
362
    return m.group(0) if m else v
363

364

365
# "EOS64-4.34.9M.swi" -> "4.34.9M"
366
def image_version(file):
367
    m = re.search(r"-(\d+\.\d+\.\d+(\.\d+)?[A-Z]*)\.swi$", file)
368
    return m.group(1) if m else None
369

370

371
# The 128-hex digest in a sha512sum line, a BSD "SHA512 (f) = <hash>" line or verify /sha512 output.
372
def sha512(text):
373
    m = re.search(r"\b[0-9a-fA-F]{128}\b", text)
374
    return m.group(0).lower() if m else ""
375

376

377
class DeviceUnreachable(Exception):
378
    pass
379

380

381
class AuthFailure(Exception):
382
    pass
383

384

385
def named(cls, message: str):
386
    return cls(f"{cls.__name__}: {message}")
387

388

389
# pyeapi logs a traceback for every socket error, and the post-check polls a switch that is rebooting.
390
logging.getLogger("pyeapi").addHandler(logging.NullHandler())
391

392

393
def eapi(auth: arista_eos, cmds: list, fmt: str = "json", timeout: float = 60):
394
    eos = node(auth, timeout)
395
    try:
396
        return eos.run_commands(cmds, fmt)
397
    except pyeapi.eapilib.ConnectionError as e:
398
        raise connection_error(auth, eos, e)
399

400

401
# pyeapi 1.0.4 raises ConnectionError for a socket error (kept in socket_error), an answer that isn't JSON (eAPI
402
# still starting, 5xx) and a 401, the only one with neither.
403
def connection_error(auth, eos, e):
404
    if eos.connection.socket_error is None and e.message != "unable to connect to eAPI":
405
        return named(
406
            AuthFailure, f"{auth['host']}: eAPI refused the credential ({e.message})"
407
        )
408
    return named(DeviceUnreachable, f"{auth['host']}: {e.message}")
409

410

411
# host is an address or a URL (https://10.0.0.11:8443). pyeapi skips certificate checks unless it is given a context,
412
# so one is always passed: verifying, or not when the resource unticks verify_tls (EOS ships a self-signed cert).
413
def node(auth: arista_eos, timeout: float = 60):
414
    url = urlsplit(auth["host"] if "://" in auth["host"] else f"https://{auth['host']}")
415
    context = ssl.create_default_context()
416
    if auth.get("verify_tls") is False:
417
        context.check_hostname = False
418
        context.verify_mode = ssl.CERT_NONE
419
    return pyeapi.connect(
420
        transport=url.scheme,
421
        host=url.hostname,
422
        port=url.port or auth.get("port") or None,
423
        username=auth["username"],
424
        password=auth["password"],
425
        context=context,
426
        timeout=timeout,
427
        return_node=True,
428
    )
429