tests/topotests/lib/micronet_cli.py

   1 # -*- coding: utf-8 eval: (blacken-mode 1) -*-
   2 #
   3 # July 24 2021, Christian Hopps <chopps@labn.net>
   4 #
   5 # Copyright (c) 2021, LabN Consulting, L.L.C.
   6 #
   7 # This program is free software; you can redistribute it and/or
   8 # modify it under the terms of the GNU General Public License
   9 # as published by the Free Software Foundation; either version 2
  10 # of the License, or (at your option) any later version.
  11 #
  12 # This program is distributed in the hope that it will be useful,
  13 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 # GNU General Public License for more details.
  16 #
  17 # You should have received a copy of the GNU General Public License along
  18 # with this program; see the file COPYING; if not, write to the Free Software
  19 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
  20 #
  21 import argparse
  22 import logging
  23 import os
  24 import pty
  25 import re
  26 import readline
  27 import select
  28 import socket
  29 import subprocess
  30 import sys
  31 import tempfile
  32 import termios
  33 import tty
  34
  35
  36 ENDMARKER = b"\x00END\x00"
  37
  38
  39 def lineiter(sock):
  40     s = ""
  41     while True:
  42         sb = sock.recv(256)
  43         if not sb:
  44             return
  45
  46         s += sb.decode("utf-8")
  47         i = s.find("\n")
  48         if i != -1:
  49             yield s[:i]
  50             s = s[i + 1 :]
  51
  52
  53 def spawn(unet, host, cmd):
  54     if sys.stdin.isatty():
  55         old_tty = termios.tcgetattr(sys.stdin)
  56         tty.setraw(sys.stdin.fileno())
  57     try:
  58         master_fd, slave_fd = pty.openpty()
  59
  60         # use os.setsid() make it run in a new process group, or bash job
  61         # control will not be enabled
  62         p = unet.hosts[host].popen(
  63             cmd,
  64             preexec_fn=os.setsid,
  65             stdin=slave_fd,
  66             stdout=slave_fd,
  67             stderr=slave_fd,
  68             universal_newlines=True,
  69         )
  70
  71         while p.poll() is None:
  72             r, w, e = select.select([sys.stdin, master_fd], [], [], 0.25)
  73             if sys.stdin in r:
  74                 d = os.read(sys.stdin.fileno(), 10240)
  75                 os.write(master_fd, d)
  76             elif master_fd in r:
  77                 o = os.read(master_fd, 10240)
  78                 if o:
  79                     os.write(sys.stdout.fileno(), o)
  80     finally:
  81         # restore tty settings back
  82         if sys.stdin.isatty():
  83             termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_tty)
  84
  85
  86 def doline(unet, line, writef):
  87     def host_cmd_split(unet, cmd):
  88         csplit = cmd.split()
  89         for i, e in enumerate(csplit):
  90             if e not in unet.hosts:
  91                 break
  92         hosts = csplit[:i]
  93         if not hosts:
  94             hosts = sorted(unet.hosts.keys())
  95         cmd = " ".join(csplit[i:])
  96         return hosts, cmd
  97
  98     line = line.strip()
  99     m = re.match(r"^(\S+)(?:\s+(.*))?$", line)
 100     if not m:
 101         return True
 102
 103     cmd = m.group(1)
 104     oargs = m.group(2) if m.group(2) else ""
 105     if cmd == "q" or cmd == "quit":
 106         return False
 107     if cmd == "hosts":
 108         writef("%% hosts: %s\n" % " ".join(sorted(unet.hosts.keys())))
 109     elif cmd in ["term", "vtysh", "xterm"]:
 110         args = oargs.split()
 111         if not args or (len(args) == 1 and args[0] == "*"):
 112             args = sorted(unet.hosts.keys())
 113         hosts = [unet.hosts[x] for x in args]
 114         for host in hosts:
 115             if cmd == "t" or cmd == "term":
 116                 host.run_in_window("bash", title="sh-%s" % host)
 117             elif cmd == "v" or cmd == "vtysh":
 118                 host.run_in_window("vtysh", title="vt-%s" % host)
 119             elif cmd == "x" or cmd == "xterm":
 120                 host.run_in_window("bash", title="sh-%s" % host, forcex=True)
 121     elif cmd == "sh":
 122         hosts, cmd = host_cmd_split(unet, oargs)
 123         for host in hosts:
 124             if sys.stdin.isatty():
 125                 spawn(unet, host, cmd)
 126             else:
 127                 if len(hosts) > 1:
 128                     writef("------ Host: %s ------\n" % host)
 129                 output = unet.hosts[host].cmd_legacy(cmd)
 130                 writef(output)
 131                 if len(hosts) > 1:
 132                     writef("------- End: %s ------\n" % host)
 133         writef("\n")
 134     elif cmd == "h" or cmd == "help":
 135         writef(
 136             """
 137 Commands:
 138   help                       :: this help
 139   sh [hosts] <shell-command> :: execute <shell-command> on <host>
 140   term [hosts]               :: open shell terminals for hosts
 141   vtysh [hosts]              :: open vtysh terminals for hosts
 142   [hosts] <vtysh-command>    :: execute vtysh-command on hosts\n\n"""
 143         )
 144     else:
 145         hosts, cmd = host_cmd_split(unet, line)
 146         for host in hosts:
 147             if len(hosts) > 1:
 148                 writef("------ Host: %s ------\n" % host)
 149             output = unet.hosts[host].cmd_legacy('vtysh -c "{}"'.format(cmd))
 150             writef(output)
 151             if len(hosts) > 1:
 152                 writef("------- End: %s ------\n" % host)
 153         writef("\n")
 154     return True
 155
 156
 157 def cli_server_setup(unet):
 158     sockdir = tempfile.mkdtemp("-sockdir", "pyt")
 159     sockpath = os.path.join(sockdir, "cli-server.sock")
 160     try:
 161         sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 162         sock.settimeout(10)
 163         sock.bind(sockpath)
 164         sock.listen(1)
 165         return sock, sockdir, sockpath
 166     except Exception:
 167         unet.cmd_status("rm -rf " + sockdir)
 168         raise
 169
 170
 171 def cli_server(unet, server_sock):
 172     sock, addr = server_sock.accept()
 173
 174     # Go into full non-blocking mode now
 175     sock.settimeout(None)
 176
 177     for line in lineiter(sock):
 178         line = line.strip()
 179
 180         def writef(x):
 181             xb = x.encode("utf-8")
 182             sock.send(xb)
 183
 184         if not doline(unet, line, writef):
 185             return
 186         sock.send(ENDMARKER)
 187
 188
 189 def cli_client(sockpath, prompt="unet> "):
 190     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
 191     sock.settimeout(10)
 192     sock.connect(sockpath)
 193
 194     # Go into full non-blocking mode now
 195     sock.settimeout(None)
 196
 197     print("\n--- Micronet CLI Starting ---\n\n")
 198     while True:
 199         if sys.version_info[0] == 2:
 200             line = raw_input(prompt)  # pylint: disable=E0602
 201         else:
 202             line = input(prompt)
 203         if line is None:
 204             return
 205
 206         # Need to put \n back
 207         line += "\n"
 208
 209         # Send the CLI command
 210         sock.send(line.encode("utf-8"))
 211
 212         def bendswith(b, sentinel):
 213             slen = len(sentinel)
 214             return len(b) >= slen and b[-slen:] == sentinel
 215
 216         # Collect the output
 217         rb = b""
 218         while not bendswith(rb, ENDMARKER):
 219             lb = sock.recv(4096)
 220             if not lb:
 221                 return
 222             rb += lb
 223
 224         # Remove the marker
 225         rb = rb[: -len(ENDMARKER)]
 226
 227         # Write the output
 228         sys.stdout.write(rb.decode("utf-8"))
 229
 230
 231 def local_cli(unet, outf, prompt="unet> "):
 232     print("\n--- Micronet CLI Starting ---\n\n")
 233     while True:
 234         if sys.version_info[0] == 2:
 235             line = raw_input(prompt)  # pylint: disable=E0602
 236         else:
 237             line = input(prompt)
 238         if line is None:
 239             return
 240         if not doline(unet, line, outf.write):
 241             return
 242
 243
 244 def cli(
 245     unet,
 246     histfile=None,
 247     sockpath=None,
 248     force_window=False,
 249     title=None,
 250     prompt=None,
 251     background=True,
 252 ):
 253     if prompt is None:
 254         prompt = "unet> "
 255
 256     if force_window or not sys.stdin.isatty():
 257         # Run CLI in another window b/c we have no tty.
 258         sock, sockdir, sockpath = cli_server_setup(unet)
 259
 260         python_path = unet.get_exec_path(["python3", "python"])
 261         us = os.path.realpath(__file__)
 262         cmd = "{} {}".format(python_path, us)
 263         if histfile:
 264             cmd += " --histfile=" + histfile
 265         if title:
 266             cmd += " --prompt={}".format(title)
 267         cmd += " " + sockpath
 268
 269         try:
 270             unet.run_in_window(cmd, new_window=True, title=title, background=background)
 271             return cli_server(unet, sock)
 272         finally:
 273             unet.cmd_status("rm -rf " + sockdir)
 274
 275     if not unet:
 276         logger.debug("client-cli using sockpath %s", sockpath)
 277
 278     try:
 279         if histfile is None:
 280             histfile = os.path.expanduser("~/.micronet-history.txt")
 281             if not os.path.exists(histfile):
 282                 if unet:
 283                     unet.cmd("touch " + histfile)
 284                 else:
 285                     subprocess.run("touch " + histfile)
 286         if histfile:
 287             readline.read_history_file(histfile)
 288     except Exception:
 289         pass
 290
 291     try:
 292         if sockpath:
 293             cli_client(sockpath, prompt=prompt)
 294         else:
 295             local_cli(unet, sys.stdout, prompt=prompt)
 296     except EOFError:
 297         pass
 298     except Exception as ex:
 299         logger.critical("cli: got exception: %s", ex, exc_info=True)
 300         raise
 301     finally:
 302         readline.write_history_file(histfile)
 303
 304
 305 if __name__ == "__main__":
 306     logging.basicConfig(level=logging.DEBUG, filename="/tmp/topotests/cli-client.log")
 307     logger = logging.getLogger("cli-client")
 308     logger.info("Start logging cli-client")
 309
 310     parser = argparse.ArgumentParser()
 311     parser.add_argument("--histfile", help="file to user for history")
 312     parser.add_argument("--prompt-text", help="prompt string to use")
 313     parser.add_argument("socket", help="path to pair of sockets to communicate over")
 314     args = parser.parse_args()
 315
 316     prompt = "{}> ".format(args.prompt_text) if args.prompt_text else "unet> "
 317     cli(None, args.histfile, args.socket, prompt=prompt)