Add support for SSH protocol in base_url

Signed-off-by: Joffrey F <joffrey@docker.com>

4 files changed, 145 insertions, 5 deletions

diff --git a/docker/api/client.py b/docker/api/client.py
index 91da1c8..197846d 100644
--- a/docker/api/client.py
+++ b/docker/api/client.py
@@ -39,6 +39,11 @@ try:
 except ImportError:
     pass
 
+try:
+    from ..transport import SSHAdapter
+except ImportError:
+    pass
+
 
 class APIClient(
         requests.Session,
@@ -141,6 +146,18 @@ class APIClient(
                 )
             self.mount('http+docker://', self._custom_adapter)
             self.base_url = 'http+docker://localnpipe'
+        elif base_url.startswith('ssh://'):
+            try:
+                self._custom_adapter = SSHAdapter(
+                    base_url, timeout, pool_connections=num_pools
+                )
+            except NameError:
+                raise DockerException(
+                    'Install paramiko package to enable ssh:// support'
+                )
+            self.mount('http+docker://ssh', self._custom_adapter)
+            self._unmount('http://', 'https://')
+            self.base_url = 'http+docker://ssh'
         else:
             # Use SSLAdapter for the ability to specify SSL version
             if isinstance(tls, TLSConfig):
@@ -279,6 +296,8 @@ class APIClient(
         self._raise_for_status(response)
         if self.base_url == "http+docker://localnpipe":
             sock = response.raw._fp.fp.raw.sock
+        elif self.base_url.startswith('http+docker://ssh'):
+            sock = response.raw._fp.fp.channel
         elif six.PY3:
             sock = response.raw._fp.fp.raw
             if self.base_url.startswith("https://"):
diff --git a/docker/transport/__init__.py b/docker/transport/__init__.py
index abbee18..d2cf2a7 100644
--- a/docker/transport/__init__.py
+++ b/docker/transport/__init__.py
@@ -6,3 +6,8 @@ try:
     from .npipesocket import NpipeSocket
 except ImportError:
     pass
+
+try:
+    from .sshconn import SSHAdapter
+except ImportError:
+    pass
diff --git a/docker/transport/sshconn.py b/docker/transport/sshconn.py
new file mode 100644
index 0000000..6c9c119
--- /dev/null
+++ b/docker/transport/sshconn.py
@@ -0,0 +1,110 @@
+import urllib.parse
+
+import paramiko
+import requests.adapters
+import six
+
+
+from .. import constants
+
+if six.PY3:
+    import http.client as httplib
+else:
+    import httplib
+
+try:
+    import requests.packages.urllib3 as urllib3
+except ImportError:
+    import urllib3
+
+RecentlyUsedContainer = urllib3._collections.RecentlyUsedContainer
+
+
+class SSHConnection(httplib.HTTPConnection, object):
+    def __init__(self, ssh_transport, timeout=60):
+        super(SSHConnection, self).__init__(
+            'localhost', timeout=timeout
+        )
+        self.ssh_transport = ssh_transport
+        self.timeout = timeout
+
+    def connect(self):
+        sock = self.ssh_transport.open_session()
+        sock.settimeout(self.timeout)
+        sock.exec_command('docker system dial-stdio')
+        self.sock = sock
+
+
+class SSHConnectionPool(urllib3.connectionpool.HTTPConnectionPool):
+    scheme = 'ssh'
+
+    def __init__(self, ssh_client, timeout=60, maxsize=10):
+        super(SSHConnectionPool, self).__init__(
+            'localhost', timeout=timeout, maxsize=maxsize
+        )
+        self.ssh_transport = ssh_client.get_transport()
+        self.timeout = timeout
+
+    def _new_conn(self):
+        return SSHConnection(self.ssh_transport, self.timeout)
+
+    # When re-using connections, urllib3 calls fileno() on our
+    # SSH channel instance, quickly overloading our fd limit. To avoid this,
+    # we override _get_conn
+    def _get_conn(self, timeout):
+        conn = None
+        try:
+            conn = self.pool.get(block=self.block, timeout=timeout)
+
+        except AttributeError:  # self.pool is None
+            raise urllib3.exceptions.ClosedPoolError(self, "Pool is closed.")
+
+        except six.moves.queue.Empty:
+            if self.block:
+                raise urllib3.exceptions.EmptyPoolError(
+                    self,
+                    "Pool reached maximum size and no more "
+                    "connections are allowed."
+                )
+            pass  # Oh well, we'll create a new connection then
+
+        return conn or self._new_conn()
+
+
+class SSHAdapter(requests.adapters.HTTPAdapter):
+
+    __attrs__ = requests.adapters.HTTPAdapter.__attrs__ + [
+        'pools', 'timeout', 'ssh_client',
+    ]
+
+    def __init__(self, base_url, timeout=60,
+                 pool_connections=constants.DEFAULT_NUM_POOLS):
+        self.ssh_client = paramiko.SSHClient()
+        self.ssh_client.load_system_host_keys()
+
+        parsed = urllib.parse.urlparse(base_url)
+        self.ssh_client.connect(
+            parsed.hostname, parsed.port, parsed.username,
+        )
+        self.timeout = timeout
+        self.pools = RecentlyUsedContainer(
+            pool_connections, dispose_func=lambda p: p.close()
+        )
+        super(SSHAdapter, self).__init__()
+
+    def get_connection(self, url, proxies=None):
+        with self.pools.lock:
+            pool = self.pools.get(url)
+            if pool:
+                return pool
+
+            pool = SSHConnectionPool(
+                self.ssh_client, self.timeout
+            )
+            self.pools[url] = pool
+
+        return pool
+
+    def close(self):
+        self.pools.clear()
+        self.ssh_client.close()
diff --git a/docker/utils/utils.py b/docker/utils/utils.py
index fe3b9a5..f8f7123 100644
--- a/docker/utils/utils.py
+++ b/docker/utils/utils.py
@@ -250,6 +250,9 @@ def parse_host(addr, is_win32=False, tls=False):
         addr = addr[8:]
     elif addr.startswith('fd://'):
         raise errors.DockerException("fd protocol is not implemented")
+    elif addr.startswith('ssh://'):
+        proto = 'ssh'
+        addr = addr[6:]
     else:
         if "://" in addr:
             raise errors.DockerException(
@@ -257,17 +260,20 @@ def parse_host(addr, is_win32=False, tls=False):
             )
         proto = "https" if tls else "http"
 
-    if proto in ("http", "https"):
+    if proto in ("http", "https", "ssh"):
         address_parts = addr.split('/', 1)
         host = address_parts[0]
         if len(address_parts) == 2:
             path = '/' + address_parts[1]
         host, port = splitnport(host)
 
-        if port is None:
-            raise errors.DockerException(
-                "Invalid port: {0}".format(addr)
-            )
+        if port is None or port < 0:
+            if proto == 'ssh':
+                port = 22
+            else:
+                raise errors.DockerException(
+                    "Invalid port: {0}".format(addr)
+                )
 
         if not host:
             host = DEFAULT_HTTP_HOST