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         jujudata = JujuData()
 267
 268         if ':' in model:
 269             controller_name, model_name = model.split(':')
 270         else:
 271             controller_name = jujudata.current_controller()
 272             model_name = model
 273
 274         controller = jujudata.controllers()[controller_name]
 275         endpoint = controller['api-endpoints'][0]
 276         cacert = controller.get('ca-cert')
 277         accounts = jujudata.accounts()[controller_name]
 278         username = accounts['user']
 279         password = accounts.get('password')
 280         models = jujudata.models()[controller_name]
 281         if '/' not in model_name:
 282             model_name = '{}/{}'.format(username, model_name)
 283         model_uuid = models['models'][model_name]['uuid']
 284         macaroons = get_macaroons() if not password else None
 285
 286         return await cls.connect(
 287             endpoint, model_uuid, username, password, cacert, macaroons, loop)
 288
 289     def build_facades(self, info):
 290         self.facades.clear()
 291         for facade in info:
 292             self.facades[facade['name']] = facade['versions'][-1]
 293
 294     async def login(self, username, password, macaroons=None):
 295         if macaroons:
 296             username = ''
 297             password = ''
 298
 299         if username and not username.startswith('user-'):
 300             username = 'user-{}'.format(username)
 301
 302         result = await self.rpc({
 303             "type": "Admin",
 304             "request": "Login",
 305             "version": 3,
 306             "params": {
 307                 "auth-tag": username,
 308                 "credentials": password,
 309                 "nonce": "".join(random.sample(string.printable, 12)),
 310                 "macaroons": macaroons or []
 311             }})
 312         response = result['response']
 313         self.build_facades(response.get('facades', {}))
 314         self.info = response.copy()
 315         return response
 316
 317     async def redirect_info(self):
 318         try:
 319             result = await self.rpc({
 320                 "type": "Admin",
 321                 "request": "RedirectInfo",
 322                 "version": 3,
 323             })
 324         except JujuAPIError as e:
 325             if e.message == 'not redirected':
 326                 return None
 327             raise
 328         return result['response']
 329
 330
 331 class JujuData:
 332     def __init__(self):
 333         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 334         self.path = os.path.abspath(os.path.expanduser(self.path))
 335
 336     def current_controller(self):
 337         cmd = shlex.split('juju list-controllers --format yaml')
 338         output = subprocess.check_output(cmd)
 339         output = yaml.safe_load(output)
 340         return output.get('current-controller', '')
 341
 342     def controllers(self):
 343         return self._load_yaml('controllers.yaml', 'controllers')
 344
 345     def models(self):
 346         return self._load_yaml('models.yaml', 'controllers')
 347
 348     def accounts(self):
 349         return self._load_yaml('accounts.yaml', 'controllers')
 350
 351     def _load_yaml(self, filename, key):
 352         filepath = os.path.join(self.path, filename)
 353         with io.open(filepath, 'rt') as f:
 354             return yaml.safe_load(f)[key]
 355
 356
 357 def get_macaroons():
 358     """Decode and return macaroons from default ~/.go-cookies
 359
 360     """
 361     try:
 362         cookie_file = os.path.expanduser('~/.go-cookies')
 363         with open(cookie_file, 'r') as f:
 364             cookies = json.load(f)
 365     except (OSError, ValueError):
 366         log.warn("Couldn't load macaroons from %s", cookie_file)
 367         return []
 368
 369     base64_macaroons = [
 370         c['Value'] for c in cookies
 371         if c['Name'].startswith('macaroon-') and c['Value']
 372     ]
 373
 374     return [
 375         json.loads(base64.b64decode(value).decode('utf-8'))
 376         for value in base64_macaroons
 377     ]