Adds support for juju client proxies

juju/client/connection.py

@@ -231,6 +231,7 @@ async def connect(
             retries=3,
             retry_backoff=10,
             specified_facades=None,
+            proxy=None,
     ):
         """Connect to the websocket.
 
@@ -308,6 +309,10 @@ async def connect(
             max_frame_size = self.MAX_FRAME_SIZE
         self.max_frame_size = max_frame_size
 
+        self.proxy = proxy
+        if self.proxy is not None:
+            self.proxy.connect()
+
         _endpoints = [(endpoint, cacert)] if isinstance(endpoint, str) else [(e, cacert) for e in endpoint]
         for _ep in _endpoints:
             try:
@@ -348,12 +353,23 @@ async def _open(self, endpoint, cacert):
         else:
             url = "wss://{}/api".format(endpoint)
 
+        # We need to establish a server_hostname here for TLS sni if we are
+        # connecting through a proxy as the Juju controller certificates will
+        # not be covering the proxy
+        sock = None
+        server_hostname = None
+        if self.proxy is not None:
+            sock = self.proxy.socket()
+            server_hostname = "juju-app"
+
         return (await websockets.connect(
             url,
             ssl=self._get_ssl(cacert),
             loop=self.loop,
             max_size=self.max_frame_size,
-        ), url, endpoint, cacert)
+            server_hostname=server_hostname,
+            sock=sock,
+        )), url, endpoint, cacert
 
     async def close(self):
         if not self.ws:
@@ -364,6 +380,9 @@ async def close(self):
         await self.ws.close()
         self.ws = None
 
+        if self.proxy is not None:
+            self.proxy.close()
+
     async def _recv(self, request_id):
         if not self.is_open:
             raise websockets.exceptions.ConnectionClosed(0, 'websocket closed')
@@ -551,11 +570,9 @@ async def clone(self):
         return await Connection.connect(**self.connect_params())
 
     def connect_params(self):
-        """Return a tuple of parameters suitable for passing to
+        """Return a dict of parameters suitable for passing to
         Connection.connect that can be used to make a new connection
-        to the same controller (and model if specified. The first
-        element in the returned tuple holds the endpoint argument;
-        the other holds a dict of the keyword args.
+        to the same controller (and model if specified).
         """
         return {
             'endpoint': self.endpoint,
@@ -566,6 +583,7 @@ def connect_params(self):
             'bakery_client': self.bakery_client,
             'loop': self.loop,
             'max_frame_size': self.max_frame_size,
+            'proxy': self.proxy,
         }
 
     async def controller(self):

juju/client/connector.py

@@ -6,6 +6,7 @@
 from juju.client.connection import Connection
 from juju.client.gocookies import GoCookieJar, go_to_py_cookie
 from juju.client.jujudata import FileJujuData
+from juju.client.proxy.factory import proxy_from_config
 from juju.errors import JujuConnectionError, JujuError
 
 log = logging.getLogger('connector')
@@ -55,6 +56,7 @@ async def connect(self, **kwargs):
 
         kwargs are passed through to Connection.connect()
         """
+
         kwargs.setdefault('loop', self.loop)
         kwargs.setdefault('max_frame_size', self.max_frame_size)
         kwargs.setdefault('bakery_client', self.bakery_client)
@@ -89,6 +91,8 @@ async def connect_controller(self, controller_name=None, specified_facades=None)
         endpoints = controller['api-endpoints']
         accounts = self.jujudata.accounts().get(controller_name, {})
 
+        proxy = proxy_from_config(controller.get('proxy-config', None))
+
         await self.connect(
             endpoint=endpoints,
             uuid=None,
@@ -97,6 +101,7 @@ async def connect_controller(self, controller_name=None, specified_facades=None)
             cacert=controller.get('ca-cert'),
             bakery_client=self.bakery_client_for_controller(controller_name),
             specified_facades=specified_facades,
+            proxy=proxy,
         )
         self.controller_name = controller_name
         self.controller_uuid = controller["uuid"]
@@ -121,9 +126,12 @@ async def connect_model(self, model_name=None):
         account = self.jujudata.accounts().get(controller_name, {})
         models = self.jujudata.models().get(controller_name, {}).get('models',
                                                                      {})
+
         if model_name not in models:
             raise JujuConnectionError('Model not found: {}'.format(model_name))
 
+        proxy = proxy_from_config(controller.get('proxy-config', None))
+
         # TODO if there's no record for the required model name, connect
         # to the controller to find out the model's uuid, then connect
         # to that. This will let connect_model work with models that
@@ -137,6 +145,7 @@ async def connect_model(self, model_name=None):
             password=account.get('password'),
             cacert=controller.get('ca-cert'),
             bakery_client=self.bakery_client_for_controller(controller_name),
+            proxy=proxy,
         )
         self.controller_name = controller_name
         self.model_name = controller_name + ':' + model_name

juju/client/proxy/factory.py

@@ -0,0 +1,26 @@
+from juju.client.proxy.kubernetes.proxy import KubernetesProxy
+
+
+def proxy_from_config(conf):
+    if conf is None:
+        return None
+
+    if 'type' not in conf:
+        return None
+
+    proxy_type = conf['type']
+    if proxy_type != 'kubernetes-port-forward':
+        raise ValueError('unknown proxy type %s' % proxy_type)
+
+    return _construct_kube_proxy(conf['config'])
+
+
+def _construct_kube_proxy(config):
+    return KubernetesProxy(
+        config.get('api-host', ''),
+        config.get('namespace', ''),
+        config.get('remote-port', ''),
+        config.get('service', ''),
+        config.get('service-account-token', ''),
+        config.get('ca-cert', None),
+    )

juju/client/proxy/kubernetes/proxy.py

@@ -0,0 +1,71 @@
+import tempfile
+
+from juju.client.proxy.proxy import Proxy, ProxyNotConnectedError
+from kubernetes import client
+from kubernetes.stream import portforward
+
+
+class KubernetesProxy(Proxy):
+    def __init__(
+            self,
+            api_host,
+            namespace,
+            remote_port,
+            service,
+            service_account_token,
+            ca_cert=None,
+    ):
+        config = client.Configuration()
+        config.host = api_host
+        config.ssl_ca_cert = ca_cert
+        config.api_key = {"authorization": "Bearer " + service_account_token}
+
+        self.namespace = namespace
+        self.remote_port = remote_port
+        self.service = service
+
+        try:
+            self.remote_port = int(remote_port)
+        except ValueError:
+            raise ValueError("Invalid port number: {}".format(remote_port))
+
+        if ca_cert:
+            self.temp_ca_file = tempfile.NamedTemporaryFile()
+            self.temp_ca_file.write(bytes(ca_cert, 'utf-8'))
+            self.temp_ca_file.flush()
+            config.ssl_ca_cert = self.temp_ca_file.name
+
+        self.api_client = client.ApiClient(config)
+
+    def connect(self):
+        corev1 = client.CoreV1Api(self.api_client)
+        service = corev1.read_namespaced_service(self.service, self.namespace)
+
+        label_selector = ','.join(k + '=' + v for k, v in service.spec.selector.items())
+
+        pods = corev1.list_namespaced_pod(
+            namespace=self.namespace,
+            label_selector=label_selector,
+        )
+
+        self.port_forwarder = portforward(
+            corev1.connect_get_namespaced_pod_portforward,
+            pods.items[0].metadata.name,
+            self.namespace,
+            ports=str(self.remote_port),
+        )
+
+    def __del__(self):
+        self.close()
+
+    def close(self):
+        try:
+            self.port_forwarder.close()
+            self.temp_ca_file.close()
+        except AttributeError:
+            pass
+
+    def socket(self):
+        if self.port_forwarder is not None:
+            return self.port_forwarder.socket(self.remote_port)._socket
+        raise ProxyNotConnectedError()

juju/client/proxy/proxy.py

@@ -0,0 +1,23 @@
+from abc import abstractmethod
+
+
+class ProxyNotConnectedError(Exception):
+    pass
+
+
+class Proxy():
+    """
+    Abstract class to represent a generic controller connection proxy
+    """
+
+    @abstractmethod
+    def connect(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def close(self):
+        raise NotImplementedError()
+
+    @abstractmethod
+    def socket(self):
+        raise NotImplementedError()

setup.py

@@ -40,7 +40,8 @@
         'paramiko>=2.4.0,<3.0.0',
         'pyasn1>=0.4.4',
         'toposort>=1.5,<2',
-        'typing_inspect>=0.6.0'
+        'typing_inspect>=0.6.0',
+        'kubernetes>=12.0.1',
     ],
     include_package_data=True,
     maintainer='Juju Ecosystem Engineering',

tests/unit/test_proxy.py

@@ -0,0 +1,49 @@
+import unittest
+
+from juju.client.proxy.factory import proxy_from_config
+from juju.client.proxy.kubernetes.proxy import KubernetesProxy
+
+
+class TestJujuDataFactory(unittest.TestCase):
+
+    def test_proxy_from_config_unknown_type(self):
+        """
+        Test that a unknown proxy type results in a UnknownProxyTypeError
+        exception
+        """
+        self.assertRaises(ValueError, proxy_from_config, {
+            "config": {},
+            "type": "does-not-exists",
+        })
+
+    def test_proxy_from_config_missing_type(self):
+        """
+        Test that a nil proxy type returns None
+        """
+        self.assertIsNone(proxy_from_config({
+            "config": {},
+        }))
+
+    def test_proxy_from_config_non_arg(self):
+        """
+        Tests that providing an empty proxy config results in a None proxy
+        """
+        self.assertIsNone(proxy_from_config(None))
+
+    def test_proxy_from_config_kubernetes(self):
+        """
+        Tests that a Kubernetes proxy is correctly created from config
+        """
+        proxy = proxy_from_config({
+            "type": "kubernetes-port-forward",
+            "config": {
+                "api-host": "https://localhost:8456",
+                "namespace": "controller-python-test",
+                "remote-port": "1234",
+                "service": "controller",
+                "service-account-token": "==AA",
+                "ca-cert": "==AA",
+            },
+        })
+
+        self.assertIs(type(proxy), KubernetesProxy)

tests/unit/test_proxy_kubernetes.py

@@ -0,0 +1,15 @@
+import unittest
+from juju.client.proxy.kubernetes.proxy import KubernetesProxy
+
+
+class TestKubernetesProxy(unittest.TestCase):
+    def test_remote_port_error(self):
+        self.assertRaises(
+            ValueError,
+            KubernetesProxy,
+            api_host="https://localhost:1234",
+            namespace="controller",
+            remote_port="not-a-integer-port",
+            service="service",
+            service_account_token="==AA",
+        )