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