ssht

#!/bin/bash

####################
# GLOBAL VARIABLES
####################

declare -A tunnel_params
DEFAULT_CONFIG_PATH="$HOME/.config/ssht/ssht.yml"
GREEN='\033[0;32m' # Used for std output
RED='\033[0;31m'   # Used for std output
NC='\033[0m' # No Color (Used for std output)

####################

##
# Entry point of the script.
main() {
    parse_args $@

    # If no file is specified the default file is chosen
    if [ -z "$FILE" ] ; then
        FILE=$DEFAULT_CONFIG_PATH
    else
        if [[ ! -f "$FILE" ]] ; then
            error_unknown_file $FILE
        fi
    fi

    run_command $@
}

##
# Parses command line arguments.
parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -f|--file)
                FILE="$2"
                OPTIONS+=("$1")
                shift # past argument
                shift # past value
            ;;
            -*|--*)
                error_unknown_option $1
            ;;
            *)
                POSITIONAL_ARGS+=("$1") # save positional arg
                shift # past argument
            ;;
        esac
    done
}

##
# Runs the cli command provided by the user.
run_command() {
    command=$1
    case $command in
        "") help_main ;;
        "open")  tunnel_open  ${POSITIONAL_ARGS[@]:1} ;;
        "close") tunnel_close ${POSITIONAL_ARGS[@]:1} ;;
        "status") tunnel_status ${POSITIONAL_ARGS[@]:1} ;;
        "show") tunnel_show ${POSITIONAL_ARGS[@]:1} ;;
        "list") tunnel_list  ;;
        "help") help ${@:2}  ;;
        *) 
            echo "Error - Command not found." 
            print_available_commands
        ;;
    esac
}

####################
# TUNNEL FUNCTIONS
####################
tunnel_open() {
    if [ ${#POSITIONAL_ARGS[@]} -gt 2 ]; then
        error_too_many_params ${#POSITIONAL_ARGS[@]}
    fi

    process_name=$1

    if [ -z $process_name ] ; then
        tunnel_open_all
        exit 0
    fi

    if [[ $process_name == *\** ]]; then
        for i in $(tunnel_search_wildcard $process_name) ; do
            tunnel_open $i
        done
    else
        tunnel_read_config $process_name
        key=""
        if [ -n "${tunnel_params['ssh_key_path']}" ]; then
            key="-i ${tunnel_params['ssh_key_path']}"
        fi

        ssh_cmd="ssh -N -L ${tunnel_params['port_forward']}:${tunnel_params['host_destination']}:${tunnel_params['port_destination']} ${tunnel_params['user_server']}@${tunnel_params['host_server']} -f $key"
        echo "Running: $ssh_cmd"
        eval $ssh_cmd
    fi


}

tunnel_open_all() {
        all_configs=$(yq '.tunnels | keys | .[]' $FILE)
        for config in $all_configs ; do
            tunnel_open $config
        done
}

tunnel_search_wildcard() {
    yq ".tunnels | keys | .[]" $FILE | grep "^[^#]" | grep "$1"
}

# Reads the YAML config file and loads it into an array to be used globally.
tunnel_read_config() {
    process_name=$1
    tunnel_params=()

    while IFS="=" read -r key value; do tunnel_params["$key"]=$value; done < <(
        yq ".tunnels.\"$process_name\" | to_entries | map([.key, .value] | join(\"=\")) | .[]" $FILE
    )

    if [ ${#tunnel_params[@]} -eq 0 ] ; then
        error_process_name_not_found $1
    fi
}

tunnel_close() {
    if [ ${#POSITIONAL_ARGS[@]} -gt 2 ]; then
        error_too_many_params ${#POSITIONAL_ARGS[@]}
    fi
    process_name=$1

    if [ -z $process_name ] ; then
        tunnel_close_all
        exit 0
    fi
    if [[ $process_name == *\** ]]; then
        for i in $(tunnel_search_wildcard $process_name) ; do
            tunnel_close $i
        done
    else
        port=$(tunnel_get_port_forward $process_name)

        if [[ $port == "null" ]] ; then
            error_process_name_not_found $process_name
        fi
        cmd="fuser $port/tcp -k"
        echo "Running (killing the process_name): $cmd"
        eval $cmd
    fi

}

tunnel_close_all() {
    all_configs=$(yq '.tunnels | keys | .[]' $FILE)
    for config in $all_configs ; do
        tunnel_close $config
    done
}

tunnel_get_port_forward() {
    process_name=$1
    yq ".tunnels.\"$process_name\".port_forward" $FILE
}

tunnel_status() {
    if [ ${#POSITIONAL_ARGS[@]} -gt 2 ]; then
        error_too_many_params ${#POSITIONAL_ARGS[@]}
    fi

    process_name=$1

    if [ -z $process_name ] ; then
        tunnel_status_all
        exit 0
    fi

    port=$(tunnel_get_port_forward $process_name)
    if [[ $port == "null" ]] ; then
        error_process_name_not_found $process_name
    fi
    
    # We check if the port is from is LISTENING and it's an ssh command
    if nc -z localhost $port ; then
        echo -e "${GREEN}open${NC}"
    else
        echo -e "${RED}closed${NC}"
    fi
    
}

tunnel_status_all() {
    all_configs=$(yq '.tunnels | keys | .[]' $FILE)
    for config in $all_configs ; do
        echo "- $config:"
        echo "  $(tunnel_status $config)"
    done
}

tunnel_show() {
    if [ ${#POSITIONAL_ARGS[@]} -gt 2 ]; then
        error_too_many_params ${#POSITIONAL_ARGS[@]}
    fi

    process_name=$1
    if [ -z $process_name ] ; then
        yq $FILE
        exit 0
    fi

    yq ".tunnels.\"$process_name\"" $FILE
}

tunnel_exists() {
    process_name=$1
    result=$(yq ".tunnels.\"$process_name\" | key" $FILE)
    if [[ $result == "null" ]] ; then
        return 1
    else
        return 0
    fi
}

tunnel_list() {
    if [ ${#POSITIONAL_ARGS[@]} -gt 1 ]; then
        error_too_many_params ${#POSITIONAL_ARGS[@]}
    fi

    yq '... comments="" | .tunnels | keys ' $FILE
}

####################
# ERROR FUNCTIONS
####################
error_process_name_not_found() {
    process_name=$1
    >&2 echo "Error - The name \"$process_name\" couldn't be found in the config file."
    echo "Use \"ssht list\" to display the available tunnels."
    exit 1
}

error_too_many_params() {
    param_count=$1
    >&2 echo "Error - Too many positional params for this command ($param_count)"
    exit 2
}

error_unknown_option() {
    option=$1
    >&2 echo "Error - Unknown option \"$option\""
    exit 3
}

error_unknown_file() {
    file=$1
    >&2 echo "Error - The file \"$file\" couldn't be found"
    exit 4
}
error_help_command_options() {
    option=$1
    >&2 echo "Error - The help command does't allow any options. You provided \"$option\""
    exit 5
}

####################
# HELP FUNCTIONS
####################
help() {
    if [ ${#OPTIONS[@]} -gt 0 ]; then
        error_help_command_options ${OPTIONS[0]}
    fi

    if [ ${#POSITIONAL_ARGS[@]} -gt 2 ]; then
        error_too_many_params ${#POSITIONAL_ARGS[@]}
    fi

    command=$1
    case $command in
        "") help_main  ;;
        "open") help_open  ;;
        "close") help_close  ;;
        "status") help_status  ;;
        "show") help_show  ;;
        "list") help_list  ;;
        "help") help_help  ;;
        *) 
            echo "Error - Command not found." 
            print_available_commands
        ;;
    esac
}

help_main() {
    cat << EOF
Manages SSH tunnels from a YAML configuration.

Usage: 
ssht <command> [-f|--file <file>]

Commands:
    - open [query]    Opens ssh tunnels
    - close [query]   Closes ssh tunnels
    - status [query]  Tells if a tunnel is opened or closed.
    - show [query]    Shows an ssht.yml configuration
    - list            Lists the available tunnels in the configuration
    - help [command]  Shows the help description for a command

Options:
    - -f|--file   The configuration file. Default if none is 
                  provided: ~/.config/ssht/ssht.yml

Help with specific command: 
ssht help <command>

Example: 
ssht help open
ssht help help

For more info visit: https://github.com/markelca/ssh-tunnels#example
EOF
}

help_open() {
    cat << EOF
Opens an ssh tunnel from a YAML configuration given an optional query.

Usage:
ssht open [query] [-f|--file <file>]

The query must match the name of some tunnel in your config file.
If no query is provided it will open all the tunnels available.

Examples:
ssht open remote_database
ssht open private_database -f my-config.yml
EOF
}

help_status() {
    echo "Status help"
    cat << EOF
Shows if a tunnel is opened or closed.

Usage:
ssht status [query] [-f|--file <file>]

It only shows if a process is running on the forwaded port, so it
could give false positives (e.g. zombie process, network related 
problems). To ensure connectivity use tools like autossh.

The query must match the name of some tunnel in your config file.
If no query is provided it will show the status for all the tunnels 
available.

Examples:
ssht status remote_database
ssht status private_database -f my-config.yml
EOF
}
help_close() {
    cat << EOF
Closes an ssh tunnel from a YAML configuration given a specific query.

Usage:
ssht close [query] [-f|--file <file>]

The query must match the name of some tunnel in your config file.
If no query is provided it will close all the tunnels available.

Examples:
ssht close remote_database
ssht close private_database -f my-config.yml

The example above is using the ssht.example.yml config file from the repository.
EOF

}

help_show() {
    cat << EOF
Displays the ssht.yml configuration for a given a query.

The query must match the name of some tunnel in your config file.
If no query is provided it will show the entire config file.

Usage:
ssht show [query] [-f|--file <file>]

Examples:
ssht show
ssht show remote_database
ssht show -f ssht.example.yml
EOF
}

help_list() {
    cat << EOF
Lists the available ssht tunnels from a YAML configuration.

Usage: 
ssht list [-f|--file <file>]

Examples:
ssht list
ssht list -f ssht.example.yml
EOF
}

print_available_commands() {
    cat << EOF
The available commands are:
    - open
    - close
    - status
    - show
    - list
    - help
EOF
}

help_help() {
    cat << EOF
Shows help text for a specific command.

Usage: 
ssht help [command]

Examples:
ssht help
ssht help open
EOF
}

main $@