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