#!/bin/bash
#
# Author: Martin 'BruXy' Bruchanov, bruchy at gmail
# Version: 1.6 (Sat 22 Jul 2023 10:25:48 PM ADT)
#
# URL: http://bruxy.regnet.cz/web/linux/EN/socks-via-ssh/
# GitHub: https://github.com/BruXy/bash-utils/tree/master/socks-via-ssh
#
#%A
# SOCKS SSH Tunnels
# =================
#
# Usage:
# ------
#
#   sstun start|stop|restart|status|help|list
#
#         help    ... Show this help
#         start   ... Enable all tunnels listed in ~/.sstunrc.
#         stop    ... Disconnect all tunnels.
#         restart ... Cycle stop->start.
#         status  ... Get PIDs of ssh processes and proxy info.
#         list    ... List enabled proxies.
#         alias   ... Create aliases for nmap, ssh, curl.
#
# Configuration:
# --------------
#
#   1. Create your SSH server entries in ~/.ssh/config.
#   2. Specify option DynamicForward <PORT>, to open SOCKS tunnel
#   3. Create ~/.sstunrc and list all hosts you will use as proxy.
#   4. Run this script ./sstun start.
#
# Trick for SOCKS proxy binded to localhost:1080
# ----------------------------------------------
#
# 1. Scan host via proxy:
#
#    nmap -sV -Pn -n --proxies socks4://127.0.0.1:1080 scanme.nmap.org
#
# 2. HTTP request via proxy:
#
#    curl --user-agent "Mozilla" --socks4 localhost:1080 http://ifconfig.co
#
# 3. SSH via proxy:
#
#    ssh -o ProxyCommand='nc --proxy-type socks4 --proxy 127.0.0.1:1080 %h %p' user@target
#
# 4. Some programs can use SOCKS via system proxy settings:
#
#    export http_proxy=socks5://127.0.0.1:1080
#    export https_proxy=socks5://127.0.0.1:1080
#
#    youtube-dl "youtube.com/watch?V=..."
#
# Aliases:
# --------
#
# Use: eval $(sstun alias)
#
# It will import aliases for curl, nmap and ssh to your current shell.
# Format of alias is for example: curl_<proxy_id>.
#
# Disable aliases: unalias $(sstun alias | sed 's/alias \([^=]*\)=.*/\1/')
#
#%B

####################
# Global variables #
####################

LOG="/tmp/${0}.log"
LOCK="$HOME/.sstun.run"
TUN_BIND="localhost"
HOST_LIST="$HOME/.sstunrc"
SSH_HOSTS=()
AGENT="Mozilla/5.0 (Macintosh; Intel Mac OS X 10_9_3) AppleWebKit/537.75.14"
AGENT+=" (KHTML, like Gecko) Version/7.0.3 Safari/7046A194A"
declare -A PID_LIST # hashmap [host]=PID

#############
# Functions #
#############

#=== FUNCTION ============================================================
#        Name: print_help
# Description: Display header help text between #%A and #%B tokens.
#=========================================================================

function print_help() {
    sed -ne '/^#%A/,/^#%B/{s/\(^#\|^#%[AB]\)//p}' $0
}


#=== FUNCTION ============================================================
#        Name: get_child_pid
# Description: Print PID of child process, only 1 child process expected.
# Parameter 1: Parent process ID (PPID)
#     Returns: 0 OK, 1 on error.
#=========================================================================

function get_child_pid() {
    local pid=$1
    # check if $pid of parent exists
    if ! kill -0 "$pid" >& /dev/null ; then
        printf "$0 ERROR -- Process PID %d does not exists!\n" "$pid"
        return 1
    fi
    # get child PID
    local retval=$(ps --ppid $pid -h -o pid)
    if [ -z "$retval" ] ; then
        printf "$0 WARNING -- Cannot find child process of %d!\n" $pid
        return 1
    else
        echo $retval
        return 0
    fi
}


#=== FUNCTION ============================================================
#        Name: read_host_list
# Description: Read list of host from $HOST_LIST file. Each line has one
#              hostname defined in ~/.ssh/config. Lines can be commented
#              by '#' sign.
#     Returns: 0 when list exists and non empty, 1 when not found.
#=========================================================================

function read_host_list() {
    local list=''
    if [ -f "$HOST_LIST" ] ; then
        SSH_HOSTS=( $(grep -Ev "^\s*#|^$" $HOST_LIST) )
        list=$( sed -e 's/ /, /g' <<< ${SSH_HOSTS[*]} )
        if [ ${#list} -lt 1 ] ; then
            printf "$0: ERROR -- $HOST_LIST does not exist!\n" >&2
        fi
        printf "$0: INFO -- Host list contains: $list.\n"
    else
        printf "$0: ERROR -- $HOST_LIST does not exist!\n" >&2
        exit 1
    fi
}


#=== FUNCTION ============================================================
#        Name: read_lock_file
# Description: Read lock file with host=PID records and convert it into
#              hash array stored in $PID_LIST global variable.
#=========================================================================

function read_lock_file() {
    declare -A host_list
    for line in $(cat $LOCK)
    do
        read host pid<<<$(sed 's/=/ /'<<<$line)
        PID_LIST[$host]=$pid
    done
    printf "$0: INFO -- %s\n" "$(set | grep ^PID_LIST)"
}


#=== FUNCTION ============================================================
#        Name: check_lockfile
# Description: Check if lock file exist, when it exist it will display
#              details about TCP connections provided by SSH processes.
#     Returns: 1 when lockfile is present, 0 when does not.
#=========================================================================

function check_lockfile() {
    local exists=0
    if [ -f "$LOCK" ] ; then
        exists=1
        printf "$0: INFO -- Lock file exists ($LOCK)!\n" 2>&1
        printf "$0: INFO -- List of running processes: \n" 2>&1
        for pid in $(sed -ne '/^SSH:/s///p' $LOCK)
        do
            conn=$(lsof -P -p $pid | grep -E "ESTABLISHED|LISTEN")
            if [ -z "$conn" ] ; then
                printf "$0: $pid is not running!\n"  2>&1
            else
                printf "$0: PID $pid is running:\n" 2>&1
                printf "$conn\n" 2>&1
            fi
        done
    fi
    return $exists
}


#=== FUNCTION ============================================================
#        Name: create_tunnels
# Description: Starts autossh in background, stores its PID in $LOCK file
#              and also PID of its child process ssh which maintains SSH
#              communication itself. Lock file is later used for connection
#              check and for disabling tunnels.
# Parameter 1: Reads list of hosts from SSH_HOSTS array.
#=========================================================================

function create_tunnels() {
    for host in ${SSH_HOSTS[*]}
    do
        # Autossh will create a new process and its PID is not
        # the !$ but it is stored in AUTOSSH_PIDFILE. Autossh
        # will fork child process which executes ssh connection
        # itself.

        export AUTOSSH_PIDFILE=/tmp/autossh-$[RANDOM].pid
        autossh -M 0 -v -t -f -N $host &
        wait $! # it will take same time to initialize autossh
        pid=$(cat $AUTOSSH_PIDFILE)
        rm -f $AUTOSSH_PIDFILE
        echo "$0: started tunnel on host $host with PID = $pid"

        # Store autossh PIDs:
        echo "$host=$pid" >> $LOCK

        # Note: Only autossh PID is important, it is starting child
        #       process with actuall ssh connection and this process
        #       my time out and autossh will automatically start a
        #       new ssh process to reconnect.
    done
}


#=== FUNCTION ============================================================
#        Name: stop_tunnels
# Description: Send SIGINT signal to all autossh processes listed in lock
#              file. Delete lock file.
#=========================================================================

function stop_tunnels() {
    read_lock_file
    for host in ${!PID_LIST[*]}
    do
        pid=${PID_LIST[$host]}
        printf "$0 INFO -- Closing connection to '%s'\n" $host
        printf "$0 INFO -- Sending SIGINT to PID %d\n" $pid
        kill -s SIGINT $pid
        retval=$?
        if [ $retval != 0 ] ; then
            printf "$0 WARNING -- kill retval = $retval\n"
        fi
    done
    rm -f $LOCK
}


#=== FUNCTION ============================================================
#        Name: get_hostname
# Description: Query DNS and return hostname for given IP.
# Parameter 1: IP Address.
#     Returns: List of hostnames separated by comma.
#=========================================================================

function get_hostname() {
    local ip=$1
    dig +short -x $ip | \
        sed -e "s/\.%//" -e "s/\.$/, /" -e '$s/,\s*$//' | \
        tr -d '\n'
}


#=== FUNCTION ============================================================
#        Name: ip_status
# Description: Check remote ip over SOCKS proxy. It will check all enabled
#              SSH processes, detects open ports and get public IP by
#              HTTP request on http://ifconfig.co/ site.
#=========================================================================

function ip_status() {
    read_lock_file
    if [ -f $LOCK ] ; then
        for host in ${!PID_LIST[*]}
        do
            ppid=${PID_LIST[$host]}
            child_pid=$(get_child_pid $ppid)
            # Test if parent and child process exist.
            if [ $? -ne 0 ] ; then
                printf "$0: ERROR -- %s ssh process does not exist!\n" $host >&2
                continue
            fi
            printf "$0: INFO -- %s autossh (ppid: %d) -> ssh (pid: %d)\n" \
                "$host" "$ppid" "$child_pid"
            # Note 1: lsof on Ubuntu complains about tracefs to stderr,
            #         filtered by '|&'
            # Note 2: I have some connections with more then one tunnel.
            ports=$(lsof -P -p $child_pid |& \
                    sed -ne 's/.* localhost:\([0-9]*\) .*/\1/p' | uniq)
            printf "$0: INFO --   TCP port: %s\n" $ports
            for p in $ports
            do
                remote_ip=$(curl --silent --socks5 $TUN_BIND:$p \
                    --user-agent "$AGENT" http://ifconfig.co/json | \
                    jq -r '.ip' )
                if [ ! -z "$remote_ip" ] ; then
                    printf "$0: SOCKS proxy %s (%s) connects via: %s, %s; %s\n"\
                        "$host" "$TUN_BIND:$p" "$remote_ip" \
                        "$(get_hostname $remote_ip)" \
                        "$(geo_location $TUN_BIND:$p)"
                else
                    printf "$0: $host has no proxy enabled on port '%d', (pid=%d)!\n" "$p" "$pid"
                fi
            done
        done
    else
        printf "$0: ERROR -- No lock file found, was sstun started?\n" >&2
    fi
}


#=== FUNCTION ============================================================
#        Name: format_table
# Description: Format output of 'status' command as table.
#=========================================================================

function format_table() {
    sed -r -e '1i Id; Port; IP Address; Hostname; Location' \
       -e 's/^.*proxy ([^ ]*)/\1;/' \
       -e 's/\(([^)]+)\)/\1;/' \
       -e 's/connects via: ([0-9.]+),/\1;/' | column -s';' -t | \
       sed '1{p;s/./-/g}'
}


#=== FUNCTION ============================================================
#        Name: geo_location
# Description: Query ifconfig.co to get geo location.
# Parameter 1: Proxy bind, e.g $TUN_BIND:port
#     Returns: Country,City to stdout
#=========================================================================

function geo_location() {
    local socks_proxy=$1

    # Parse output with jq
    if jq --help &> /dev/null ; then
        curl --socks5 $socks_proxy --user-agent "$AGENT" \
            -s http://ifconfig.co/json | \
            jq -r '[.country,.city] | join(", ")'
    else
    # Without jq
        curl --socks5 $socks_proxy --user-agent "$AGENT" \
            -s http://ifconfig.co/json | \
            tr ',' '\n' | \
            sed -ne '/"country"\|"city"/{s/.*:"\(.*\)"$/\1/;p}' | \
            sed "N;s/\n/, /"
    fi
}


#=== FUNCTION ============================================================
#        Name: create_aliases
# Description: Create aliases for curl and nmap to use SOCKS proxy
#     Returns: N/A
#=========================================================================

function create_aliases() {
    ip_status | sed -rne '/SOCKS/{s/.* ([^ ]*) \(([^)]*)\).*/\1 \2/;p}' |
        while read proxy_id port
        do
            # Nmap does not like 'localhost', rather use 127.0.0.1:
            printf "alias nmap_%s='nmap --proxies socks4://127.1:%s ';\n" \
                $proxy_id ${port/localhost:/}
            printf "alias curl_%s='curl --socks4 %s ';\n" $proxy_id $port
            printf "alias ssh_%s='ssh -o ProxyCommand=\"nc --proxy-type socks4 --proxy %s %%h %%p\"';\n" $proxy_id $port
        done
}


#=== FUNCTION ============================================================
#        Name: wait_spinner
# Description: Display message and spin while waiting to finish something.
#              Run in background (&), get its PID ($!), run task and kill
#              the PID once the task is finished.
# Parameter 1: Printed message
#     Returns: N/A
#=========================================================================

function wait_spinner() {
    local msg="$*"
    local sp="◴◷◶◵"
    while true
    do
        printf "\b${msg} ${sp:i++%${#sp}:1}\r"
        sleep 0.25
    done
}


########
# Main #
########

# Display help
if [ $# -eq 0 ] ; then print_help 1>&2; exit 1; fi
if [[ "$*" =~ -h ]] ; then print_help; exit 0; fi

case $1 in
    start)
        if check_lockfile ; then
            read_host_list
            create_tunnels
            sleep 2 # wait for connections to establish
            ip_status
            exit 0
        else
            printf "No tunnel(s) enabled, restarting!\n"
            rm -v -f $LOCK
            $0 start
            exit 0
        fi
        ;;
    stop)
        read_lock_file
        stop_tunnels
        ;;
    restart)
        $0 stop
        $0 start
        ;;
    status)
        check_lockfile
        ip_status
        ;;
    list|socks|info)
        wait_spinner "Checking connection..." &
        pid=$!
        ip_status | grep SOCKS | sort -n | format_table
        kill $pid
        ;;
    alias)
        create_aliases
        ;;
    help)
        print_help
        exit 1
        ;;
    *)
        printf "Uknown command: '$1'!\n" 2>&1
        exit 1
        ;;
esac