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", outgoing)
  77         #log.debug("Recv: %s", result)
  78         if result and 'error' in result:
  79             raise RuntimeError(result)
  80         return result
  81
  82     async def clone(self):
  83         """Return a new Connection, connected to the same websocket endpoint
  84         as this one.
  85
  86         """
  87         return await Connection.connect(
  88             self.endpoint,
  89             self.uuid,
  90             self.username,
  91             self.password,
  92             self.cacert,
  93         )
  94
  95     async def controller(self):
  96         """Return a Connection to the controller at self.endpoint
  97
  98         """
  99         return await Connection.connect(
 100             self.endpoint,
 101             None,
 102             self.username,
 103             self.password,
 104             self.cacert,
 105         )
 106
 107     @classmethod
 108     async def connect(cls, endpoint, uuid, username, password, cacert=None):
 109         """Connect to the websocket.
 110
 111         If uuid is None, the connection will be to the controller. Otherwise it
 112         will be to the model.
 113
 114         """
 115         if uuid:
 116             url = "wss://{}/model/{}/api".format(endpoint, uuid)
 117         else:
 118             url = "wss://{}/api".format(endpoint)
 119         client = cls(endpoint, uuid, username, password, cacert)
 120         await client.open(url, cacert)
 121         server_info = await client.login(username, password)
 122         client.build_facades(server_info['facades'])
 123         log.info("Driver connected to juju %s", url)
 124
 125         return client
 126
 127     @classmethod
 128     async def connect_current(cls):
 129         """Connect to the currently active model.
 130
 131         """
 132         jujudata = JujuData()
 133         controller_name = jujudata.current_controller()
 134         controller = jujudata.controllers()[controller_name]
 135         endpoint = controller['api-endpoints'][0]
 136         cacert = controller.get('ca-cert')
 137         accounts = jujudata.accounts()[controller_name]
 138         username = accounts['user']
 139         password = accounts['password']
 140         models = jujudata.models()[controller_name]
 141         model_name = models['current-model']
 142         model_uuid = models['models'][model_name]['uuid']
 143
 144         return await cls.connect(
 145             endpoint, model_uuid, username, password, cacert)
 146
 147     @classmethod
 148     async def connect_model(cls, model):
 149         """Connect to a model by name.
 150
 151         :param str model: <controller>:<model>
 152
 153         """
 154         controller_name, model_name = model.split(':')
 155
 156         jujudata = JujuData()
 157         controller = jujudata.controllers()[controller_name]
 158         endpoint = controller['api-endpoints'][0]
 159         cacert = controller.get('ca-cert')
 160         accounts = jujudata.accounts()[controller_name]
 161         username = accounts['user']
 162         password = accounts['password']
 163         models = jujudata.models()[controller_name]
 164         model_uuid = models['models'][model_name]['uuid']
 165
 166         return await cls.connect(
 167             endpoint, model_uuid, username, password, cacert)
 168
 169     def build_facades(self, info):
 170         self.facades.clear()
 171         for facade in info:
 172             self.facades[facade['name']] = facade['versions'][-1]
 173
 174     async def login(self, username, password):
 175         if not username.startswith('user-'):
 176             username = 'user-{}'.format(username)
 177
 178         result = await self.rpc({
 179             "type": "Admin",
 180             "request": "Login",
 181             "version": 3,
 182             "params": {
 183                 "auth-tag": username,
 184                 "credentials": password,
 185                 "nonce": "".join(random.sample(string.printable, 12)),
 186             }})
 187         return result['response']
 188
 189
 190 class JujuData:
 191     def __init__(self):
 192         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 193         self.path = os.path.abspath(os.path.expanduser(self.path))
 194
 195     def current_controller(self):
 196         cmd = shlex.split('juju show-controller --format yaml')
 197         output = subprocess.check_output(cmd)
 198         output = yaml.safe_load(output)
 199         return list(output.keys())[0]
 200
 201     def controllers(self):
 202         return self._load_yaml('controllers.yaml', 'controllers')
 203
 204     def models(self):
 205         return self._load_yaml('models.yaml', 'controllers')
 206
 207     def accounts(self):
 208         return self._load_yaml('accounts.yaml', 'controllers')
 209
 210     def _load_yaml(self, filename, key):
 211         filepath = os.path.join(self.path, filename)
 212         with io.open(filepath, 'rt') as f:
 213             return yaml.safe_load(f)[key]