juju/client/connection.py

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