dotfiles/_custom/.local/bin/ssh-tun

#!/usr/bin/env python3

"""
Convenience wrapper for SSH tunnel creation and control.

Allows spawning and exiting SSH tunnel processes in the background easily.
Reverse tunnels work too. (see help text)
Currently only supports port-based forwarding, no sockets.

Example usage:
$ ssh-tun user@hostname start 8000 9020
Sets up a forwarding tunnel at user@hostname from local port 8000 to remote port 9020.
$ ssh-tun user@hostname check
Returns the PID of the tunnel to user@hostname (if one exists)
$ ssh-tun user@hostname stop
Exits the process of the tunnel to user@hostname (if one is running)

More help:
$ ssh-tun -h
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
from subprocess import run


# Mandatory always:
DESTINATION = 'destination'
# Optional to specify control socket:
CTL_SOCK_PATH, CTL_SOCK_PATH_SHORT = 'control_socket_path', 'S'
CTL_SOCK_DEFAULT_FILE_NAME = '.ssh-tunnel-ctl'
CTL_SOCK_PATH_DEFAULT = Path(Path(), CTL_SOCK_DEFAULT_FILE_NAME)
# Sub-commands:
CMD = 'cmd'
START, CHECK, STOP = 'start', 'check', 'stop'
# Mandatory for starting:
PORT, HOSTPORT = 'entry_port', 'exit_port'
# Optional for starting:
BIND_ADDRESS, BIND_ADDRESS_SHORT = 'bind_address', 'b'
HOST, HOST_SHORT = 'host', 'H'
REVERSE, REVERSE_SHORT = 'reverse', 'R'

# SSH constants:
SSH = 'ssh'
CON_CTL_CMD = 'ctl_cmd'
CON_CTL_CMD_SHORT = 'O'
CTL_CHECK = 'check'
CTL_EXIT = 'exit'
FORWARD_SHORT = 'L'
TUNNEL_FLAGS = 'fNM'  # run in the background (f), don't execute remote command (N), and use "master" mode (M)


def main() -> None:
    parser = ArgumentParser(description="Work with ssh tunnels in the background")
    parser.add_argument(
        DESTINATION,
        help="The destination that ssh is supposed to connect to, e.g. `[user@]hostname`"
    )
    parser.add_argument(
        f'-{CTL_SOCK_PATH_SHORT}', f'--{CTL_SOCK_PATH.replace("_", "-")}',
        type=Path,
        default=CTL_SOCK_PATH_DEFAULT,
        help=f"Specifies the location of a control socket; defaults to {CTL_SOCK_PATH_DEFAULT}"
    )
    sub_parsers = parser.add_subparsers(dest=CMD)
    # Starting a new tunnel:
    parser_start = sub_parsers.add_parser(START, help="Start a new ssh tunnel")
    parser_start.add_argument(
        PORT,
        help="The entry port of the tunnel (called `port` in the ssh manual)"
    )
    parser_start.add_argument(
        HOSTPORT,
        help="The exit port of the tunnel (called `hostport` in the ssh manual)"
    )
    parser_start.add_argument(
        f'-{HOST_SHORT}', f'--{HOST}',
        default='localhost',
        help="The host to connect to at the tunnel exit; defaults to `localhost`"
    )
    parser_start.add_argument(
        f'-{BIND_ADDRESS_SHORT}', f'--{BIND_ADDRESS.replace("_", "-")}',
        help="The address to listen to at the tunnel entrance. When doing regular forwarding, by default the local "
             "port is bound in accordance with the GatewayPorts setting; when doing reverse forwarding listening "
             "socket on the server will be bound to the loopback interface only."
    )
    parser_start.add_argument(
        f'-{REVERSE_SHORT}', f'--{REVERSE}',
        action='store_true',
        help="If set, a reverse tunnel is created, where connections to the given TCP port on the remote (server) host "
             "are to be forwarded to the local side; if not set, connections to the given TCP port on the local "
             "(client) host are to be forwarded to the given host and port on the remote side."
    )
    parser_start.set_defaults(func=handle_start)
    parser_check = sub_parsers.add_parser(CHECK, help="Check the status of an existing ssh tunnel")
    parser_check.set_defaults(func=handle_check)
    parser_stop = sub_parsers.add_parser(STOP, help="Close an existing ssh tunnel")
    parser_stop.set_defaults(func=handle_stop)

    args = parser.parse_args()
    try:
        args.func(args)
    except AttributeError:
        print("Command not specified! \n")
        parser.parse_args([])


def handle_start(args: Namespace) -> None:
    args = vars(args)
    tunnel_spec = f'{args[PORT]}:{args[HOST]}:{args[HOSTPORT]}'
    if args[BIND_ADDRESS]:
        tunnel_spec = f"{args[BIND_ADDRESS]}:{tunnel_spec}"
    tunnel_type = f'-{REVERSE_SHORT if args[REVERSE] else FORWARD_SHORT}'
    run([
        SSH,
        tunnel_type, tunnel_spec,
        f'-{TUNNEL_FLAGS}',
        f'-{CTL_SOCK_PATH_SHORT}', args[CTL_SOCK_PATH],
        args[DESTINATION]
    ])


def handle_check(args: Namespace) -> None:
    setattr(args, CON_CTL_CMD, CTL_CHECK)
    _handle_control(args)


def handle_stop(args: Namespace) -> None:
    setattr(args, CON_CTL_CMD, CTL_EXIT)
    _handle_control(args)


def _handle_control(args: Namespace) -> None:
    args = vars(args)
    run([
        SSH,
        f'-{CTL_SOCK_PATH_SHORT}', args[CTL_SOCK_PATH],
        f'-{CON_CTL_CMD_SHORT}', args[CON_CTL_CMD],
        args[DESTINATION]
    ])


if __name__ == '__main__':
    main()