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     @classmethod
  95     async def connect(cls, endpoint, uuid, username, password, cacert=None):
  96         url = "wss://{}/model/{}/api".format(endpoint, uuid)
  97         client = cls(endpoint, uuid, username, password, cacert)
  98         await client.open(url, cacert)
  99         server_info = await client.login(username, password)
 100         client.build_facades(server_info['facades'])
 101         log.info("Driver connected to juju %s", endpoint)
 102
 103         return client
 104
 105     @classmethod
 106     async def connect_current(cls):
 107         """Connect to the currently active model.
 108
 109         """
 110         jujudata = JujuData()
 111         controller_name = jujudata.current_controller()
 112         controller = jujudata.controllers()[controller_name]
 113         endpoint = controller['api-endpoints'][0]
 114         cacert = controller.get('ca-cert')
 115         accounts = jujudata.accounts()[controller_name]
 116         username = accounts['current-account']
 117         password = accounts['accounts'][username]['password']
 118         models = jujudata.models()[controller_name]['accounts'][username]
 119         model_name = models['current-model']
 120         model_uuid = models['models'][model_name]['uuid']
 121
 122         return await cls.connect(
 123             endpoint, model_uuid, username, password, cacert)
 124
 125     @classmethod
 126     async def connect_model(cls, model):
 127         """Connect to a model by name.
 128
 129         :param str model: <controller>:<model>
 130
 131         """
 132         controller_name, model_name = model.split(':')
 133
 134         jujudata = JujuData()
 135         controller = jujudata.controllers()[controller_name]
 136         endpoint = controller['api-endpoints'][0]
 137         cacert = controller.get('ca-cert')
 138         accounts = jujudata.accounts()[controller_name]
 139         username = accounts['current-account']
 140         password = accounts['accounts'][username]['password']
 141         models = jujudata.models()[controller_name]['accounts'][username]
 142         model_uuid = models['models'][model_name]['uuid']
 143
 144         return await cls.connect(
 145             endpoint, model_uuid, username, password, cacert)
 146
 147     def build_facades(self, info):
 148         self.facades.clear()
 149         for facade in info:
 150             self.facades[facade['name']] = facade['versions'][-1]
 151
 152     async def login(self, username, password):
 153         if not username.startswith('user-'):
 154             username = 'user-{}'.format(username)
 155
 156         result = await self.rpc({
 157             "type": "Admin",
 158             "request": "Login",
 159             "version": 3,
 160             "params": {
 161                 "auth-tag": username,
 162                 "credentials": password,
 163                 "nonce": "".join(random.sample(string.printable, 12)),
 164             }})
 165         return result['response']
 166
 167
 168 class JujuData:
 169     def __init__(self):
 170         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 171         self.path = os.path.abspath(os.path.expanduser(self.path))
 172
 173     def current_controller(self):
 174         cmd = shlex.split('juju show-controller --format yaml')
 175         output = subprocess.check_output(cmd)
 176         output = yaml.safe_load(output)
 177         return list(output.keys())[0]
 178
 179     def controllers(self):
 180         return self._load_yaml('controllers.yaml', 'controllers')
 181
 182     def models(self):
 183         return self._load_yaml('models.yaml', 'controllers')
 184
 185     def accounts(self):
 186         return self._load_yaml('accounts.yaml', 'controllers')
 187
 188     def _load_yaml(self, filename, key):
 189         filepath = os.path.join(self.path, filename)
 190         with io.open(filepath, 'rt') as f:
 191             return yaml.safe_load(f)[key]