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 from http.client import HTTPSConnection
  13
  14 import asyncio
  15 import yaml
  16
  17 from juju import tag
  18 from juju.errors import JujuError, JujuAPIError, JujuConnectionError
  19
  20 log = logging.getLogger("websocket")
  21
  22
  23 class Connection:
  24     """
  25     Usage::
  26
  27         # Connect to an arbitrary api server
  28         client = await Connection.connect(
  29             api_endpoint, model_uuid, username, password, cacert)
  30
  31         # Connect using a controller/model name
  32         client = await Connection.connect_model('local.local:default')
  33
  34         # Connect to the currently active model
  35         client = await Connection.connect_current()
  36
  37     Note: Any connection method or constructor can accept an optional `loop`
  38     argument to override the default event loop from `asyncio.get_event_loop`.
  39     """
  40     def __init__(
  41             self, endpoint, uuid, username, password, cacert=None,
  42             macaroons=None, loop=None):
  43         self.endpoint = endpoint
  44         self.uuid = uuid
  45         self.username = username
  46         self.password = password
  47         self.macaroons = macaroons
  48         self.cacert = cacert
  49         self.loop = loop or asyncio.get_event_loop()
  50
  51         self.__request_id__ = 0
  52         self.addr = None
  53         self.ws = None
  54         self.facades = {}
  55
  56     @property
  57     def is_open(self):
  58         if self.ws:
  59             return self.ws.open
  60         return False
  61
  62     def _get_ssl(self, cert=None):
  63         return ssl.create_default_context(
  64             purpose=ssl.Purpose.CLIENT_AUTH, cadata=cert)
  65
  66     async def open(self):
  67         if self.uuid:
  68             url = "wss://{}/model/{}/api".format(self.endpoint, self.uuid)
  69         else:
  70             url = "wss://{}/api".format(self.endpoint)
  71
  72         kw = dict()
  73         kw['ssl'] = self._get_ssl(self.cacert)
  74         kw['loop'] = self.loop
  75         self.addr = url
  76         self.ws = await websockets.connect(url, **kw)
  77         log.info("Driver connected to juju %s", url)
  78         return self
  79
  80     async def close(self):
  81         await self.ws.close()
  82
  83     async def recv(self):
  84         result = await self.ws.recv()
  85         if result is not None:
  86             result = json.loads(result)
  87         return result
  88
  89     async def rpc(self, msg, encoder=None):
  90         self.__request_id__ += 1
  91         msg['request-id'] = self.__request_id__
  92         if'params' not in msg:
  93             msg['params'] = {}
  94         if "version" not in msg:
  95             msg['version'] = self.facades[msg['type']]
  96         outgoing = json.dumps(msg, indent=2, cls=encoder)
  97         await self.ws.send(outgoing)
  98         result = await self.recv()
  99
 100         if not result:
 101             return result
 102
 103         if 'error' in result:
 104             # API Error Response
 105             raise JujuAPIError(result)
 106
 107         if not 'response' in result:
 108             # This may never happen
 109             return result
 110
 111         if 'results' in result['response']:
 112             # Check for errors in a result list.
 113             errors = []
 114             for res in result['response']['results']:
 115                 if res.get('error', {}).get('message'):
 116                     errors.append(res['error']['message'])
 117             if errors:
 118                 raise JujuError(errors)
 119
 120         elif result['response'].get('error', {}).get('message'):
 121             raise JujuError(result['response']['error']['message'])
 122
 123         return result
 124
 125     def http_headers(self):
 126         """Return dictionary of http headers necessary for making an http
 127         connection to the endpoint of this Connection.
 128
 129         :return: Dictionary of headers
 130
 131         """
 132         if not self.username:
 133             return {}
 134
 135         creds = u'{}:{}'.format(
 136             tag.user(self.username),
 137             self.password or ''
 138         )
 139         token = base64.b64encode(creds.encode())
 140         return {
 141             'Authorization': 'Basic {}'.format(token.decode())
 142         }
 143
 144     def https_connection(self):
 145         """Return an https connection to this Connection's endpoint.
 146
 147         Returns a 3-tuple containing::
 148
 149             1. The :class:`HTTPSConnection` instance
 150             2. Dictionary of auth headers to be used with the connection
 151             3. The root url path (str) to be used for requests.
 152
 153         """
 154         endpoint = self.endpoint
 155         host, remainder = endpoint.split(':', 1)
 156         port = remainder
 157         if '/' in remainder:
 158             port, _ = remainder.split('/', 1)
 159
 160         conn = HTTPSConnection(
 161             host, int(port),
 162             context=self._get_ssl(self.cacert),
 163         )
 164
 165         path = (
 166             "/model/{}".format(self.uuid)
 167             if self.uuid else ""
 168         )
 169         return conn, self.http_headers(), path
 170
 171     async def clone(self):
 172         """Return a new Connection, connected to the same websocket endpoint
 173         as this one.
 174
 175         """
 176         return await Connection.connect(
 177             self.endpoint,
 178             self.uuid,
 179             self.username,
 180             self.password,
 181             self.cacert,
 182             self.macaroons,
 183             self.loop,
 184         )
 185
 186     async def controller(self):
 187         """Return a Connection to the controller at self.endpoint
 188
 189         """
 190         return await Connection.connect(
 191             self.endpoint,
 192             None,
 193             self.username,
 194             self.password,
 195             self.cacert,
 196             self.macaroons,
 197             self.loop,
 198         )
 199
 200     @classmethod
 201     async def connect(
 202             cls, endpoint, uuid, username, password, cacert=None,
 203             macaroons=None, loop=None):
 204         """Connect to the websocket.
 205
 206         If uuid is None, the connection will be to the controller. Otherwise it
 207         will be to the model.
 208
 209         """
 210         client = cls(endpoint, uuid, username, password, cacert, macaroons,
 211                      loop)
 212         await client.open()
 213
 214         redirect_info = await client.redirect_info()
 215         if not redirect_info:
 216             await client.login(username, password, macaroons)
 217             return client
 218
 219         await client.close()
 220         servers = [
 221             s for servers in redirect_info['servers']
 222             for s in servers if s["scope"] == 'public'
 223         ]
 224         for server in servers:
 225             client = cls(
 226                 "{value}:{port}".format(**server), uuid, username,
 227                 password, redirect_info['ca-cert'], macaroons)
 228             await client.open()
 229             try:
 230                 result = await client.login(username, password, macaroons)
 231                 if 'discharge-required-error' in result:
 232                     continue
 233                 return client
 234             except Exception as e:
 235                 await client.close()
 236                 log.exception(e)
 237
 238         raise Exception(
 239             "Couldn't authenticate to %s", endpoint)
 240
 241     @classmethod
 242     async def connect_current(cls, loop=None):
 243         """Connect to the currently active model.
 244
 245         """
 246         jujudata = JujuData()
 247         controller_name = jujudata.current_controller()
 248         model_name = jujudata.current_model()
 249
 250         return await cls.connect_model(
 251             '{}:{}'.format(controller_name, model_name), loop)
 252
 253     @classmethod
 254     async def connect_current_controller(cls, loop=None):
 255         """Connect to the currently active controller.
 256
 257         """
 258         jujudata = JujuData()
 259         controller_name = jujudata.current_controller()
 260         if not controller_name:
 261             raise JujuConnectionError('No current controller')
 262
 263         return await cls.connect_controller(controller_name, loop)
 264
 265     @classmethod
 266     async def connect_controller(cls, controller_name, loop=None):
 267         """Connect to a controller by name.
 268
 269         """
 270         jujudata = JujuData()
 271         controller = jujudata.controllers()[controller_name]
 272         endpoint = controller['api-endpoints'][0]
 273         cacert = controller.get('ca-cert')
 274         accounts = jujudata.accounts()[controller_name]
 275         username = accounts['user']
 276         password = accounts.get('password')
 277         macaroons = get_macaroons() if not password else None
 278
 279         return await cls.connect(
 280             endpoint, None, username, password, cacert, macaroons, loop)
 281
 282     @classmethod
 283     async def connect_model(cls, model, loop=None):
 284         """Connect to a model by name.
 285
 286         :param str model: [<controller>:]<model>
 287
 288         """
 289         jujudata = JujuData()
 290
 291         if ':' in model:
 292             # explicit controller given
 293             controller_name, model_name = model.split(':')
 294         else:
 295             # use the current controller if one isn't explicitly given
 296             controller_name = jujudata.current_controller()
 297             model_name = model
 298
 299         accounts = jujudata.accounts()[controller_name]
 300         username = accounts['user']
 301         # model name must include a user prefix, so add it if it doesn't
 302         if '/' not in model_name:
 303             model_name = '{}/{}'.format(username, model_name)
 304
 305         controller = jujudata.controllers()[controller_name]
 306         endpoint = controller['api-endpoints'][0]
 307         cacert = controller.get('ca-cert')
 308         password = accounts.get('password')
 309         models = jujudata.models()[controller_name]
 310         model_uuid = models['models'][model_name]['uuid']
 311         macaroons = get_macaroons() if not password else None
 312
 313         return await cls.connect(
 314             endpoint, model_uuid, username, password, cacert, macaroons, loop)
 315
 316     def build_facades(self, info):
 317         self.facades.clear()
 318         for facade in info:
 319             self.facades[facade['name']] = facade['versions'][-1]
 320
 321     async def login(self, username, password, macaroons=None):
 322         if macaroons:
 323             username = ''
 324             password = ''
 325
 326         if username and not username.startswith('user-'):
 327             username = 'user-{}'.format(username)
 328
 329         result = await self.rpc({
 330             "type": "Admin",
 331             "request": "Login",
 332             "version": 3,
 333             "params": {
 334                 "auth-tag": username,
 335                 "credentials": password,
 336                 "nonce": "".join(random.sample(string.printable, 12)),
 337                 "macaroons": macaroons or []
 338             }})
 339         response = result['response']
 340         self.build_facades(response.get('facades', {}))
 341         self.info = response.copy()
 342         return response
 343
 344     async def redirect_info(self):
 345         try:
 346             result = await self.rpc({
 347                 "type": "Admin",
 348                 "request": "RedirectInfo",
 349                 "version": 3,
 350             })
 351         except JujuAPIError as e:
 352             if e.message == 'not redirected':
 353                 return None
 354             raise
 355         return result['response']
 356
 357
 358 class JujuData:
 359     def __init__(self):
 360         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 361         self.path = os.path.abspath(os.path.expanduser(self.path))
 362
 363     def current_controller(self):
 364         cmd = shlex.split('juju list-controllers --format yaml')
 365         output = subprocess.check_output(cmd)
 366         output = yaml.safe_load(output)
 367         return output.get('current-controller', '')
 368
 369     def current_model(self, controller_name=None):
 370         if not controller_name:
 371             controller_name = self.current_controller()
 372         models = self.models()[controller_name]
 373         if 'current-model' not in models:
 374             raise JujuError('No current model')
 375         return models['current-model']
 376
 377     def controllers(self):
 378         return self._load_yaml('controllers.yaml', 'controllers')
 379
 380     def models(self):
 381         return self._load_yaml('models.yaml', 'controllers')
 382
 383     def accounts(self):
 384         return self._load_yaml('accounts.yaml', 'controllers')
 385
 386     def _load_yaml(self, filename, key):
 387         filepath = os.path.join(self.path, filename)
 388         with io.open(filepath, 'rt') as f:
 389             return yaml.safe_load(f)[key]
 390
 391
 392 def get_macaroons():
 393     """Decode and return macaroons from default ~/.go-cookies
 394
 395     """
 396     try:
 397         cookie_file = os.path.expanduser('~/.go-cookies')
 398         with open(cookie_file, 'r') as f:
 399             cookies = json.load(f)
 400     except (OSError, ValueError):
 401         log.warn("Couldn't load macaroons from %s", cookie_file)
 402         return []
 403
 404     base64_macaroons = [
 405         c['Value'] for c in cookies
 406         if c['Name'].startswith('macaroon-') and c['Value']
 407     ]
 408
 409     return [
 410         json.loads(base64.b64decode(value).decode('utf-8'))
 411         for value in base64_macaroons
 412     ]