#!/usr/bin/env python """A top-like program for monitoring ssh tunnels""" import argparse import asyncio import enum import sys import typing import tomllib class Argparser: # pylint: disable=too-few-public-methods """Argparser class.""" def __init__(self): self.parser = argparse.ArgumentParser() self.parser.add_argument( "--config", "-c", type=str, help="the path to the .tunneltop.toml file", default="/home/devi/.tunneltop.toml", ) self.parser.add_argument( "--noheader", "-n", type=bool, help="dont print the header", default=False, ) self.parser.add_argument( "--delay", "-d", type=float, help="The delay between updates in seconds", default=5, ) self.args = self.parser.parse_args() # pylint: disable=too-few-public-methods class Colors(enum.EnumType): """static color definitions""" purple = "\033[95m" blue = "\033[94m" green = "\033[92m" yellow = "\033[93m" red = "\033[91m" grey = "\033[1;37m" darkgrey = "\033[1;30m" cyan = "\033[1;36m" ENDC = "\033[0m" BOLD = "\033[1m" UNDERLINE = "\033[4m" blueblue = "\x1b[38;5;24m" greenie = "\x1b[38;5;23m" goo = "\x1b[38;5;22m" screen_clear = "\033c\033[3J" hide_cursor = "\033[?25l" # pylint: disable=too-many-locals def ffs( offset: int, header_list: typing.Optional[typing.List[str]], numbered: bool, *args, ) -> typing.List[str]: """A simple columnar printer""" max_column_width = [] lines = [] numbers_f: typing.List[int] = [] dummy = [] if sys.stdout.isatty(): greenie = Colors.greenie bold = Colors.BOLD endc = Colors.ENDC goo = Colors.goo blueblue = Colors.blueblue else: greenie = "" bold = "" endc = "" goo = "" blueblue = "" for arg in args: max_column_width.append(max(len(repr(argette)) for argette in arg)) if header_list is not None: if numbered: numbers_f.extend(range(1, len(args[-1]) + 1)) max_column_width.append( max(len(repr(number)) for number in numbers_f) ) header_list.insert(0, "idx") index = range(0, len(header_list)) for header, width, i in zip(header_list, max_column_width, index): max_column_width[i] = max(len(header), width) + offset for i in index: dummy.append( greenie + bold + header_list[i].ljust(max_column_width[i]) + endc ) lines.append("".join(dummy)) dummy.clear() index2 = range(0, len(args[-1])) for i in index2: if numbered: dummy.append( goo + bold + repr(i).ljust(max_column_width[0]) + endc ) for arg, width in zip(args, max_column_width[1:]): dummy.append(blueblue + (arg[i]).ljust(width) + endc) else: for arg, width in zip(args, max_column_width): dummy.append(blueblue + (arg[i]).ljust(width) + endc) lines.append("".join(dummy)) dummy.clear() return lines async def run_subshell(cmd: str) -> typing.Tuple[bytes, bytes]: """Run a command in a subshell""" proc = await asyncio.create_subprocess_shell( cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE ) # return stdout and stderr return await proc.communicate() def tunnel_test_callback(task: asyncio.Task) -> None: """Tunnel test callback function.""" task_name = task.get_name() # data_cols["stdout"] = task.result()[0] # data_cols["stderr"] = task.result()[1] if ( task.result()[0].decode("utf-8").strip("\n") == data_cols[task_name]["test_command_result"] ): data_cols[task_name]["status"] = "UP" else: data_cols[task_name]["status"] = "DOWN" async def tunnel_test_procs() -> typing.List[asyncio.Task]: """run all the tunnel tests in the background as separate tasks""" tasks: typing.List[asyncio.Task] = [] for _, value in data_cols.items(): if value["test_command"] != "": tasks.append( asyncio.create_task( run_subshell(value["test_command"]), name=value["name"] ) ) tasks[-1].add_done_callback(tunnel_test_callback) await asyncio.sleep(0) return tasks async def tunnel_procs(commands: typing.List[str]) -> None: """run all the tunnels in the background as separate tasks""" for command in commands: asyncio.create_task(run_subshell(command)) await asyncio.sleep(0) data_cols: typing.Dict[str, typing.Dict] = {} async def main() -> None: """entrypoint""" argparser = Argparser() print(Colors.screen_clear, end="") print(Colors.hide_cursor, end="") with open(argparser.args.config, "rb") as conf_file: data = tomllib.load(conf_file) for key, value in data.items(): data_cols[key] = { "name": key, "address": value["address"], "port": value["port"], "command": value["command"], "status": "UNKN", "test_command": value["test_command"], "test_command_result": value["test_command_result"], "test_interval": value["test_interval"], "test_timeout": value["test_timeout"], "stdout": "", "stderr": "", } await tunnel_procs([v["command"] for _, v in data_cols.items()]) while True: await tunnel_test_procs() lines = ffs( 2, ["NAME", "ADDRESS", "PORT", "STATUS", "STDOUT", "STDERR"] if not argparser.args.noheader else None, False, [v["name"] for _, v in data_cols.items()], [v["address"] for _, v in data_cols.items()], [repr(v["port"]) for _, v in data_cols.items()], [v["status"] for _, v in data_cols.items()], [v["stdout"] for _, v in data_cols.items()], [v["stderr"] for _, v in data_cols.items()], ) for line in lines: print(line) await asyncio.sleep(argparser.args.delay) print(Colors.screen_clear, end="") print(Colors.hide_cursor, end="") if __name__ == "__main__": asyncio.run(main())