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