juju/client/connection.py

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