juju/client/connection.py

   1 import base64
   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 from juju.errors import JujuAPIError
  16
  17 log = logging.getLogger("websocket")
  18
  19
  20 class Connection:
  21     """
  22     Usage::
  23
  24         # Connect to an arbitrary api server
  25         client = await Connection.connect(
  26             api_endpoint, model_uuid, username, password, cacert)
  27
  28         # Connect using a controller/model name
  29         client = await Connection.connect_model('local.local:default')
  30
  31         # Connect to the currently active model
  32         client = await Connection.connect_current()
  33
  34     """
  35     def __init__(
  36             self, endpoint, uuid, username, password, cacert=None,
  37             macaroons=None):
  38         self.endpoint = endpoint
  39         self.uuid = uuid
  40         self.username = username
  41         self.password = password
  42         self.macaroons = macaroons
  43         self.cacert = cacert
  44
  45         self.__request_id__ = 0
  46         self.addr = None
  47         self.ws = None
  48         self.facades = {}
  49
  50     @property
  51     def is_open(self):
  52         if self.ws:
  53             return self.ws.open
  54         return False
  55
  56     def _get_ssl(self, cert=None):
  57         return ssl.create_default_context(
  58             purpose=ssl.Purpose.CLIENT_AUTH, cadata=cert)
  59
  60     async def open(self):
  61         if self.uuid:
  62             url = "wss://{}/model/{}/api".format(self.endpoint, self.uuid)
  63         else:
  64             url = "wss://{}/api".format(self.endpoint)
  65
  66         kw = dict()
  67         kw['ssl'] = self._get_ssl(self.cacert)
  68         self.addr = url
  69         self.ws = await websockets.connect(url, **kw)
  70         log.info("Driver connected to juju %s", url)
  71         return self
  72
  73     async def close(self):
  74         await self.ws.close()
  75
  76     async def recv(self):
  77         result = await self.ws.recv()
  78         if result is not None:
  79             result = json.loads(result)
  80         return result
  81
  82     async def rpc(self, msg, encoder=None):
  83         self.__request_id__ += 1
  84         msg['request-id'] = self.__request_id__
  85         if'params' not in msg:
  86             msg['params'] = {}
  87         if "version" not in msg:
  88             msg['version'] = self.facades[msg['type']]
  89         outgoing = json.dumps(msg, indent=2, cls=encoder)
  90         await self.ws.send(outgoing)
  91         result = await self.recv()
  92         if result and 'error' in result:
  93             raise JujuAPIError(result)
  94         return result
  95
  96     async def clone(self):
  97         """Return a new Connection, connected to the same websocket endpoint
  98         as this one.
  99
 100         """
 101         return await Connection.connect(
 102             self.endpoint,
 103             self.uuid,
 104             self.username,
 105             self.password,
 106             self.cacert,
 107             self.macaroons,
 108         )
 109
 110     async def controller(self):
 111         """Return a Connection to the controller at self.endpoint
 112
 113         """
 114         return await Connection.connect(
 115             self.endpoint,
 116             None,
 117             self.username,
 118             self.password,
 119             self.cacert,
 120             self.macaroons,
 121         )
 122
 123     @classmethod
 124     async def connect(
 125             cls, endpoint, uuid, username, password, cacert=None,
 126             macaroons=None):
 127         """Connect to the websocket.
 128
 129         If uuid is None, the connection will be to the controller. Otherwise it
 130         will be to the model.
 131
 132         """
 133         client = cls(endpoint, uuid, username, password, cacert, macaroons)
 134         await client.open()
 135
 136         redirect_info = await client.redirect_info()
 137         if not redirect_info:
 138             await client.login(username, password, macaroons)
 139             return client
 140
 141         await client.close()
 142         servers = [
 143             s for servers in redirect_info['servers']
 144             for s in servers if s["scope"] == 'public'
 145         ]
 146         for server in servers:
 147             client = cls(
 148                 "{value}:{port}".format(**server), uuid, username,
 149                 password, redirect_info['ca-cert'], macaroons)
 150             await client.open()
 151             try:
 152                 result = await client.login(username, password, macaroons)
 153                 if 'discharge-required-error' in result:
 154                     continue
 155                 return client
 156             except Exception as e:
 157                 await client.close()
 158                 log.exception(e)
 159
 160         raise Exception(
 161             "Couldn't authenticate to %s", endpoint)
 162
 163     @classmethod
 164     async def connect_current(cls):
 165         """Connect to the currently active model.
 166
 167         """
 168         jujudata = JujuData()
 169         controller_name = jujudata.current_controller()
 170         models = jujudata.models()[controller_name]
 171         model_name = models['current-model']
 172
 173         return await cls.connect_model(
 174             '{}:{}'.format(controller_name, model_name))
 175
 176     @classmethod
 177     async def connect_current_controller(cls):
 178         """Connect to the currently active controller.
 179
 180         """
 181         jujudata = JujuData()
 182         controller_name = jujudata.current_controller()
 183
 184         return await cls.connect_controller(controller_name)
 185
 186     @classmethod
 187     async def connect_controller(cls, controller_name):
 188         """Connect to a controller by name.
 189
 190         """
 191         jujudata = JujuData()
 192         controller = jujudata.controllers()[controller_name]
 193         endpoint = controller['api-endpoints'][0]
 194         cacert = controller.get('ca-cert')
 195         accounts = jujudata.accounts()[controller_name]
 196         username = accounts['user']
 197         password = accounts.get('password')
 198         macaroons = get_macaroons() if not password else None
 199
 200         return await cls.connect(
 201             endpoint, None, username, password, cacert, macaroons)
 202
 203     @classmethod
 204     async def connect_model(cls, model):
 205         """Connect to a model by name.
 206
 207         :param str model: <controller>:<model>
 208
 209         """
 210         controller_name, model_name = model.split(':')
 211
 212         jujudata = JujuData()
 213         controller = jujudata.controllers()[controller_name]
 214         endpoint = controller['api-endpoints'][0]
 215         cacert = controller.get('ca-cert')
 216         accounts = jujudata.accounts()[controller_name]
 217         username = accounts['user']
 218         password = accounts.get('password')
 219         models = jujudata.models()[controller_name]
 220         model_uuid = models['models'][model_name]['uuid']
 221         macaroons = get_macaroons() if not password else None
 222
 223         return await cls.connect(
 224             endpoint, model_uuid, username, password, cacert, macaroons)
 225
 226     def build_facades(self, info):
 227         self.facades.clear()
 228         for facade in info:
 229             self.facades[facade['name']] = facade['versions'][-1]
 230
 231     async def login(self, username, password, macaroons=None):
 232         if macaroons:
 233             username = ''
 234             password = ''
 235
 236         if username and not username.startswith('user-'):
 237             username = 'user-{}'.format(username)
 238
 239         result = await self.rpc({
 240             "type": "Admin",
 241             "request": "Login",
 242             "version": 3,
 243             "params": {
 244                 "auth-tag": username,
 245                 "credentials": password,
 246                 "nonce": "".join(random.sample(string.printable, 12)),
 247                 "macaroons": macaroons or []
 248             }})
 249         response = result['response']
 250         self.build_facades(response.get('facades', {}))
 251         self.info = response.copy()
 252         return response
 253
 254     async def redirect_info(self):
 255         try:
 256             result = await self.rpc({
 257                 "type": "Admin",
 258                 "request": "RedirectInfo",
 259                 "version": 3,
 260             })
 261         except JujuAPIError as e:
 262             if e.message == 'not redirected':
 263                 return None
 264             raise
 265         return result['response']
 266
 267
 268 class JujuData:
 269     def __init__(self):
 270         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 271         self.path = os.path.abspath(os.path.expanduser(self.path))
 272
 273     def current_controller(self):
 274         cmd = shlex.split('juju show-controller --format yaml')
 275         output = subprocess.check_output(cmd)
 276         output = yaml.safe_load(output)
 277         return list(output.keys())[0]
 278
 279     def controllers(self):
 280         return self._load_yaml('controllers.yaml', 'controllers')
 281
 282     def models(self):
 283         return self._load_yaml('models.yaml', 'controllers')
 284
 285     def accounts(self):
 286         return self._load_yaml('accounts.yaml', 'controllers')
 287
 288     def _load_yaml(self, filename, key):
 289         filepath = os.path.join(self.path, filename)
 290         with io.open(filepath, 'rt') as f:
 291             return yaml.safe_load(f)[key]
 292
 293
 294 def get_macaroons():
 295     """Decode and return macaroons from default ~/.go-cookies
 296
 297     """
 298     try:
 299         cookie_file = os.path.expanduser('~/.go-cookies')
 300         with open(cookie_file, 'r') as f:
 301             cookies = json.load(f)
 302     except (OSError, ValueError) as e:
 303         log.warn("Couldn't load macaroons from %s", cookie_file)
 304         return []
 305
 306     base64_macaroons = [
 307         c['Value'] for c in cookies
 308         if c['Name'].startswith('macaroon-') and c['Value']
 309     ]
 310
 311     return [
 312         json.loads(base64.b64decode(value).decode('utf-8'))
 313         for value in base64_macaroons
 314     ]