tests/topotests/munet/cli.py

   1 # -*- coding: utf-8 eval: (blacken-mode 1) -*-
   2 # SPDX-License-Identifier: GPL-2.0-or-later
   3 #
   4 # July 24 2021, Christian Hopps <chopps@labn.net>
   5 #
   6 # Copyright 2021, LabN Consulting, L.L.C.
   7 #
   8 """A module that implements a CLI."""
   9 import argparse
  10 import asyncio
  11 import functools
  12 import logging
  13 import multiprocessing
  14 import os
  15 import pty
  16 import re
  17 import readline
  18 import select
  19 import shlex
  20 import socket
  21 import subprocess
  22 import sys
  23 import tempfile
  24 import termios
  25 import tty
  26
  27
  28 try:
  29     from . import linux
  30     from .config import list_to_dict_with_key
  31 except ImportError:
  32     # We cannot use relative imports and still run this module directly as a script, and
  33     # there are some use cases where we want to run this file as a script.
  34     sys.path.append(os.path.dirname(os.path.realpath(__file__)))
  35     import linux
  36
  37     from config import list_to_dict_with_key
  38
  39
  40 ENDMARKER = b"\x00END\x00"
  41
  42 logger = logging.getLogger(__name__)
  43
  44
  45 def lineiter(sock):
  46     s = ""
  47     while True:
  48         sb = sock.recv(256)
  49         if not sb:
  50             return
  51
  52         s += sb.decode("utf-8")
  53         i = s.find("\n")
  54         if i != -1:
  55             yield s[:i]
  56             s = s[i + 1 :]
  57
  58
  59 # Would be nice to convert to async, but really not needed as used
  60 def spawn(unet, host, cmd, iow, ns_only):
  61     if sys.stdin.isatty():
  62         old_tty = termios.tcgetattr(sys.stdin)
  63         tty.setraw(sys.stdin.fileno())
  64
  65     try:
  66         master_fd, slave_fd = pty.openpty()
  67
  68         ns = unet.hosts[host] if host and host != unet else unet
  69         popenf = ns.popen_nsonly if ns_only else ns.popen
  70
  71         # use os.setsid() make it run in a new process group, or bash job
  72         # control will not be enabled
  73         p = popenf(
  74             cmd,
  75             # _common_prologue, later in call chain, only does this for use_pty == False
  76             preexec_fn=os.setsid,
  77             stdin=slave_fd,
  78             stdout=slave_fd,
  79             stderr=slave_fd,
  80             universal_newlines=True,
  81             use_pty=True,
  82             # XXX this is actually implementing "run on host" for real
  83             # skip_pre_cmd=ns_only,
  84         )
  85         iow.write("\r")
  86         iow.flush()
  87
  88         while p.poll() is None:
  89             r, _, _ = select.select([sys.stdin, master_fd], [], [], 0.25)
  90             if sys.stdin in r:
  91                 d = os.read(sys.stdin.fileno(), 10240)
  92                 os.write(master_fd, d)
  93             elif master_fd in r:
  94                 o = os.read(master_fd, 10240)
  95                 if o:
  96                     iow.write(o.decode("utf-8", "ignore"))
  97                     iow.flush()
  98     finally:
  99         # restore tty settings back
 100         if sys.stdin.isatty():
 101             termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_tty)
 102
 103
 104 def is_host_regex(restr):
 105     return len(restr) > 2 and restr[0] == "/" and restr[-1] == "/"
 106
 107
 108 def get_host_regex(restr):
 109     if len(restr) < 3 or restr[0] != "/" or restr[-1] != "/":
 110         return None
 111     return re.compile(restr[1:-1])
 112
 113
 114 def host_in(restr, names):
 115     """Determine if matcher is a regex that matches one of names."""
 116     if not (regexp := get_host_regex(restr)):
 117         return restr in names
 118     for name in names:
 119         if regexp.fullmatch(name):
 120             return True
 121     return False
 122
 123
 124 def expand_host(restr, names):
 125     """Expand name or regexp into list of hosts."""
 126     hosts = []
 127     regexp = get_host_regex(restr)
 128     if not regexp:
 129         assert restr in names
 130         hosts.append(restr)
 131     else:
 132         for name in names:
 133             if regexp.fullmatch(name):
 134                 hosts.append(name)
 135     return sorted(hosts)
 136
 137
 138 def expand_hosts(restrs, names):
 139     """Expand list of host names or regex into list of hosts."""
 140     hosts = []
 141     for restr in restrs:
 142         hosts += expand_host(restr, names)
 143     return sorted(hosts)
 144
 145
 146 def host_cmd_split(unet, line, toplevel):
 147     all_hosts = set(unet.hosts)
 148     csplit = line.split()
 149     i = 0
 150     banner = False
 151     for i, e in enumerate(csplit):
 152         if is_re := is_host_regex(e):
 153             banner = True
 154         if not host_in(e, all_hosts):
 155             if not is_re:
 156                 break
 157     else:
 158         i += 1
 159
 160     if i == 0 and csplit and csplit[0] == "*":
 161         hosts = sorted(all_hosts)
 162         csplit = csplit[1:]
 163         banner = True
 164     elif i == 0 and csplit and csplit[0] == ".":
 165         hosts = [unet]
 166         csplit = csplit[1:]
 167     else:
 168         hosts = expand_hosts(csplit[:i], all_hosts)
 169         csplit = csplit[i:]
 170
 171         if not hosts and not csplit[:i]:
 172             if toplevel:
 173                 hosts = [unet]
 174             else:
 175                 hosts = sorted(all_hosts)
 176                 banner = True
 177
 178     if not csplit:
 179         return hosts, "", "", True
 180
 181     i = line.index(csplit[0])
 182     i += len(csplit[0])
 183     return hosts, csplit[0], line[i:].strip(), banner
 184
 185
 186 def win_cmd_host_split(unet, cmd, kinds, defall):
 187     if kinds:
 188         all_hosts = {
 189             x for x in unet.hosts if unet.hosts[x].config.get("kind", "") in kinds
 190         }
 191     else:
 192         all_hosts = set(unet.hosts)
 193
 194     csplit = cmd.split()
 195     i = 0
 196     for i, e in enumerate(csplit):
 197         if not host_in(e, all_hosts):
 198             if not is_host_regex(e):
 199                 break
 200     else:
 201         i += 1
 202
 203     if i == 0 and csplit and csplit[0] == "*":
 204         hosts = sorted(all_hosts)
 205         csplit = csplit[1:]
 206     elif i == 0 and csplit and csplit[0] == ".":
 207         hosts = [unet]
 208         csplit = csplit[1:]
 209     else:
 210         hosts = expand_hosts(csplit[:i], all_hosts)
 211
 212         if not hosts and defall and not csplit[:i]:
 213             hosts = sorted(all_hosts)
 214
 215     # Filter hosts based on cmd
 216     cmd = " ".join(csplit[i:])
 217     return hosts, cmd
 218
 219
 220 def proc_readline(fd, prompt, histfile):
 221     """Read a line of input from user while running in a sub-process."""
 222     # How do we change the command though, that's what's displayed in ps normally
 223     linux.set_process_name("Munet CLI")
 224     try:
 225         # For some reason sys.stdin is fileno == 16 and useless
 226         sys.stdin = os.fdopen(0)
 227         histfile = init_history(None, histfile)
 228         line = input(prompt)
 229         readline.write_history_file(histfile)
 230         if line is None:
 231             os.write(fd, b"\n")
 232         os.write(fd, bytes(f":{str(line)}\n", encoding="utf-8"))
 233     except EOFError:
 234         os.write(fd, b"\n")
 235     except KeyboardInterrupt:
 236         os.write(fd, b"I\n")
 237     except Exception as error:
 238         os.write(fd, bytes(f"E{str(error)}\n", encoding="utf-8"))
 239
 240
 241 async def async_input_reader(rfd):
 242     """Read a line of input from the user input sub-process pipe."""
 243     rpipe = os.fdopen(rfd, mode="r")
 244     reader = asyncio.StreamReader()
 245
 246     def protocol_factory():
 247         return asyncio.StreamReaderProtocol(reader)
 248
 249     loop = asyncio.get_event_loop()
 250     transport, _ = await loop.connect_read_pipe(protocol_factory, rpipe)
 251     o = await reader.readline()
 252     transport.close()
 253
 254     o = o.decode("utf-8").strip()
 255     if not o:
 256         return None
 257     if o[0] == "I":
 258         raise KeyboardInterrupt()
 259     if o[0] == "E":
 260         raise Exception(o[1:])
 261     assert o[0] == ":"
 262     return o[1:]
 263
 264
 265 #
 266 # A lot of work to add async `input` handling without creating a thread. We cannot use
 267 # threads when unshare_inline is used with pid namespace per kernel clone(2)
 268 # restriction.
 269 #
 270 async def async_input(prompt, histfile):
 271     """Asynchronously read a line from the user."""
 272     rfd, wfd = os.pipe()
 273     p = multiprocessing.Process(target=proc_readline, args=(wfd, prompt, histfile))
 274     p.start()
 275     logging.debug("started async_input input process: %s", p)
 276     try:
 277         return await async_input_reader(rfd)
 278     finally:
 279         logging.debug("joining async_input input process")
 280         p.join()
 281
 282
 283 def make_help_str(unet):
 284
 285     w = sorted([x if x else "" for x in unet.cli_in_window_cmds])
 286     ww = unet.cli_in_window_cmds
 287     u = sorted([x if x else "" for x in unet.cli_run_cmds])
 288     uu = unet.cli_run_cmds
 289
 290     s = (
 291         """
 292 Basic Commands:
 293   cli   :: open a secondary CLI window
 294   help  :: this help
 295   hosts :: list hosts
 296   quit  :: quit the cli
 297
 298   HOST can be a host or one of the following:
 299     - '*' for all hosts
 300     - '.' for the parent munet
 301     - a regex specified between '/' (e.g., '/rtr.*/')
 302
 303 New Window Commands:\n"""
 304         + "\n".join([f"  {ww[v][0]}\t:: {ww[v][1]}" for v in w])
 305         + """\nInline Commands:\n"""
 306         + "\n".join([f"  {uu[v][0]}\t:: {uu[v][1]}" for v in u])
 307         + "\n"
 308     )
 309     return s
 310
 311
 312 def get_shcmd(unet, host, kinds, execfmt, ucmd):
 313     if host is None:
 314         h = None
 315         kind = None
 316     elif host is unet or host == "":
 317         h = unet
 318         kind = ""
 319     else:
 320         h = unet.hosts[host]
 321         kind = h.config.get("kind", "")
 322         if kinds and kind not in kinds:
 323             return ""
 324     if not isinstance(execfmt, str):
 325         execfmt = execfmt.get(kind, {}).get("exec", "")
 326     if not execfmt:
 327         return ""
 328
 329     # Do substitutions for {} in string
 330     numfmt = len(re.findall(r"{\d*}", execfmt))
 331     if numfmt > 1:
 332         ucmd = execfmt.format(*shlex.split(ucmd))
 333     elif numfmt:
 334         ucmd = execfmt.format(ucmd)
 335     elif len(re.findall(r"{[a-zA-Z_][0-9a-zA-Z_\.]*}", execfmt)):
 336         if execfmt.endswith('"'):
 337             fstring = "f'''" + execfmt + "'''"
 338         else:
 339             fstring = 'f"""' + execfmt + '"""'
 340         ucmd = eval(  # pylint: disable=W0123
 341             fstring,
 342             globals(),
 343             {"host": h, "unet": unet, "user_input": ucmd},
 344         )
 345     else:
 346         # No variable or usercmd substitution at all.
 347         ucmd = execfmt
 348
 349     # Do substitution for munet variables
 350     ucmd = ucmd.replace("%CONFIGDIR%", str(unet.config_dirname))
 351     if host is None or host is unet:
 352         ucmd = ucmd.replace("%RUNDIR%", str(unet.rundir))
 353         return ucmd.replace("%NAME%", ".")
 354     ucmd = ucmd.replace("%RUNDIR%", str(os.path.join(unet.rundir, host)))
 355     if h.mgmt_ip:
 356         ucmd = ucmd.replace("%IPADDR%", str(h.mgmt_ip))
 357     elif h.mgmt_ip6:
 358         ucmd = ucmd.replace("%IPADDR%", str(h.mgmt_ip6))
 359     if h.mgmt_ip6:
 360         ucmd = ucmd.replace("%IP6ADDR%", str(h.mgmt_ip6))
 361     return ucmd.replace("%NAME%", str(host))
 362
 363
 364 async def run_command(
 365     unet,
 366     outf,
 367     line,
 368     execfmt,
 369     banner,
 370     hosts,
 371     toplevel,
 372     kinds,
 373     ns_only=False,
 374     interactive=False,
 375 ):
 376     """Runs a command on a set of hosts.
 377
 378     Runs `execfmt`. Prior to executing the string the following transformations are
 379     performed on it.
 380
 381     `execfmt` may also be a dictionary of dicitonaries keyed on kind with `exec` holding
 382     the kind's execfmt string.
 383
 384     - if `{}` is present then `str.format` is called to replace `{}` with any extra
 385        input values after the command and hosts are removed from the input.
 386     - else if any `{digits}` are present then `str.format` is called to replace
 387       `{digits}` with positional args obtained from the addittional user input
 388       first passed to `shlex.split`.
 389     - else f-string style interpolation is performed on the string with
 390       the local variables `host` (the current node object or None),
 391       `unet` (the Munet object), and `user_input` (the additional command input)
 392       defined.
 393
 394     The output is sent to `outf`.  If `ns_only` is True then the `execfmt` is
 395     run using `Commander.cmd_status_nsonly` otherwise it is run with
 396     `Commander.cmd_status`.
 397     """
 398     if kinds:
 399         logging.info("Filtering hosts to kinds: %s", kinds)
 400         hosts = [x for x in hosts if unet.hosts[x].config.get("kind", "") in kinds]
 401         logging.info("Filtered hosts: %s", hosts)
 402
 403     if not hosts:
 404         if not toplevel:
 405             return
 406         hosts = [unet]
 407
 408     # if unknowns := [x for x in hosts if x not in unet.hosts]:
 409     #     outf.write("%% Unknown host[s]: %s\n" % ", ".join(unknowns))
 410     #     return
 411
 412     # if sys.stdin.isatty() and interactive:
 413     if interactive:
 414         for host in hosts:
 415             shcmd = get_shcmd(unet, host, kinds, execfmt, line)
 416             if not shcmd:
 417                 continue
 418             if len(hosts) > 1 or banner:
 419                 outf.write(f"------ Host: {host} ------\n")
 420             spawn(unet, host if not toplevel else unet, shcmd, outf, ns_only)
 421             if len(hosts) > 1 or banner:
 422                 outf.write(f"------- End: {host} ------\n")
 423         outf.write("\n")
 424         return
 425
 426     aws = []
 427     for host in hosts:
 428         shcmd = get_shcmd(unet, host, kinds, execfmt, line)
 429         if not shcmd:
 430             continue
 431         if toplevel:
 432             ns = unet
 433         else:
 434             ns = unet.hosts[host] if host and host != unet else unet
 435         if ns_only:
 436             cmdf = ns.async_cmd_status_nsonly
 437         else:
 438             cmdf = ns.async_cmd_status
 439         aws.append(cmdf(shcmd, warn=False, stderr=subprocess.STDOUT))
 440
 441     results = await asyncio.gather(*aws, return_exceptions=True)
 442     for host, result in zip(hosts, results):
 443         if isinstance(result, Exception):
 444             o = str(result) + "\n"
 445             rc = -1
 446         else:
 447             rc, o, _ = result
 448         if len(hosts) > 1 or banner:
 449             outf.write(f"------ Host: {host} ------\n")
 450         if rc:
 451             outf.write(f"*** non-zero exit status: {rc}\n")
 452         outf.write(o)
 453         if len(hosts) > 1 or banner:
 454             outf.write(f"------- End: {host} ------\n")
 455
 456
 457 cli_builtins = ["cli", "help", "hosts", "quit"]
 458
 459
 460 class Completer:
 461     """A completer class for the CLI."""
 462
 463     def __init__(self, unet):
 464         self.unet = unet
 465
 466     def complete(self, text, state):
 467         line = readline.get_line_buffer()
 468         tokens = line.split()
 469         # print(f"\nXXX: tokens: {tokens} text: '{text}' state: {state}'\n")
 470
 471         first_token = not tokens or (text and len(tokens) == 1)
 472
 473         # If we have already have a builtin command we are done
 474         if tokens and tokens[0] in cli_builtins:
 475             return [None]
 476
 477         cli_run_cmds = set(self.unet.cli_run_cmds.keys())
 478         top_run_cmds = {x for x in cli_run_cmds if self.unet.cli_run_cmds[x][3]}
 479         cli_run_cmds -= top_run_cmds
 480         cli_win_cmds = set(self.unet.cli_in_window_cmds.keys())
 481         hosts = set(self.unet.hosts.keys())
 482         is_window_cmd = bool(tokens) and tokens[0] in cli_win_cmds
 483         done_set = set()
 484         if bool(tokens):
 485             if text:
 486                 done_set = set(tokens[:-1])
 487             else:
 488                 done_set = set(tokens)
 489
 490         # Determine the domain for completions
 491         if not tokens or first_token:
 492             all_cmds = (
 493                 set(cli_builtins) | hosts | cli_run_cmds | cli_win_cmds | top_run_cmds
 494             )
 495         elif is_window_cmd:
 496             all_cmds = hosts
 497         elif tokens and tokens[0] in top_run_cmds:
 498             # nothing to complete if a top level command
 499             pass
 500         elif not bool(done_set & cli_run_cmds):
 501             all_cmds = hosts | cli_run_cmds
 502
 503         if not text:
 504             completes = all_cmds
 505         else:
 506             # print(f"\nXXX: all_cmds: {all_cmds} text: '{text}'\n")
 507             completes = {x + " " for x in all_cmds if x.startswith(text)}
 508
 509         # print(f"\nXXX: completes: {completes} text: '{text}' state: {state}'\n")
 510         # remove any completions already present
 511         completes -= done_set
 512         completes = sorted(completes) + [None]
 513         return completes[state]
 514
 515
 516 async def doline(
 517     unet, line, outf, background=False, notty=False
 518 ):  # pylint: disable=R0911
 519
 520     line = line.strip()
 521     m = re.fullmatch(r"^(\S+)(?:\s+(.*))?$", line)
 522     if not m:
 523         return True
 524
 525     cmd = m.group(1)
 526     nline = m.group(2) if m.group(2) else ""
 527
 528     if cmd in ("q", "quit"):
 529         return False
 530
 531     if cmd == "help":
 532         outf.write(make_help_str(unet))
 533         return True
 534     if cmd in ("h", "hosts"):
 535         outf.write(f"% Hosts:\t{' '.join(sorted(unet.hosts.keys()))}\n")
 536         return True
 537     if cmd == "cli":
 538         await remote_cli(
 539             unet,
 540             "secondary> ",
 541             "Secondary CLI",
 542             background,
 543         )
 544         return True
 545
 546     #
 547     # In window commands
 548     #
 549
 550     if cmd in unet.cli_in_window_cmds:
 551         execfmt, toplevel, kinds, kwargs = unet.cli_in_window_cmds[cmd][2:]
 552
 553         # if toplevel:
 554         #     ucmd = " ".join(nline.split())
 555         # else:
 556         hosts, ucmd = win_cmd_host_split(unet, nline, kinds, False)
 557         if not hosts:
 558             if not toplevel:
 559                 return True
 560             hosts = [unet]
 561
 562         if isinstance(execfmt, str):
 563             found_brace = "{}" in execfmt
 564         else:
 565             found_brace = False
 566             for d in execfmt.values():
 567                 if "{}" in d["exec"]:
 568                     found_brace = True
 569                     break
 570         if not found_brace and ucmd and not toplevel:
 571             # CLI command does not expect user command so treat as hosts of which some
 572             # must be unknown
 573             unknowns = [x for x in ucmd.split() if x not in unet.hosts]
 574             outf.write(f"% Unknown host[s]: {' '.join(unknowns)}\n")
 575             return True
 576
 577         try:
 578             if not hosts and toplevel:
 579                 hosts = [unet]
 580
 581             for host in hosts:
 582                 shcmd = get_shcmd(unet, host, kinds, execfmt, ucmd)
 583                 if toplevel or host == unet:
 584                     unet.run_in_window(shcmd, **kwargs)
 585                 else:
 586                     unet.hosts[host].run_in_window(shcmd, **kwargs)
 587         except Exception as error:
 588             outf.write(f"% Error: {error}\n")
 589         return True
 590
 591     #
 592     # Inline commands
 593     #
 594
 595     toplevel = unet.cli_run_cmds[cmd][3] if cmd in unet.cli_run_cmds else False
 596     # if toplevel:
 597     #     logging.debug("top-level: cmd: '%s' nline: '%s'", cmd, nline)
 598     #     hosts = None
 599     #     banner = False
 600     # else:
 601
 602     hosts, cmd, nline, banner = host_cmd_split(unet, line, toplevel)
 603     hoststr = "munet" if hosts == [unet] else f"{hosts}"
 604     logging.debug("hosts: '%s' cmd: '%s' nline: '%s'", hoststr, cmd, nline)
 605
 606     if cmd in unet.cli_run_cmds:
 607         pass
 608     elif "" in unet.cli_run_cmds:
 609         nline = f"{cmd} {nline}"
 610         cmd = ""
 611     else:
 612         outf.write(f"% Unknown command: {cmd} {nline}\n")
 613         return True
 614
 615     execfmt, toplevel, kinds, ns_only, interactive = unet.cli_run_cmds[cmd][2:]
 616     if interactive and notty:
 617         outf.write("% Error: interactive command must be run from primary CLI\n")
 618         return True
 619
 620     await run_command(
 621         unet,
 622         outf,
 623         nline,
 624         execfmt,
 625         banner,
 626         hosts,
 627         toplevel,
 628         kinds,
 629         ns_only,
 630         interactive,
 631     )
 632
 633     return True
 634
 635
 636 async def cli_client(sockpath, prompt="munet> "):
 637     """Implement the user-facing CLI for a remote munet reached by a socket."""
 638     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 639     sock.settimeout(10)
 640     sock.connect(sockpath)
 641
 642     # Go into full non-blocking mode now
 643     sock.settimeout(None)
 644
 645     print("\n--- Munet CLI Starting ---\n\n")
 646     while True:
 647         line = input(prompt)
 648         if line is None:
 649             return
 650
 651         # Need to put \n back
 652         line += "\n"
 653
 654         # Send the CLI command
 655         sock.send(line.encode("utf-8"))
 656
 657         def bendswith(b, sentinel):
 658             slen = len(sentinel)
 659             return len(b) >= slen and b[-slen:] == sentinel
 660
 661         # Collect the output
 662         rb = b""
 663         while not bendswith(rb, ENDMARKER):
 664             lb = sock.recv(4096)
 665             if not lb:
 666                 return
 667             rb += lb
 668
 669         # Remove the marker
 670         rb = rb[: -len(ENDMARKER)]
 671
 672         # Write the output
 673         sys.stdout.write(rb.decode("utf-8"))
 674
 675
 676 async def local_cli(unet, outf, prompt, histfile, background):
 677     """Implement the user-side CLI for local munet."""
 678     assert unet is not None
 679     completer = Completer(unet)
 680     readline.parse_and_bind("tab: complete")
 681     readline.set_completer(completer.complete)
 682
 683     print("\n--- Munet CLI Starting ---\n\n")
 684     while True:
 685         try:
 686             line = await async_input(prompt, histfile)
 687             if line is None:
 688                 return
 689
 690             if not await doline(unet, line, outf, background):
 691                 return
 692         except KeyboardInterrupt:
 693             outf.write("%% Caught KeyboardInterrupt\nUse ^D or 'quit' to exit")
 694
 695
 696 def init_history(unet, histfile):
 697     try:
 698         if histfile is None:
 699             histfile = os.path.expanduser("~/.munet-history.txt")
 700             if not os.path.exists(histfile):
 701                 if unet:
 702                     unet.cmd("touch " + histfile)
 703                 else:
 704                     subprocess.run("touch " + histfile, shell=True, check=True)
 705         if histfile:
 706             readline.read_history_file(histfile)
 707         return histfile
 708     except Exception as error:
 709         logging.warning("init_history failed: %s", error)
 710     return None
 711
 712
 713 async def cli_client_connected(unet, background, reader, writer):
 714     """Handle CLI commands inside the munet process from a socket."""
 715     # # Go into full non-blocking mode now
 716     # client.settimeout(None)
 717     logging.debug("cli client connected")
 718     while True:
 719         line = await reader.readline()
 720         if not line:
 721             logging.debug("client closed cli connection")
 722             break
 723         line = line.decode("utf-8").strip()
 724
 725         class EncodingFile:
 726             """Wrap a writer to encode in utf-8."""
 727
 728             def __init__(self, writer):
 729                 self.writer = writer
 730
 731             def write(self, x):
 732                 self.writer.write(x.encode("utf-8"))
 733
 734             def flush(self):
 735                 self.writer.flush()
 736
 737         if not await doline(unet, line, EncodingFile(writer), background, notty=True):
 738             logging.debug("server closing cli connection")
 739             return
 740
 741         writer.write(ENDMARKER)
 742         await writer.drain()
 743
 744
 745 async def remote_cli(unet, prompt, title, background):
 746     """Open a CLI in a new window."""
 747     try:
 748         if not unet.cli_sockpath:
 749             sockpath = os.path.join(tempfile.mkdtemp("-sockdir", "pty-"), "cli.sock")
 750             ccfunc = functools.partial(cli_client_connected, unet, background)
 751             s = await asyncio.start_unix_server(ccfunc, path=sockpath)
 752             unet.cli_server = asyncio.create_task(s.serve_forever(), name="cli-task")
 753             unet.cli_sockpath = sockpath
 754             logging.info("server created on :\n%s\n", sockpath)
 755
 756         # Open a new window with a new CLI
 757         python_path = await unet.async_get_exec_path(["python3", "python"])
 758         us = os.path.realpath(__file__)
 759         cmd = f"{python_path} {us}"
 760         if unet.cli_histfile:
 761             cmd += " --histfile=" + unet.cli_histfile
 762         if prompt:
 763             cmd += f" --prompt='{prompt}'"
 764         cmd += " " + unet.cli_sockpath
 765         unet.run_in_window(cmd, title=title, background=False)
 766     except Exception as error:
 767         logging.error("cli server: unexpected exception: %s", error)
 768
 769
 770 def add_cli_in_window_cmd(
 771     unet, name, helpfmt, helptxt, execfmt, toplevel, kinds, **kwargs
 772 ):
 773     """Adds a CLI command to the CLI.
 774
 775     The command `cmd` is added to the commands executable by the user from the CLI.  See
 776     `base.Commander.run_in_window` for the arguments that can be passed in `args` and
 777     `kwargs` to this function.
 778
 779     Args:
 780         unet: unet object
 781         name: command string (no spaces)
 782         helpfmt: format of command to display in help (left side)
 783         helptxt: help string for command (right side)
 784         execfmt: interpreter `cmd` to pass to `host.run_in_window()`, if {} present then
 785           allow for user commands to be entered and inserted. May also be a dict of dict
 786           keyed on kind with sub-key of "exec" providing the `execfmt` string for that
 787           kind.
 788         toplevel: run command in common top-level namespaec not inside hosts
 789         kinds: limit CLI command to nodes which match list of kinds.
 790         **kwargs: keyword args to pass to `host.run_in_window()`
 791     """
 792     unet.cli_in_window_cmds[name] = (helpfmt, helptxt, execfmt, toplevel, kinds, kwargs)
 793
 794
 795 def add_cli_run_cmd(
 796     unet,
 797     name,
 798     helpfmt,
 799     helptxt,
 800     execfmt,
 801     toplevel,
 802     kinds,
 803     ns_only=False,
 804     interactive=False,
 805 ):
 806     """Adds a CLI command to the CLI.
 807
 808     The command `cmd` is added to the commands executable by the user from the CLI.
 809     See `run_command` above in the `doline` function and for the arguments that can
 810     be passed in to this function.
 811
 812     Args:
 813         unet: unet object
 814         name: command string (no spaces)
 815         helpfmt: format of command to display in help (left side)
 816         helptxt: help string for command (right side)
 817         execfmt: format string to insert user cmds into for execution. May also be a
 818           dict of dict keyed on kind with sub-key of "exec" providing the `execfmt`
 819           string for that kind.
 820         toplevel: run command in common top-level namespaec not inside hosts
 821         kinds: limit CLI command to nodes which match list of kinds.
 822         ns_only: Should execute the command on the host vs in the node namespace.
 823         interactive: Should execute the command inside an allocated pty (interactive)
 824     """
 825     unet.cli_run_cmds[name] = (
 826         helpfmt,
 827         helptxt,
 828         execfmt,
 829         toplevel,
 830         kinds,
 831         ns_only,
 832         interactive,
 833     )
 834
 835
 836 def add_cli_config(unet, config):
 837     """Adds CLI commands based on config.
 838
 839     All exec strings will have %CONFIGDIR%, %NAME% and %RUNDIR% replaced with the
 840     corresponding config directory and the current nodes `name` and `rundir`.
 841     Additionally, the exec string will have f-string style interpolation performed
 842     with the local variables `host` (node object or None), `unet` (Munet object) and
 843     `user_input` (if provided to the CLI command) defined.
 844
 845     The format of the config dictionary can be seen in the following example.
 846     The first list entry represents the default command because it has no `name` key.
 847
 848       commands:
 849         - help: "run the given FRR command using vtysh"
 850           format: "[HOST ...] FRR-CLI-COMMAND"
 851           exec: "vtysh -c {}"
 852           ns-only: false        # the default
 853           interactive: false    # the default
 854         - name: "vtysh"
 855           help: "Open a FRR CLI inside new terminal[s] on the given HOST[s]"
 856           format: "vtysh HOST [HOST ...]"
 857           exec: "vtysh"
 858           new-window: true
 859         - name: "capture"
 860           help: "Capture packets on a given network"
 861           format: "pcap NETWORK"
 862           exec: "tshark -s 9200 -i {0} -w /tmp/capture-{0}.pcap"
 863           new-window: true
 864           top-level: true # run in top-level container namespace, above hosts
 865
 866     The `new_window` key can also be a dictionary which will be passed as keyward
 867     arguments to the `Commander.run_in_window()` function.
 868
 869     Args:
 870         unet: unet object
 871         config: dictionary of cli config
 872     """
 873     for cli_cmd in config.get("commands", []):
 874         name = cli_cmd.get("name", None)
 875         helpfmt = cli_cmd.get("format", "")
 876         helptxt = cli_cmd.get("help", "")
 877         execfmt = list_to_dict_with_key(cli_cmd.get("exec-kind"), "kind")
 878         if not execfmt:
 879             execfmt = cli_cmd.get("exec", "bash -c '{}'")
 880         toplevel = cli_cmd.get("top-level", False)
 881         kinds = cli_cmd.get("kinds", [])
 882         stdargs = (unet, name, helpfmt, helptxt, execfmt, toplevel, kinds)
 883         new_window = cli_cmd.get("new-window", None)
 884         if isinstance(new_window, dict):
 885             add_cli_in_window_cmd(*stdargs, **new_window)
 886         elif bool(new_window):
 887             add_cli_in_window_cmd(*stdargs)
 888         else:
 889             # on-host is deprecated it really implemented "ns-only"
 890             add_cli_run_cmd(
 891                 *stdargs,
 892                 cli_cmd.get("ns-only", cli_cmd.get("on-host")),
 893                 cli_cmd.get("interactive", False),
 894             )
 895
 896
 897 def cli(
 898     unet,
 899     histfile=None,
 900     sockpath=None,
 901     force_window=False,
 902     title=None,
 903     prompt=None,
 904     background=True,
 905 ):
 906     asyncio.run(
 907         async_cli(unet, histfile, sockpath, force_window, title, prompt, background)
 908     )
 909
 910
 911 async def async_cli(
 912     unet,
 913     histfile=None,
 914     sockpath=None,
 915     force_window=False,
 916     title=None,
 917     prompt=None,
 918     background=True,
 919 ):
 920     if prompt is None:
 921         prompt = "munet> "
 922
 923     if force_window or not sys.stdin.isatty():
 924         await remote_cli(unet, prompt, title, background)
 925
 926     if not unet:
 927         logger.debug("client-cli using sockpath %s", sockpath)
 928
 929     try:
 930         if sockpath:
 931             await cli_client(sockpath, prompt)
 932         else:
 933             await local_cli(unet, sys.stdout, prompt, histfile, background)
 934     except KeyboardInterrupt:
 935         print("\n...^C exiting CLI")
 936     except EOFError:
 937         pass
 938     except Exception as ex:
 939         logger.critical("cli: got exception: %s", ex, exc_info=True)
 940         raise
 941
 942
 943 if __name__ == "__main__":
 944     # logging.basicConfig(level=logging.DEBUG, filename="/tmp/topotests/cli-client.log")
 945     logging.basicConfig(level=logging.DEBUG)
 946     logger = logging.getLogger("cli-client")
 947     logger.info("Start logging cli-client")
 948
 949     parser = argparse.ArgumentParser()
 950     parser.add_argument("--histfile", help="file to user for history")
 951     parser.add_argument("--prompt", help="prompt string to use")
 952     parser.add_argument("socket", help="path to pair of sockets to communicate over")
 953     cli_args = parser.parse_args()
 954
 955     cli_prompt = cli_args.prompt if cli_args.prompt else "munet> "
 956     asyncio.run(
 957         async_cli(
 958             None,
 959             cli_args.histfile,
 960             cli_args.socket,
 961             prompt=cli_prompt,
 962             background=False,
 963         )
 964     )