From 31d783b9d28659f5108121d1c2ea464116c82aa3 Mon Sep 17 00:00:00 2001
From: zjx20 <zhoujianxiong2@gmail.com>
Date: Sat, 10 Sep 2022 17:51:22 +0800
Subject: [PATCH] added ssh support

---
 README.md            | 23 +++++++++++--
 activate             | 79 +++++++++++++++++++++++++++++++++++++++++---
 deactivate           | 31 ++++++++++-------
 socksproxyenv.sample | 12 +++----
 supports/dsocks      | 63 ++---------------------------------
 supports/ssh         |  4 +++
 6 files changed, 126 insertions(+), 86 deletions(-)
 create mode 100644 supports/ssh

diff --git a/README.md b/README.md
index da00f13..360c340 100644
--- a/README.md
+++ b/README.md
@@ -13,33 +13,50 @@ Here is an incomplete list of supported commands:
 * gem
 * npm
 * mvn
+* ssh
 * ...
 
 ## Usage
 
 1. Clone the code.
+	```bash
+	git clone https://github.com/zjx20/socks-cli.git
+	```
 
 2. Copy `socksproxyenv.sample` to `socksproxyenv`, and fill your socks5 server into it.
+	```bash
+	cd socks-cli
+	cp socksproxyenv.sample socksproxyenv
+
+	# edit socksproxyenv, complete the line:
+	#   export SOCKS_PROXY=
+	```
 
-3. Call `source socks-cli/activate` before your CLI commands:
+3. Invoke `source socks-cli/activate` before running your CLI commands:
 	```
 	$ source socks-cli/activate
 	Serving HTTP proxy on 127.0.0.1 port 54967 ...
-	Done! Some environment variables have been changed to:
+	Done! Variables or aliases have been changed to:
 	  GIT_PROXY_COMMAND=/Users/x/socks-git/sh/socksified-connect.sh
 	  GIT_SSH=/Users/x/socks-git/sh/socksified-ssh.sh
 	  ALL_PROXY=socks5h://127.0.0.1:1080
 	  HTTP_PROXY=http://127.0.0.1:54967
 	  HTTPS_PROXY=http://127.0.0.1:54967
 
+	# Following commands will use the socks proxy!
+
 	$ git clone git@github.com:git/git.git
 	Cloning into 'git'...
 	remote: Counting objects: 213208, done.
 	remote: Compressing objects: 100% (372/372), done.
 	Receiving objects 2.0% (1/213208), 620.00 KiB | 121.00 KiB/s
 	...
+
+	# Check your external IP!
+	$ curl ipinfo.io
+	...
 	```
 
-4. Optionally, you can call `source socks-cli/deactivate` to deactivate `socks-cli`.
+4. Optionally, you can invoke `source socks-cli/deactivate` to deactivate `socks-cli`.
 
 For more details, please see [socksproxyenv.sample](socksproxyenv.sample).
diff --git a/activate b/activate
index 28cd75a..16d23fa 100644
--- a/activate
+++ b/activate
@@ -4,7 +4,7 @@
 function _logerror() {
     local RED='\033[0;31m'
     local NC='\033[0m'
-    echo $(echo -en "${RED}")"$@"$(echo -en "${NC}")
+    echo "$(echo -en "${RED}")$*$(echo -en "${NC}")"
 }
 
 # Determine the directory containing this script
@@ -37,8 +37,8 @@ fi
 SOCKS_CLI_DIR="$("${SOCKS_CLI_DIR}/py/py_wrapper" -c "import os,sys;print(os.path.abspath(sys.argv[1]))" "${SOCKS_CLI_DIR}")"
 
 declare -a _socks_cli_saved_vars=()
-declare -a _socks_cli_supports=()
 declare -a _socks_cli_export_vars=()
+declare -a _socks_cli_supports=()
 
 function _expand() {
     eval "(echo \"\${$1}\") 2>/dev/null"
@@ -67,6 +67,26 @@ function _contains() {
     for e in "${@:2}"; do [[ "$e" = "$1" ]] && return 0; done; return 1
 }
 
+function _clean_invalid_shadows() {
+    local tmp_path="${SOCKS_CLI_DIR}/tmp/shadows/"
+    [ ! -d "${tmp_path}" ] && return 0
+    # First pass: remove command shadows if the owning shell has been dead.
+    find "${tmp_path}" -name "_shell_pid" -print0 |
+        while IFS= read -r -d '' file; do
+            pid="$(tr -d '\n\r' < "${file}")"
+            if [ "${pid}" = "$$" ]; then
+                # Ignore the current shell
+                continue
+            fi
+            if ! pgrep -F "${file}" > /dev/null; then
+                find "$(dirname "${file}")" -type f -delete
+            fi
+        done
+
+    # Second pass: remove empty directories recursively.
+    find "${tmp_path}" -empty -type d -delete
+}
+
 function LOAD_SUPPORT() {
     if [ $# -lt 1 ]; then
         _logerror "Error: LOAD_SUPPORT() requires one parameter."
@@ -87,7 +107,7 @@ function LOAD_SUPPORT() {
 
 function EXPORT_ENV() {
     if [ $# -lt 2 ]; then
-        _logerror "Error: DEFINE_ENV() requires two parameters."
+        _logerror "Error: EXPORT_ENV() requires two parameters."
         return 1
     fi
     if _contains "$1" "${_socks_cli_export_vars[@]}" ; then
@@ -99,8 +119,59 @@ function EXPORT_ENV() {
     eval "export $1=\"$2\""
 }
 
+function SHADOW() {
+    if [ $# -lt 2 ]; then
+        _logerror "Error: SHADOW() requires at least two parameters."
+        return 1
+    fi
+
+    local shell_pid=$$
+    local tmp_path="${SOCKS_CLI_DIR}/tmp/shadows/${shell_pid}"
+    local path_preappend="${tmp_path}"
+    if [ -n "${PATH}" ]; then
+        path_preappend="${path_preappend}:"
+    fi
+
+    if [ "${_socks_cli_shadow_inited}" = "" ]; then
+        _clean_invalid_shadows
+
+        mkdir -p "${tmp_path}"
+        # Remove existing command shadows
+        [ -n "${tmp_path}" ] && find "${tmp_path}/" -type f -delete
+        echo "${shell_pid}" > "${tmp_path}/_shell_pid"
+
+        EXPORT_ENV PATH "${path_preappend}${PATH}"
+        _socks_cli_shadow_inited=1
+    fi
+
+    local shadow_name="$1"
+    local start_command="${@:2}"
+    local cmd_shadow="${tmp_path}/${shadow_name}"
+    cat <<-EOF > "${cmd_shadow}"
+		#!/bin/bash
+		function _sed_escape() {
+		    echo "\$1" | sed -e 's/\\\\/\\\\\\\\/g; s/\\//\\\\\\//g; s/&/\\\\\\&/g'
+		}
+		function _replace() {
+		    local search=\$(_sed_escape "\$2")
+		    local replace=\$(_sed_escape "\$3")
+		    echo "\$1" | sed "s/\${search}/\${replace}/g"
+		}
+		# remove the current folder from the PATH,
+		# to avoid recursively invoking this script.
+		export PATH="\$(_replace "\$PATH" "${path_preappend}" "")"
+		${start_command} "\$@"
+	EOF
+    chmod +x "${cmd_shadow}"
+}
+
 source "${SOCKS_CLI_DIR}/socksproxyenv"
 
+# Rehash executables if PATH has been changed
+if [[ " ${_socks_cli_export_vars[*]} " =~ " PATH " ]]; then
+    hash -r
+fi
+
 # Prompt that socks-cli is activated
 _save_var "PS1"
 PS1_PREFIX=$'\e[0;32msocks-cli\e[0m '
@@ -114,7 +185,7 @@ fi
 export PS1=$PS1_PREFIX$PS1
 
 _socks_cli=1
-echo "Done! Some environment variables have been changed to:"
+echo "Done! Environment variables have been changed to:"
 for var in "${_socks_cli_export_vars[@]}"; do
     echo "  $var=$(_expand "$var")"
 done
diff --git a/deactivate b/deactivate
index 3574d1f..b517f3e 100644
--- a/deactivate
+++ b/deactivate
@@ -1,19 +1,20 @@
 #!/bin/bash
 # The shebang is just a hint for editors.
 
-if [ "${_socks_cli}" = "1" ]; then
-    function _restore_var() {
-        if [ -n "$1" ]; then
-            local VAR="_socks_cli_saved_$1"
-            if _defined "${VAR}"; then
-                echo "export $1=\"\${${VAR}}\"; unset ${VAR}"
-            else
-                # Check if the variable really exists, for security.
-                # e.g. to avoid code injection or something else.
-                _defined "$1" && echo "unset $1"
-            fi
+function _restore_var() {
+    if [ -n "$1" ]; then
+        local VAR="_socks_cli_saved_$1"
+        if _defined "${VAR}"; then
+            echo "export $1=\"\${${VAR}}\"; unset ${VAR}"
+        else
+            # Check if the variable really exists, for security.
+            # e.g. to avoid code injection or something else.
+            _defined "$1" && echo "unset $1"
         fi
-    }
+    fi
+}
+
+if [ "${_socks_cli}" = "1" ]; then
 
     # Execute deactivation hooks
     for support in "${_socks_cli_supports[@]}"; do
@@ -24,7 +25,13 @@ if [ "${_socks_cli}" = "1" ]; then
     # Restore variables
     for var in "${_socks_cli_saved_vars[@]}"; do
         eval "$(_restore_var $var)"
+
+        # Rehash if PATH changed
+        if [ "$var" = "PATH" ]; then
+            hash -r
+        fi
     done
 
+    unset _socks_cli_shadow_inited
     unset _socks_cli
 fi
diff --git a/socksproxyenv.sample b/socksproxyenv.sample
index 1a59e33..61ddb8c 100644
--- a/socksproxyenv.sample
+++ b/socksproxyenv.sample
@@ -11,7 +11,7 @@
 #
 #    export SOCKS_PROXY="localhost:1080"
 #
-#  Alternatively, you can also do something complex, for example:
+#  Alternatively, you can also do something complicated, for example:
 #
 #    # Create a ssh tunnel on the fly
 #    PORT=1088
@@ -30,22 +30,22 @@ export SOCKS_PROXY=
 ###############################################################################
 
 # socks-cli has a few pluggable supports for various kinds of CLI commands.
-# You can disable some of them as you wish. For example, if you don't run git
-# command at all, you can just comment out the line "LOAD_SUPPORT git".
+# You can disable some of them as you wish. For example, if you don't want to
+# run the git command with the socks proxy, just remove "LOAD_SUPPORT git".
 
 LOAD_SUPPORT git
 LOAD_SUPPORT http
 LOAD_SUPPORT wget
 LOAD_SUPPORT mvn
+LOAD_SUPPORT ssh
 
 
 # Further more, you can extend socks-cli to support more CLI commands, by
 # exporting dedicated variables with the "EXPORT_ENV" DSL. socks-cli will take
 # care of these variables as well.
-# For example, to export FTP_PROXY:
-#   EXPORT_ENV FTP_PROXY "http://${SOCKS_CLI_HTTP_PROXY}"
 
-# EXPORT_ENV FOO "this is an example of EXPORT_ENV!"
+# Uncomment to export FTP_PROXY:
+#   EXPORT_ENV FTP_PROXY "http://${SOCKS_CLI_HTTP_PROXY}"
 
 
 # However, not all program has built-in proxy support. To settle these
diff --git a/supports/dsocks b/supports/dsocks
index a0bb1d0..a681f40 100644
--- a/supports/dsocks
+++ b/supports/dsocks
@@ -1,27 +1,8 @@
 #!/bin/bash
 # The shebang is just a hint for editors.
 
-function _socks_cli_clean_dsocks_tmp() {
-    local tmp_path="${SOCKS_CLI_DIR}/tmp/"
-    # First pass: remove command shadows if the shell has been dead.
-    find "${tmp_path}" -name "_shell_pid" -print0 |
-        while IFS= read -r -d '' file; do
-            pid="$(tr -d '\n\r' < "${file}")"
-            if [ "${pid}" = "$$" ]; then
-                # Ignore the current shell
-                continue
-            fi
-            if ! pgrep -F "${file}" > /dev/null; then
-                find "$(dirname "${file}")" -type f -delete
-            fi
-        done
-
-    # Second pass: remove empty directories recursively.
-    find "${tmp_path}" -empty -type d -delete
-}
-
 function socks_cli_activate_dsocks() {
-    _socks_cli_dsocks_lib="${SOCKS_CLI_DIR}/dsocks/libdsocks.so"
+    local _socks_cli_dsocks_lib="${SOCKS_CLI_DIR}/dsocks/libdsocks.so"
     "${SOCKS_CLI_DIR}/sh/bootstrap-dsocks.sh"
     if [ ! -f "${_socks_cli_dsocks_lib}" ]; then
         _logerror "Error: ${_socks_cli_dsocks_lib} isn't exist,"
@@ -29,34 +10,10 @@ function socks_cli_activate_dsocks() {
         return 1
     fi
 
-    _socks_cli_clean_dsocks_tmp
-
-    # Create a dedicated temporary folder for each shell process.
-    # Since the maximum pid number in unix-like OS is typically
-    # not greater than 999999, we can distribute the temporary
-    # folder into a two level structure. So it will not result in
-    # a huge number of subfolders in a single folder.
-    local shell_pid=$$
-    local level_one=$((shell_pid/10000))
-    local level_two=$((shell_pid%10000/100))
-    local name=$((shell_pid%100))
-    local tmp_path="${SOCKS_CLI_DIR}/tmp/${level_one}/${level_two}/${name}"
-    mkdir -p "${tmp_path}"
-    # Remove existing command shadows
-    [ -n "${tmp_path}" ] && find "${tmp_path}/" -type f -delete
-    echo "${shell_pid}" > "${tmp_path}/_shell_pid"
-
-    local preappend="${tmp_path}"
-    if [ -n "${PATH}" ]; then
-        preappend="${preappend}:"
-    fi
-    EXPORT_ENV PATH "${preappend}${PATH}"
     EXPORT_ENV DSOCKS_LIBRARY "${_socks_cli_dsocks_lib}"
     EXPORT_ENV DSOCKS_PROXY "${SOCKS_PROXY}"
 
     _socks_cli_socksify=1
-    _socks_cli_socksify_tmp_path="${tmp_path}"
-    _socks_cli_socksify_path_preappend="${preappend}"
 }
 
 function socks_cli_deactivate_dsocks() {
@@ -77,21 +34,5 @@ function SOCKSIFY() {
     if [ -z "$start_command" ]; then
         start_command="${hook_command}"
     fi
-    local cmd_shadow="${_socks_cli_socksify_tmp_path}/${hook_command}"
-    cat <<-EOF > "${cmd_shadow}"
-		#!/bin/bash
-		function _sed_escape() {
-		    echo "\$1" | sed -e 's/\\\\/\\\\\\\\/g; s/\\//\\\\\\//g; s/&/\\\\\\&/g'
-		}
-		function socks_cli_replace() {
-		    local search=\$(_sed_escape "\$2")
-		    local replace=\$(_sed_escape "\$3")
-		    echo "\$1" | sed "s/\${search}/\${replace}/g"
-		}
-		# remove the current folder from the PATH,
-		# to avoid recursively invoking this script.
-		export PATH="\$(socks_cli_replace "\$PATH" "${_socks_cli_socksify_path_preappend}" "")"
-		"${SOCKS_CLI_DIR}/dsocks/socksify" ${start_command} "\$@"
-	EOF
-    chmod +x "${cmd_shadow}"
+    SHADOW "${hook_command}" "${SOCKS_CLI_DIR}/dsocks/socksify" "${start_command}"
 }
diff --git a/supports/ssh b/supports/ssh
new file mode 100644
index 0000000..dd817eb
--- /dev/null
+++ b/supports/ssh
@@ -0,0 +1,4 @@
+#!/bin/bash
+# The shebang is just a hint for editors.
+
+SHADOW ssh "${SOCKS_CLI_DIR}/sh/socksified-ssh.sh"