summaryrefslogtreecommitdiffstats
path: root/tests/topotests/lib/micronet_cli.py
diff options
context:
space:
mode:
Diffstat (limited to 'tests/topotests/lib/micronet_cli.py')
-rw-r--r--tests/topotests/lib/micronet_cli.py319
1 files changed, 319 insertions, 0 deletions
diff --git a/tests/topotests/lib/micronet_cli.py b/tests/topotests/lib/micronet_cli.py
new file mode 100644
index 0000000..ef804f6
--- /dev/null
+++ b/tests/topotests/lib/micronet_cli.py
@@ -0,0 +1,319 @@
+# -*- coding: utf-8 eval: (blacken-mode 1) -*-
+#
+# July 24 2021, Christian Hopps <chopps@labn.net>
+#
+# Copyright (c) 2021, LabN Consulting, L.L.C.
+#
+# This program is free software; you can redistribute it and/or
+# modify it under the terms of the GNU General Public License
+# as published by the Free Software Foundation; either version 2
+# of the License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License along
+# with this program; see the file COPYING; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA
+#
+import argparse
+import logging
+import os
+import pty
+import re
+import readline
+import select
+import socket
+import subprocess
+import sys
+import tempfile
+import termios
+import tty
+
+
+ENDMARKER = b"\x00END\x00"
+
+
+def lineiter(sock):
+ s = ""
+ while True:
+ sb = sock.recv(256)
+ if not sb:
+ return
+
+ s += sb.decode("utf-8")
+ i = s.find("\n")
+ if i != -1:
+ yield s[:i]
+ s = s[i + 1 :]
+
+
+def spawn(unet, host, cmd):
+ if sys.stdin.isatty():
+ old_tty = termios.tcgetattr(sys.stdin)
+ tty.setraw(sys.stdin.fileno())
+ try:
+ master_fd, slave_fd = pty.openpty()
+
+ # use os.setsid() make it run in a new process group, or bash job
+ # control will not be enabled
+ p = unet.hosts[host].popen(
+ cmd,
+ preexec_fn=os.setsid,
+ stdin=slave_fd,
+ stdout=slave_fd,
+ stderr=slave_fd,
+ universal_newlines=True,
+ )
+
+ while p.poll() is None:
+ r, w, e = select.select([sys.stdin, master_fd], [], [], 0.25)
+ if sys.stdin in r:
+ d = os.read(sys.stdin.fileno(), 10240)
+ os.write(master_fd, d)
+ elif master_fd in r:
+ o = os.read(master_fd, 10240)
+ if o:
+ os.write(sys.stdout.fileno(), o)
+ finally:
+ # restore tty settings back
+ if sys.stdin.isatty():
+ termios.tcsetattr(sys.stdin, termios.TCSADRAIN, old_tty)
+
+
+def doline(unet, line, writef):
+ def host_cmd_split(unet, cmd):
+ csplit = cmd.split()
+ for i, e in enumerate(csplit):
+ if e not in unet.hosts:
+ break
+ hosts = csplit[:i]
+ if not hosts:
+ hosts = sorted(unet.hosts.keys())
+ cmd = " ".join(csplit[i:])
+ return hosts, cmd
+
+ line = line.strip()
+ m = re.match(r"^(\S+)(?:\s+(.*))?$", line)
+ if not m:
+ return True
+
+ cmd = m.group(1)
+ oargs = m.group(2) if m.group(2) else ""
+ if cmd == "q" or cmd == "quit":
+ return False
+ if cmd == "hosts":
+ writef("%% hosts: %s\n" % " ".join(sorted(unet.hosts.keys())))
+ elif cmd in ["term", "vtysh", "xterm"]:
+ args = oargs.split()
+ if not args or (len(args) == 1 and args[0] == "*"):
+ args = sorted(unet.hosts.keys())
+ hosts = [unet.hosts[x] for x in args if x in unet.hosts]
+ for host in hosts:
+ if cmd == "t" or cmd == "term":
+ host.run_in_window("bash", title="sh-%s" % host)
+ elif cmd == "v" or cmd == "vtysh":
+ host.run_in_window("vtysh", title="vt-%s" % host)
+ elif cmd == "x" or cmd == "xterm":
+ host.run_in_window("bash", title="sh-%s" % host, forcex=True)
+ elif cmd == "sh":
+ hosts, cmd = host_cmd_split(unet, oargs)
+ for host in hosts:
+ if sys.stdin.isatty():
+ spawn(unet, host, cmd)
+ else:
+ if len(hosts) > 1:
+ writef("------ Host: %s ------\n" % host)
+ output = unet.hosts[host].cmd_legacy(cmd)
+ writef(output)
+ if len(hosts) > 1:
+ writef("------- End: %s ------\n" % host)
+ writef("\n")
+ elif cmd == "h" or cmd == "help":
+ writef(
+ """
+Commands:
+ help :: this help
+ sh [hosts] <shell-command> :: execute <shell-command> on <host>
+ term [hosts] :: open shell terminals for hosts
+ vtysh [hosts] :: open vtysh terminals for hosts
+ [hosts] <vtysh-command> :: execute vtysh-command on hosts\n\n"""
+ )
+ else:
+ hosts, cmd = host_cmd_split(unet, line)
+ for host in hosts:
+ if len(hosts) > 1:
+ writef("------ Host: %s ------\n" % host)
+ output = unet.hosts[host].cmd_legacy('vtysh -c "{}"'.format(cmd))
+ writef(output)
+ if len(hosts) > 1:
+ writef("------- End: %s ------\n" % host)
+ writef("\n")
+ return True
+
+
+def cli_server_setup(unet):
+ sockdir = tempfile.mkdtemp("-sockdir", "pyt")
+ sockpath = os.path.join(sockdir, "cli-server.sock")
+ try:
+ sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+ sock.settimeout(10)
+ sock.bind(sockpath)
+ sock.listen(1)
+ return sock, sockdir, sockpath
+ except Exception:
+ unet.cmd_status("rm -rf " + sockdir)
+ raise
+
+
+def cli_server(unet, server_sock):
+ sock, addr = server_sock.accept()
+
+ # Go into full non-blocking mode now
+ sock.settimeout(None)
+
+ for line in lineiter(sock):
+ line = line.strip()
+
+ def writef(x):
+ xb = x.encode("utf-8")
+ sock.send(xb)
+
+ if not doline(unet, line, writef):
+ return
+ sock.send(ENDMARKER)
+
+
+def cli_client(sockpath, prompt="unet> "):
+ sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+ sock.settimeout(10)
+ sock.connect(sockpath)
+
+ # Go into full non-blocking mode now
+ sock.settimeout(None)
+
+ print("\n--- Micronet CLI Starting ---\n\n")
+ while True:
+ if sys.version_info[0] == 2:
+ line = raw_input(prompt) # pylint: disable=E0602
+ else:
+ line = input(prompt)
+ if line is None:
+ return
+
+ # Need to put \n back
+ line += "\n"
+
+ # Send the CLI command
+ sock.send(line.encode("utf-8"))
+
+ def bendswith(b, sentinel):
+ slen = len(sentinel)
+ return len(b) >= slen and b[-slen:] == sentinel
+
+ # Collect the output
+ rb = b""
+ while not bendswith(rb, ENDMARKER):
+ lb = sock.recv(4096)
+ if not lb:
+ return
+ rb += lb
+
+ # Remove the marker
+ rb = rb[: -len(ENDMARKER)]
+
+ # Write the output
+ sys.stdout.write(rb.decode("utf-8"))
+
+
+def local_cli(unet, outf, prompt="unet> "):
+ print("\n--- Micronet CLI Starting ---\n\n")
+ while True:
+ if sys.version_info[0] == 2:
+ line = raw_input(prompt) # pylint: disable=E0602
+ else:
+ line = input(prompt)
+ if line is None:
+ return
+ if not doline(unet, line, outf.write):
+ return
+
+
+def cli(
+ unet,
+ histfile=None,
+ sockpath=None,
+ force_window=False,
+ title=None,
+ prompt=None,
+ background=True,
+):
+ logger = logging.getLogger("cli-client")
+
+ if prompt is None:
+ prompt = "unet> "
+
+ if force_window or not sys.stdin.isatty():
+ # Run CLI in another window b/c we have no tty.
+ sock, sockdir, sockpath = cli_server_setup(unet)
+
+ python_path = unet.get_exec_path(["python3", "python"])
+ us = os.path.realpath(__file__)
+ cmd = "{} {}".format(python_path, us)
+ if histfile:
+ cmd += " --histfile=" + histfile
+ if title:
+ cmd += " --prompt={}".format(title)
+ cmd += " " + sockpath
+
+ try:
+ unet.run_in_window(cmd, new_window=True, title=title, background=background)
+ return cli_server(unet, sock)
+ finally:
+ unet.cmd_status("rm -rf " + sockdir)
+
+ if not unet:
+ logger.debug("client-cli using sockpath %s", sockpath)
+
+ try:
+ if histfile is None:
+ histfile = os.path.expanduser("~/.micronet-history.txt")
+ if not os.path.exists(histfile):
+ if unet:
+ unet.cmd("touch " + histfile)
+ else:
+ subprocess.run("touch " + histfile)
+ if histfile:
+ readline.read_history_file(histfile)
+ except Exception:
+ pass
+
+ try:
+ if sockpath:
+ cli_client(sockpath, prompt=prompt)
+ else:
+ local_cli(unet, sys.stdout, prompt=prompt)
+ except EOFError:
+ pass
+ except Exception as ex:
+ logger.critical("cli: got exception: %s", ex, exc_info=True)
+ raise
+ finally:
+ readline.write_history_file(histfile)
+
+
+if __name__ == "__main__":
+ logging.basicConfig(level=logging.DEBUG, filename="/tmp/topotests/cli-client.log")
+ logger = logging.getLogger("cli-client")
+ logger.info("Start logging cli-client")
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--histfile", help="file to user for history")
+ parser.add_argument("--prompt-text", help="prompt string to use")
+ parser.add_argument("socket", help="path to pair of sockets to communicate over")
+ args = parser.parse_args()
+
+ prompt = "{}> ".format(args.prompt_text) if args.prompt_text else "unet> "
+ cli(None, args.histfile, args.socket, prompt=prompt)