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