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         if result and 'error' in result:
 100             raise JujuAPIError(result)
 101         return result
 102
 103     def http_headers(self):
 104         """Return dictionary of http headers necessary for making an http
 105         connection to the endpoint of this Connection.
 106
 107         :return: Dictionary of headers
 108
 109         """
 110         if not self.username:
 111             return {}
 112
 113         creds = u'{}:{}'.format(
 114             tag.user(self.username),
 115             self.password or ''
 116         )
 117         token = base64.b64encode(creds.encode())
 118         return {
 119             'Authorization': 'Basic {}'.format(token.decode())
 120         }
 121
 122     def https_connection(self):
 123         """Return an https connection to this Connection's endpoint.
 124
 125         Returns a 3-tuple containing::
 126
 127             1. The :class:`HTTPSConnection` instance
 128             2. Dictionary of auth headers to be used with the connection
 129             3. The root url path (str) to be used for requests.
 130
 131         """
 132         endpoint = self.endpoint
 133         host, remainder = endpoint.split(':', 1)
 134         port = remainder
 135         if '/' in remainder:
 136             port, _ = remainder.split('/', 1)
 137
 138         conn = HTTPSConnection(
 139             host, int(port),
 140             context=self._get_ssl(self.cacert),
 141         )
 142
 143         path = (
 144             "/model/{}".format(self.uuid)
 145             if self.uuid else ""
 146         )
 147         return conn, self.http_headers(), path
 148
 149     async def clone(self):
 150         """Return a new Connection, connected to the same websocket endpoint
 151         as this one.
 152
 153         """
 154         return await Connection.connect(
 155             self.endpoint,
 156             self.uuid,
 157             self.username,
 158             self.password,
 159             self.cacert,
 160             self.macaroons,
 161             self.loop,
 162         )
 163
 164     async def controller(self):
 165         """Return a Connection to the controller at self.endpoint
 166
 167         """
 168         return await Connection.connect(
 169             self.endpoint,
 170             None,
 171             self.username,
 172             self.password,
 173             self.cacert,
 174             self.macaroons,
 175             self.loop,
 176         )
 177
 178     @classmethod
 179     async def connect(
 180             cls, endpoint, uuid, username, password, cacert=None,
 181             macaroons=None, loop=None):
 182         """Connect to the websocket.
 183
 184         If uuid is None, the connection will be to the controller. Otherwise it
 185         will be to the model.
 186
 187         """
 188         client = cls(endpoint, uuid, username, password, cacert, macaroons,
 189                      loop)
 190         await client.open()
 191
 192         redirect_info = await client.redirect_info()
 193         if not redirect_info:
 194             await client.login(username, password, macaroons)
 195             return client
 196
 197         await client.close()
 198         servers = [
 199             s for servers in redirect_info['servers']
 200             for s in servers if s["scope"] == 'public'
 201         ]
 202         for server in servers:
 203             client = cls(
 204                 "{value}:{port}".format(**server), uuid, username,
 205                 password, redirect_info['ca-cert'], macaroons)
 206             await client.open()
 207             try:
 208                 result = await client.login(username, password, macaroons)
 209                 if 'discharge-required-error' in result:
 210                     continue
 211                 return client
 212             except Exception as e:
 213                 await client.close()
 214                 log.exception(e)
 215
 216         raise Exception(
 217             "Couldn't authenticate to %s", endpoint)
 218
 219     @classmethod
 220     async def connect_current(cls, loop=None):
 221         """Connect to the currently active model.
 222
 223         """
 224         jujudata = JujuData()
 225         controller_name = jujudata.current_controller()
 226         model_name = jujudata.current_model()
 227
 228         return await cls.connect_model(
 229             '{}:{}'.format(controller_name, model_name), loop)
 230
 231     @classmethod
 232     async def connect_current_controller(cls, loop=None):
 233         """Connect to the currently active controller.
 234
 235         """
 236         jujudata = JujuData()
 237         controller_name = jujudata.current_controller()
 238         if not controller_name:
 239             raise JujuConnectionError('No current controller')
 240
 241         return await cls.connect_controller(controller_name, loop)
 242
 243     @classmethod
 244     async def connect_controller(cls, controller_name, loop=None):
 245         """Connect to a controller by name.
 246
 247         """
 248         jujudata = JujuData()
 249         controller = jujudata.controllers()[controller_name]
 250         endpoint = controller['api-endpoints'][0]
 251         cacert = controller.get('ca-cert')
 252         accounts = jujudata.accounts()[controller_name]
 253         username = accounts['user']
 254         password = accounts.get('password')
 255         macaroons = get_macaroons() if not password else None
 256
 257         return await cls.connect(
 258             endpoint, None, username, password, cacert, macaroons, loop)
 259
 260     @classmethod
 261     async def connect_model(cls, model, loop=None):
 262         """Connect to a model by name.
 263
 264         :param str model: [<controller>:]<model>
 265
 266         """
 267         jujudata = JujuData()
 268
 269         if ':' in model:
 270             # explicit controller given
 271             controller_name, model_name = model.split(':')
 272         else:
 273             # use the current controller if one isn't explicitly given
 274             controller_name = jujudata.current_controller()
 275             model_name = model
 276
 277         accounts = jujudata.accounts()[controller_name]
 278         username = accounts['user']
 279         # model name must include a user prefix, so add it if it doesn't
 280         if '/' not in model_name:
 281             model_name = '{}/{}'.format(username, model_name)
 282
 283         controller = jujudata.controllers()[controller_name]
 284         endpoint = controller['api-endpoints'][0]
 285         cacert = controller.get('ca-cert')
 286         password = accounts.get('password')
 287         models = jujudata.models()[controller_name]
 288         model_uuid = models['models'][model_name]['uuid']
 289         macaroons = get_macaroons() if not password else None
 290
 291         return await cls.connect(
 292             endpoint, model_uuid, username, password, cacert, macaroons, loop)
 293
 294     def build_facades(self, info):
 295         self.facades.clear()
 296         for facade in info:
 297             self.facades[facade['name']] = facade['versions'][-1]
 298
 299     async def login(self, username, password, macaroons=None):
 300         if macaroons:
 301             username = ''
 302             password = ''
 303
 304         if username and not username.startswith('user-'):
 305             username = 'user-{}'.format(username)
 306
 307         result = await self.rpc({
 308             "type": "Admin",
 309             "request": "Login",
 310             "version": 3,
 311             "params": {
 312                 "auth-tag": username,
 313                 "credentials": password,
 314                 "nonce": "".join(random.sample(string.printable, 12)),
 315                 "macaroons": macaroons or []
 316             }})
 317         response = result['response']
 318         self.build_facades(response.get('facades', {}))
 319         self.info = response.copy()
 320         return response
 321
 322     async def redirect_info(self):
 323         try:
 324             result = await self.rpc({
 325                 "type": "Admin",
 326                 "request": "RedirectInfo",
 327                 "version": 3,
 328             })
 329         except JujuAPIError as e:
 330             if e.message == 'not redirected':
 331                 return None
 332             raise
 333         return result['response']
 334
 335
 336 class JujuData:
 337     def __init__(self):
 338         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 339         self.path = os.path.abspath(os.path.expanduser(self.path))
 340
 341     def current_controller(self):
 342         cmd = shlex.split('juju list-controllers --format yaml')
 343         output = subprocess.check_output(cmd)
 344         output = yaml.safe_load(output)
 345         return output.get('current-controller', '')
 346
 347     def current_model(self, controller_name=None):
 348         if not controller_name:
 349             controller_name = self.current_controller()
 350         models = self.models()[controller_name]
 351         if 'current-model' not in models:
 352             raise JujuError('No current model')
 353         return models['current-model']
 354
 355     def controllers(self):
 356         return self._load_yaml('controllers.yaml', 'controllers')
 357
 358     def models(self):
 359         return self._load_yaml('models.yaml', 'controllers')
 360
 361     def accounts(self):
 362         return self._load_yaml('accounts.yaml', 'controllers')
 363
 364     def _load_yaml(self, filename, key):
 365         filepath = os.path.join(self.path, filename)
 366         with io.open(filepath, 'rt') as f:
 367             return yaml.safe_load(f)[key]
 368
 369
 370 def get_macaroons():
 371     """Decode and return macaroons from default ~/.go-cookies
 372
 373     """
 374     try:
 375         cookie_file = os.path.expanduser('~/.go-cookies')
 376         with open(cookie_file, 'r') as f:
 377             cookies = json.load(f)
 378     except (OSError, ValueError):
 379         log.warn("Couldn't load macaroons from %s", cookie_file)
 380         return []
 381
 382     base64_macaroons = [
 383         c['Value'] for c in cookies
 384         if c['Name'].startswith('macaroon-') and c['Value']
 385     ]
 386
 387     return [
 388         json.loads(base64.b64decode(value).decode('utf-8'))
 389         for value in base64_macaroons
 390     ]