initial commit

1 files changed, 558 insertions, 0 deletions

diff --git a/tunneltop.py b/tunneltop.py
new file mode 100755
index 0000000..b752fbf
--- /dev/null
+++ b/tunneltop.py
@@ -0,0 +1,558 @@
+#!/usr/bin/env python
+"""A top-like program for monitoring ssh tunnels or any tunnels"""
+# TODO- the disabled coloring is not working
+# TODO- quit doesnt work
+# TODO- we are reviving dead tunnels
+import argparse
+import asyncio
+import copy
+import curses
+import enum
+import os
+import signal
+import sys
+import tomllib
+import typing
+
+
+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,"
+            " defaults to $HOME/.tunneltop.toml",
+            default="~/.tunneltop.toml",
+        )
+        self.parser.add_argument(
+            "--noheader",
+            "-n",
+            type=bool,
+            help="Dont print the header in the output",
+            default=False,
+        )
+        self.parser.add_argument(
+            "--delay",
+            "-d",
+            type=float,
+            help="The delay between redraws in seconds, defaults to 5 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,
+    no_color: bool,
+    *args,
+) -> typing.List[str]:
+    """A simple columnar printer"""
+    max_column_width = []
+    lines = []
+    numbers_f: typing.List[int] = []
+    dummy = []
+
+    if no_color or not sys.stdout.isatty():
+        greenie = ""
+        bold = ""
+        endc = ""
+        goo = ""
+        blueblue = ""
+    else:
+        greenie = Colors.greenie
+        bold = Colors.BOLD
+        endc = Colors.ENDC
+        goo = Colors.goo
+        blueblue = Colors.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
+
+
+def render(
+    data_cols: typing.Dict[str, typing.Dict[str, str]],
+    tasks: typing.List[asyncio.Task],
+    stdscr,
+    sel: int,
+):
+    """Render the text"""
+    lines = ffs(
+        2,
+        ["NAME", "ADDRESS", "PORT", "STATUS", "STDOUT", "STDERR"],
+        False,
+        True,
+        [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()],
+    )
+    iterator = iter(lines)
+    stdscr.addstr(1, 1, lines[0], curses.color_pair(3))
+    next(iterator)
+    for i, line in enumerate(iterator):
+        try:
+            line_content = stdscr.instr(sel + 2, 1).decode("utf-8")
+            name: str = line_content[: line_content.find(" ")]
+        finally:
+            name = ""
+        if i == sel:
+            stdscr.addstr(
+                (2 + i) % (len(lines) + 1),
+                1,
+                line,
+                curses.color_pair(2)
+                if name not in tasks
+                else curses.color_pair(5),
+            )
+        else:
+            stdscr.addstr(
+                2 + i,
+                1,
+                line,
+                curses.color_pair(1)
+                if name not in tasks
+                else curses.color_pair(4),
+            )
+        stdscr.addstr("\n")
+    stdscr.box()
+
+
+def curses_init():
+    """Initialize ncurses"""
+    stdscr = curses.initscr()
+    curses.start_color()
+    curses.use_default_colors()
+    curses.curs_set(False)
+    curses.noecho()
+    curses.cbreak()
+    stdscr.keypad(True)
+    curses.halfdelay(20)
+    curses.init_pair(1, curses.COLOR_GREEN, curses.COLOR_BLACK)
+    curses.init_pair(2, curses.COLOR_BLACK, curses.COLOR_GREEN)
+    curses.init_pair(3, curses.COLOR_BLUE, curses.COLOR_BLACK)
+    curses.init_pair(4, curses.COLOR_CYAN, curses.COLOR_BLACK)
+    curses.init_pair(5, curses.COLOR_BLACK, curses.COLOR_CYAN)
+
+    return stdscr
+
+
+class TunnelManager:
+    """The tunnel top class"""
+
+    def __init__(self):
+        self.argparser = Argparser()
+        self.data_cols: typing.Dict[
+            str, typing.Dict[str, str]
+        ] = self.read_conf()
+        self.tunnel_tasks: typing.List[asyncio.Task] = []
+        self.tunnel_test_tasks: typing.List[asyncio.Task] = []
+        self.scheduler_task: asyncio.Task
+        self.scheduler_table: typing.Dict[
+            str, int
+        ] = self.init_scheduler_table()
+        # we use this when its time to quit. this will prevent any
+        # new tasks from being scheduled
+        self.are_we_dying: bool = False
+
+    def init_scheduler_table(self) -> typing.Dict[str, int]:
+        """initialize the scheduler table"""
+        result: typing.Dict[str, int] = {}
+        for key, value in self.data_cols.items():
+            if "test_interval" in value and value["test_command"] != "":
+                result[key] = 0
+
+        return result
+
+    async def stop_task(
+        self,
+        delete_task: asyncio.Task,
+        task_list: typing.List[asyncio.Task],
+        delete: bool = True,
+    ):
+        """Remove the reference"""
+        delete_index: int = -1
+        delete_task.cancel()
+        self.write_log(f"{delete_task.get_name()} is being cancelled\n")
+        await asyncio.sleep(0)
+        for i, task in enumerate(task_list):
+            if task.get_name() == delete_task.get_name():
+                delete_index = i
+                break
+
+        if delete and delete_index >= 0:
+            task_list.remove(self.tunnel_tasks[delete_index])
+
+    def read_conf(self) -> typing.Dict[str, typing.Dict[str, str]]:
+        """Read the config file"""
+        data_cols: typing.Dict[str, typing.Dict[str, str]] = {}
+        with open(
+            os.path.expanduser(self.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": "",
+                }
+        return data_cols
+
+    async def run_subprocess(self, cmd: str) -> typing.Tuple[bytes, bytes]:
+        """Run a command"""
+        try:
+            proc = await asyncio.create_subprocess_exec(
+                *cmd.split(" "),
+                stdout=asyncio.subprocess.PIPE,
+                stderr=asyncio.subprocess.PIPE,
+            )
+
+            return await proc.communicate()
+        except asyncio.TimeoutError:
+            proc.terminate()
+            raise
+        except asyncio.CancelledError:
+            proc.terminate()
+            raise
+
+    async def run_test_coro(
+        self, cmd: str, task_name: str
+    ) -> typing.Tuple[bytes, bytes]:
+        """Run a test command"""
+        try:
+            stdout, stderr = await self.run_subprocess(cmd)
+            stdout_str: str = stdout.decode("utf-8").strip("\n").strip('"')
+            stderr_str: str = stderr.decode("utf-8").strip("\n").strip('"')
+
+            self.data_cols[task_name]["stdout"] = stdout_str
+            self.data_cols[task_name]["stderr"] = stderr_str
+            if stdout_str == self.data_cols[task_name]["test_command_result"]:
+                self.data_cols[task_name]["status"] = "UP"
+            else:
+                self.data_cols[task_name]["status"] = "DOWN"
+
+            return stdout, stderr
+        except asyncio.TimeoutError:
+            self.data_cols[task_name]["status"] = "TMOUT"
+            self.data_cols[task_name]["stdout"] = "-"
+            self.data_cols[task_name]["stderr"] = "-"
+            raise
+
+    async def tunnel_procs(
+        self,
+    ) -> typing.List[asyncio.Task]:
+        """run all the tunnels in the background as separate subprocesses"""
+        tasks: typing.List[asyncio.Task] = []
+        for _, value in self.data_cols.items():
+            tasks.append(
+                asyncio.create_task(
+                    self.run_subprocess(value["command"]), name=value["name"]
+                ),
+            )
+            await asyncio.sleep(0)
+
+        return tasks
+
+    async def sighup_handler_async_worker(self, data_cols_new) -> None:
+        """Handles the actual updating of tasks when we get SIGTERM"""
+        delete_task: typing.Optional[asyncio.Task] = None
+        for k, value in data_cols_new.items():
+            if k not in self.data_cols:
+                self.tunnel_tasks.append(
+                    asyncio.create_task(
+                        self.run_subprocess(value["command"]), name=k
+                    )
+                )
+                await asyncio.sleep(0)
+                self.data_cols[k] = copy.deepcopy(value)
+                if k in self.scheduler_table:
+                    self.scheduler_table[k] = 0
+            else:
+                if (
+                    self.data_cols[k]["command"] != data_cols_new[k]["command"]
+                    or self.data_cols[k]["port"] != data_cols_new[k]["port"]
+                    or self.data_cols[k]["address"]
+                    != data_cols_new[k]["address"]
+                ):
+                    for task in self.tunnel_tasks:
+                        if task.get_name() == k:
+                            delete_task = task
+                            break
+                            # task.cancel()
+                            # await asyncio.sleep(0)
+
+                    if delete_task is not None:
+                        await self.stop_task(delete_task, self.tunnel_tasks)
+                        delete_task = None
+                    self.data_cols[k] = copy.deepcopy(data_cols_new[k])
+                    self.tunnel_tasks.append(
+                        asyncio.create_task(
+                            self.run_subprocess(value["command"]), name=k
+                        )
+                    )
+                    if k in self.scheduler_table:
+                        self.scheduler_table[k] = 0
+                    await asyncio.sleep(0)
+
+        for k, _ in self.data_cols.items():
+            if k not in data_cols_new:
+                for task in self.tunnel_tasks:
+                    if task.get_name() == k:
+                        # task.cancel()
+                        # await asyncio.sleep(0)
+                        delete_task = task
+                        break
+                if delete_task is not None:
+                    await self.stop_task(delete_task, self.tunnel_tasks)
+                    delete_task = None
+                del self.data_cols[k]
+                if k in self.scheduler_table:
+                    del self.scheduler_table[k]
+
+    async def sighup_handler(self) -> None:
+        """SIGHUP handler. we want to reload the config."""
+        # type: ignore # pylint: disable=E0203
+        data_cols_new: typing.Dict[str, typing.Dict[str, str]] = {}
+        data_cols_new = self.read_conf()
+        await self.sighup_handler_async_worker(data_cols_new)
+
+    def write_log(self, log: str):
+        """A simple logger"""
+        with open(
+            "/home/devi/devi/abbatoir/hole15/log",
+            "a",
+            encoding="utf-8",
+        ) as logfile:
+            logfile.write(log)
+
+    async def restart_task(self, line_content: str) -> None:
+        """restart a task"""
+        name: str = line_content[: line_content.find(" ")]
+        # was_cancelled: bool = False
+        for task in self.tunnel_tasks:
+            if task.get_name() == name:
+                # was_cancelled = task.cancel()
+                # self.write_log(f"was_cancelled: {was_cancelled}")
+                await self.stop_task(task, self.tunnel_tasks)
+                # await task
+                # await asyncio.sleep(0)
+                for _, value in self.data_cols.items():
+                    if value["name"] == name and task.cancelled():
+                        self.tunnel_tasks.append(
+                            asyncio.create_task(
+                                self.run_subprocess(value["command"]),
+                                name=value["name"],
+                            )
+                        )
+                        await asyncio.sleep(0)
+
+    async def flip_task(self, line_content: str) -> None:
+        """flip a task"""
+        name: str = line_content[: line_content.find(" ")]
+        # was_cancelled: bool = False
+        was_active: bool = False
+        for task in self.tunnel_tasks:
+            if task.get_name() == name:
+                await self.stop_task(task, self.tunnel_tasks)
+                # was_cancelled = task.cancel()
+                # await asyncio.sleep(0)
+                # self.write_log(f"was_cancelled: {was_cancelled}")
+                # await task
+                was_active = True
+                break
+
+        if not was_active:
+            for _, value in self.data_cols.items():
+                if value["name"] == name:
+                    self.tunnel_tasks.append(
+                        asyncio.create_task(
+                            self.run_subprocess(value["command"]),
+                            name=value["name"],
+                        )
+                    )
+                    await asyncio.sleep(0)
+
+    async def quit(self) -> None:
+        """Cleanly quit the applicaiton"""
+        # scheduler checks for this to stop running new tests
+        # when we want to quit
+        self.are_we_dying = True
+
+        for task in asyncio.all_tasks():
+            task.cancel()
+            await asyncio.sleep(0)
+        try:
+            await asyncio.gather(*asyncio.all_tasks())
+        finally:
+            sys.exit(0)
+
+    async def scheduler(self) -> None:
+        """scheduler manages running the tests and reviving dead tunnels"""
+        try:
+            while True:
+                if self.are_we_dying:
+                    return
+                for key, value in self.scheduler_table.items():
+                    if value == 0 and key not in self.tunnel_test_tasks:
+                        tunnel_entry = self.data_cols[key]
+                        test_task = asyncio.create_task(
+                            asyncio.wait_for(
+                                self.run_test_coro(
+                                    tunnel_entry["test_command"],
+                                    tunnel_entry["name"],
+                                ),
+                                timeout=float(tunnel_entry["test_timeout"]),
+                            ),
+                            name=key,
+                        )
+                        self.tunnel_test_tasks.append(test_task)
+                        self.scheduler_table[key] = int(
+                            tunnel_entry["test_interval"]
+                        )
+                        await asyncio.sleep(0)
+                    else:
+                        self.scheduler_table[key] = (
+                            self.scheduler_table[key] - 1
+                        )
+
+                # we are using a 1 second ticker. basically the scheduler
+                # runs every second instead of as fast as it can
+                await asyncio.sleep(1)
+        except asyncio.CancelledError:
+            pass
+
+    async def tui_loop(self) -> None:
+        """the tui loop"""
+        sel: int = 0
+        try:
+            stdscr = curses_init()
+            # we spawn the tunnels and the test scheduler put them
+            # in the background and then run the TUI loop
+            self.tunnel_tasks = await self.tunnel_procs()
+            self.scheduler_task = asyncio.create_task(
+                self.scheduler(), name="scheduler"
+            )
+            loop = asyncio.get_event_loop()
+            loop.add_signal_handler(
+                signal.SIGHUP,
+                lambda: asyncio.create_task(self.sighup_handler()),
+            )
+
+            while True:
+                stdscr.clear()
+                render(self.data_cols, self.tunnel_tasks, stdscr, sel)
+                char = stdscr.getch()
+
+                if char == ord("j") or char == curses.KEY_DOWN:
+                    sel = (sel + 1) % len(self.data_cols)
+                elif char == ord("k") or char == curses.KEY_UP:
+                    sel = (sel - 1) % len(self.data_cols)
+                elif char == ord("g") or char == curses.KEY_UP:
+                    sel = 0
+                elif char == ord("G") or char == curses.KEY_UP:
+                    sel = len(self.data_cols) - 1
+                elif char == ord("r"):
+                    line_content = stdscr.instr(sel + 2, 1)
+                    await self.restart_task(line_content.decode("utf-8"))
+                elif char == ord("q"):
+                    await self.quit()
+                elif char == ord("s"):
+                    line_content = stdscr.instr(sel + 2, 1)
+                    await self.flip_task(line_content.decode("utf-8"))
+
+                stdscr.refresh()
+                await asyncio.sleep(0)
+        finally:
+            curses.nocbreak()
+            stdscr.keypad(False)
+            curses.echo()
+            curses.endwin()
+            # tasks = asyncio.all_tasks()
+            # for task in tasks:
+            #     task.cancel()
+            await self.quit()
+
+
+def main() -> None:
+    """entry point"""
+    tunnel_manager = TunnelManager()
+    asyncio.run(tunnel_manager.tui_loop())
+
+
+if __name__ == "__main__":
+    main()