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