#!/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()
