feat(shell-proxy): support SOCKS protocol based on URI scheme (#10069)

1 files changed, 18 insertions, 3 deletions

diff --git a/plugins/shell-proxy/ssh-proxy.py b/plugins/shell-proxy/ssh-proxy.py
index 5efd5fd21..6773a77bc 100755
--- a/plugins/shell-proxy/ssh-proxy.py
+++ b/plugins/shell-proxy/ssh-proxy.py
@@ -2,15 +2,30 @@
 import os
 import subprocess
 import sys
-import urllib.parse
+from urllib.parse import urlparse
 
 proxy = next(os.environ[_] for _ in ("HTTP_PROXY", "HTTPS_PROXY") if _ in os.environ)
+
+parsed = urlparse(proxy)
+
+proxy_protocols = {
+    "http": "connect",
+    "https": "connect",
+    "socks": "5",
+    "socks5": "5",
+    "socks4": "4",
+    "socks4a": "4",
+}
+
+if parsed.scheme not in proxy_protocols:
+    raise TypeError('unsupported proxy protocol: "{}"'.format(parsed.scheme))
+
 argv = [
     "nc",
     "-X",
-    "connect",
+    proxy_protocols[parsed.scheme], # Supported protocols are 4 (SOCKS v4), 5 (SOCKS v5) and connect (HTTP proxy). Default SOCKS v5 is used.
     "-x",
-    urllib.parse.urlparse(proxy).netloc,  # proxy-host:proxy-port
+    parsed.netloc,  # proxy-host:proxy-port
     sys.argv[1],  # host
     sys.argv[2],  # port
 ]