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