diff options
Diffstat (limited to 'bin/tunneltop')
| -rwxr-xr-x | bin/tunneltop | 230 |
1 files changed, 230 insertions, 0 deletions
diff --git a/bin/tunneltop b/bin/tunneltop new file mode 100755 index 0000000..ebed5d9 --- /dev/null +++ b/bin/tunneltop @@ -0,0 +1,230 @@ +#!/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()) |