juju/client/connection.py

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