juju/client/connection.py

   1 import io
   2 import json
   3 import logging
   4 import os
   5 import random
   6 import shlex
   7 import ssl
   8 import string
   9 import subprocess
  10 import websockets
  11
  12 import yaml
  13
  14 log = logging.getLogger("websocket")
  15
  16
  17 class Connection:
  18     """
  19     Usage::
  20
  21         # Connect to an arbitrary api server
  22         client = await Connection.connect(
  23             api_endpoint, model_uuid, username, password, cacert)
  24
  25         # Connect using a controller/model name
  26         client = await Connection.connect_model('local.local:default')
  27
  28         # Connect to the currently active model
  29         client = await Connection.connect_current()
  30
  31     """
  32     def __init__(self, endpoint, uuid, username, password, cacert=None):
  33         self.endpoint = endpoint
  34         self.uuid = uuid
  35         self.username = username
  36         self.password = password
  37         self.cacert = cacert
  38
  39         self.__request_id__ = 0
  40         self.addr = None
  41         self.ws = None
  42         self.facades = {}
  43
  44     @property
  45     def is_open(self):
  46         if self.ws:
  47             return self.ws.open
  48         return False
  49
  50     def _get_ssl(self, cert):
  51         return ssl.create_default_context(
  52             purpose=ssl.Purpose.CLIENT_AUTH, cadata=cert)
  53
  54     async def open(self, addr, cert=None):
  55         kw = dict()
  56         if cert:
  57             kw['ssl'] = self._get_ssl(cert)
  58         self.addr = addr
  59         self.ws = await websockets.connect(addr, **kw)
  60         return self
  61
  62     async def close(self):
  63         await self.ws.close()
  64
  65     async def recv(self):
  66         result = await self.ws.recv()
  67         if result is not None:
  68             result = json.loads(result)
  69         return result
  70
  71     async def rpc(self, msg, encoder=None):
  72         self.__request_id__ += 1
  73         msg['request-id'] = self.__request_id__
  74         if'params' not in msg:
  75             msg['params'] = {}
  76         if "version" not in msg:
  77             msg['version'] = self.facades[msg['type']]
  78         outgoing = json.dumps(msg, indent=2, cls=encoder)
  79         await self.ws.send(outgoing)
  80         result = await self.recv()
  81         #log.debug("Send: %s", outgoing)
  82         #log.debug("Recv: %s", result)
  83         if result and 'error' in result:
  84             raise RuntimeError(result)
  85         return result
  86
  87     async def clone(self):
  88         """Return a new Connection, connected to the same websocket endpoint
  89         as this one.
  90
  91         """
  92         return await Connection.connect(
  93             self.endpoint,
  94             self.uuid,
  95             self.username,
  96             self.password,
  97             self.cacert,
  98         )
  99
 100     async def controller(self):
 101         """Return a Connection to the controller at self.endpoint
 102
 103         """
 104         return await Connection.connect(
 105             self.endpoint,
 106             None,
 107             self.username,
 108             self.password,
 109             self.cacert,
 110         )
 111
 112     @classmethod
 113     async def connect(cls, endpoint, uuid, username, password, cacert=None):
 114         """Connect to the websocket.
 115
 116         If uuid is None, the connection will be to the controller. Otherwise it
 117         will be to the model.
 118
 119         """
 120         if uuid:
 121             url = "wss://{}/model/{}/api".format(endpoint, uuid)
 122         else:
 123             url = "wss://{}/api".format(endpoint)
 124         client = cls(endpoint, uuid, username, password, cacert)
 125         await client.open(url, cacert)
 126         server_info = await client.login(username, password)
 127         client.build_facades(server_info['facades'])
 128         log.info("Driver connected to juju %s", url)
 129
 130         return client
 131
 132     @classmethod
 133     async def connect_current(cls):
 134         """Connect to the currently active model.
 135
 136         """
 137         jujudata = JujuData()
 138         controller_name = jujudata.current_controller()
 139         controller = jujudata.controllers()[controller_name]
 140         endpoint = controller['api-endpoints'][0]
 141         cacert = controller.get('ca-cert')
 142         accounts = jujudata.accounts()[controller_name]
 143         username = accounts['user']
 144         password = accounts['password']
 145         models = jujudata.models()[controller_name]
 146         model_name = models['current-model']
 147         model_uuid = models['models'][model_name]['uuid']
 148
 149         return await cls.connect(
 150             endpoint, model_uuid, username, password, cacert)
 151
 152     @classmethod
 153     async def connect_model(cls, model):
 154         """Connect to a model by name.
 155
 156         :param str model: <controller>:<model>
 157
 158         """
 159         controller_name, model_name = model.split(':')
 160
 161         jujudata = JujuData()
 162         controller = jujudata.controllers()[controller_name]
 163         endpoint = controller['api-endpoints'][0]
 164         cacert = controller.get('ca-cert')
 165         accounts = jujudata.accounts()[controller_name]
 166         username = accounts['user']
 167         password = accounts['password']
 168         models = jujudata.models()[controller_name]
 169         model_uuid = models['models'][model_name]['uuid']
 170
 171         return await cls.connect(
 172             endpoint, model_uuid, username, password, cacert)
 173
 174     def build_facades(self, info):
 175         self.facades.clear()
 176         for facade in info:
 177             self.facades[facade['name']] = facade['versions'][-1]
 178
 179     async def login(self, username, password):
 180         if not username.startswith('user-'):
 181             username = 'user-{}'.format(username)
 182
 183         result = await self.rpc({
 184             "type": "Admin",
 185             "request": "Login",
 186             "version": 3,
 187             "params": {
 188                 "auth-tag": username,
 189                 "credentials": password,
 190                 "nonce": "".join(random.sample(string.printable, 12)),
 191             }})
 192         return result['response']
 193
 194
 195 class JujuData:
 196     def __init__(self):
 197         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 198         self.path = os.path.abspath(os.path.expanduser(self.path))
 199
 200     def current_controller(self):
 201         cmd = shlex.split('juju show-controller --format yaml')
 202         output = subprocess.check_output(cmd)
 203         output = yaml.safe_load(output)
 204         return list(output.keys())[0]
 205
 206     def controllers(self):
 207         return self._load_yaml('controllers.yaml', 'controllers')
 208
 209     def models(self):
 210         return self._load_yaml('models.yaml', 'controllers')
 211
 212     def accounts(self):
 213         return self._load_yaml('accounts.yaml', 'controllers')
 214
 215     def _load_yaml(self, filename, key):
 216         filepath = os.path.join(self.path, filename)
 217         with io.open(filepath, 'rt') as f:
 218             return yaml.safe_load(f)[key]