path: root/docker/transport/sshconn.py

import paramiko
import requests.adapters
import six
import logging
import os
import signal
import socket
import subprocess

from docker.transport.basehttpadapter import BaseHTTPAdapter
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 SSHSocket(socket.socket):
    def __init__(self, host):
        super(SSHSocket, self).__init__(
            socket.AF_INET, socket.SOCK_STREAM)
        self.host = host
        self.port = None
        self.user = None
        if ':' in self.host:
            self.host, self.port = self.host.split(':')
        if '@' in self.host:
            self.user, self.host = self.host.split('@')

        self.proc = None

    def connect(self, **kwargs):
        args = ['ssh']
        if self.user:
            args = args + ['-l', self.user]

        if self.port:
            args = args + ['-p', self.port]

        args = args + ['--', self.host, 'docker system dial-stdio']

        preexec_func = None
        if not constants.IS_WINDOWS_PLATFORM:
            def f():
                signal.signal(signal.SIGINT, signal.SIG_IGN)
            preexec_func = f

        env = dict(os.environ) 

        # drop LD_LIBRARY_PATH and SSL_CERT_FILE
        env.pop('LD_LIBRARY_PATH', None)
        env.pop('SSL_CERT_FILE', None)

        self.proc = subprocess.Popen(
            ' '.join(args),
            env=env,
            shell=True,
            stdout=subprocess.PIPE,
            stdin=subprocess.PIPE,
            preexec_fn=preexec_func)

    def _write(self, data):
        if not self.proc or self.proc.stdin.closed:
            raise Exception('SSH subprocess not initiated.'
                            'connect() must be called first.')
        written = self.proc.stdin.write(data)
        self.proc.stdin.flush()
        return written

    def sendall(self, data):
        self._write(data)

    def send(self, data):
        return self._write(data)

    def recv(self, n):
        if not self.proc:
            raise Exception('SSH subprocess not initiated.'
                            'connect() must be called first.')
        return self.proc.stdout.read(n)

    def makefile(self, mode):
        if not self.proc:
            self.connect()
        if six.PY3:
            self.proc.stdout.channel = self

        return self.proc.stdout

    def close(self):
        if not self.proc or self.proc.stdin.closed:
            return
        self.proc.stdin.write(b'\n\n')
        self.proc.stdin.flush()
        self.proc.terminate()


class SSHConnection(httplib.HTTPConnection, object):
    def __init__(self, ssh_transport=None, timeout=60, host=None):
        super(SSHConnection, self).__init__(
            'localhost', timeout=timeout
        )
        self.ssh_transport = ssh_transport
        self.timeout = timeout
        self.ssh_host = host

    def connect(self):
        if self.ssh_transport:
            sock = self.ssh_transport.open_session()
            sock.settimeout(self.timeout)
            sock.exec_command('docker system dial-stdio')
        else:
            sock = SSHSocket(self.ssh_host)
            sock.settimeout(self.timeout)
            sock.connect()

        self.sock = sock


class SSHConnectionPool(urllib3.connectionpool.HTTPConnectionPool):
    scheme = 'ssh'

    def __init__(self, ssh_client=None, timeout=60, maxsize=10, host=None):
        super(SSHConnectionPool, self).__init__(
            'localhost', timeout=timeout, maxsize=maxsize
        )
        self.ssh_transport = None
        self.timeout = timeout
        if ssh_client:
            self.ssh_transport = ssh_client.get_transport()
        self.ssh_host = host

    def _new_conn(self):
        return SSHConnection(self.ssh_transport, self.timeout, self.ssh_host)

    # 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 SSHHTTPAdapter(BaseHTTPAdapter):

    __attrs__ = requests.adapters.HTTPAdapter.__attrs__ + [
        'pools', 'timeout', 'ssh_client', 'ssh_params', 'max_pool_size'
    ]

    def __init__(self, base_url, timeout=60,
                 pool_connections=constants.DEFAULT_NUM_POOLS,
                 max_pool_size=constants.DEFAULT_MAX_POOL_SIZE,
                 shell_out=False):
        self.ssh_client = None
        if not shell_out:
            self._create_paramiko_client(base_url)
            self._connect()

        self.ssh_host = base_url
        if base_url.startswith('ssh://'):
            self.ssh_host = base_url[len('ssh://'):]

        self.timeout = timeout
        self.max_pool_size = max_pool_size
        self.pools = RecentlyUsedContainer(
            pool_connections, dispose_func=lambda p: p.close()
        )
        super(SSHHTTPAdapter, self).__init__()

    def _create_paramiko_client(self, base_url):
        logging.getLogger("paramiko").setLevel(logging.WARNING)
        self.ssh_client = paramiko.SSHClient()
        base_url = six.moves.urllib_parse.urlparse(base_url)
        self.ssh_params = {
            "hostname": base_url.hostname,
            "port": base_url.port,
            "username": base_url.username
            }
        ssh_config_file = os.path.expanduser("~/.ssh/config")
        if os.path.exists(ssh_config_file):
            conf = paramiko.SSHConfig()
            with open(ssh_config_file) as f:
                conf.parse(f)
            host_config = conf.lookup(base_url.hostname)
            self.ssh_conf = host_config
            if 'proxycommand' in host_config:
                self.ssh_params["sock"] = paramiko.ProxyCommand(
                    self.ssh_conf['proxycommand']
                )
            if 'hostname' in host_config:
                self.ssh_params['hostname'] = host_config['hostname']
            if base_url.port is None and 'port' in host_config:
                self.ssh_params['port'] = self.ssh_conf['port']
            if base_url.username is None and 'user' in host_config:
                self.ssh_params['username'] = self.ssh_conf['user']

        self.ssh_client.load_system_host_keys()
        self.ssh_client.set_missing_host_key_policy(paramiko.WarningPolicy())

    def _connect(self):
        if self.ssh_client:
            self.ssh_client.connect(**self.ssh_params)

    def get_connection(self, url, proxies=None):
        if not self.ssh_client:
            return SSHConnectionPool(
                ssh_client=self.ssh_client,
                timeout=self.timeout,
                maxsize=self.max_pool_size,
                host=self.ssh_host
            )
        with self.pools.lock:
            pool = self.pools.get(url)
            if pool:
                return pool

            # Connection is closed try a reconnect
            if self.ssh_client and not self.ssh_client.get_transport():
                self._connect()

            pool = SSHConnectionPool(
                ssh_client=self.ssh_client,
                timeout=self.timeout,
                maxsize=self.max_pool_size,
                host=self.ssh_host
            )
            self.pools[url] = pool

        return pool

    def close(self):
        super(SSHHTTPAdapter, self).close()
        if self.ssh_client:
            self.ssh_client.close()