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
  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             server_info = await client.login(username, password, macaroons)
 139             client.build_facades(server_info['facades'])
 140             return client
 141
 142         await client.close()
 143         servers = [
 144             s for servers in redirect_info['servers']
 145             for s in servers if s["scope"] == 'public'
 146         ]
 147         for server in servers:
 148             client = cls(
 149                 "{value}:{port}".format(**server), uuid, username,
 150                 password, redirect_info['ca-cert'], macaroons)
 151             await client.open()
 152             try:
 153                 result = await client.login(username, password, macaroons)
 154                 if 'discharge-required-error' in result:
 155                     continue
 156                 client.build_facades(result['facades'])
 157                 return client
 158             except Exception as e:
 159                 await client.close()
 160                 log.exception(e)
 161
 162         raise Exception(
 163             "Couldn't authenticate to %s", endpoint)
 164
 165     @classmethod
 166     async def connect_current(cls):
 167         """Connect to the currently active model.
 168
 169         """
 170         jujudata = JujuData()
 171         controller_name = jujudata.current_controller()
 172         models = jujudata.models()[controller_name]
 173         model_name = models['current-model']
 174
 175         return await cls.connect_model(
 176             '{}:{}'.format(controller_name, model_name))
 177
 178     @classmethod
 179     async def connect_model(cls, model):
 180         """Connect to a model by name.
 181
 182         :param str model: <controller>:<model>
 183
 184         """
 185         controller_name, model_name = model.split(':')
 186
 187         jujudata = JujuData()
 188         controller = jujudata.controllers()[controller_name]
 189         endpoint = controller['api-endpoints'][0]
 190         cacert = controller.get('ca-cert')
 191         accounts = jujudata.accounts()[controller_name]
 192         username = accounts['user']
 193         password = accounts.get('password')
 194         models = jujudata.models()[controller_name]
 195         model_uuid = models['models'][model_name]['uuid']
 196         macaroons = get_macaroons() if not password else None
 197
 198         return await cls.connect(
 199             endpoint, model_uuid, username, password, cacert, macaroons)
 200
 201     def build_facades(self, info):
 202         self.facades.clear()
 203         for facade in info:
 204             self.facades[facade['name']] = facade['versions'][-1]
 205
 206     async def login(self, username, password, macaroons=None):
 207         if macaroons:
 208             username = ''
 209             password = ''
 210
 211         if username and not username.startswith('user-'):
 212             username = 'user-{}'.format(username)
 213
 214         result = await self.rpc({
 215             "type": "Admin",
 216             "request": "Login",
 217             "version": 3,
 218             "params": {
 219                 "auth-tag": username,
 220                 "credentials": password,
 221                 "nonce": "".join(random.sample(string.printable, 12)),
 222                 "macaroons": macaroons or []
 223             }})
 224         return result['response']
 225
 226     async def redirect_info(self):
 227         try:
 228             result = await self.rpc({
 229                 "type": "Admin",
 230                 "request": "RedirectInfo",
 231                 "version": 3,
 232             })
 233         except JujuAPIError as e:
 234             if e.message == 'not redirected':
 235                 return None
 236             raise
 237         return result['response']
 238
 239
 240 class JujuData:
 241     def __init__(self):
 242         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 243         self.path = os.path.abspath(os.path.expanduser(self.path))
 244
 245     def current_controller(self):
 246         cmd = shlex.split('juju show-controller --format yaml')
 247         output = subprocess.check_output(cmd)
 248         output = yaml.safe_load(output)
 249         return list(output.keys())[0]
 250
 251     def controllers(self):
 252         return self._load_yaml('controllers.yaml', 'controllers')
 253
 254     def models(self):
 255         return self._load_yaml('models.yaml', 'controllers')
 256
 257     def accounts(self):
 258         return self._load_yaml('accounts.yaml', 'controllers')
 259
 260     def _load_yaml(self, filename, key):
 261         filepath = os.path.join(self.path, filename)
 262         with io.open(filepath, 'rt') as f:
 263             return yaml.safe_load(f)[key]
 264
 265
 266 def get_macaroons():
 267     """Decode and return macaroons from default ~/.go-cookies
 268
 269     """
 270     try:
 271         cookie_file = os.path.expanduser('~/.go-cookies')
 272         with open(cookie_file, 'r') as f:
 273             cookies = json.load(f)
 274     except (OSError, ValueError) as e:
 275         log.warn("Couldn't load macaroons from %s", cookie_file)
 276         return []
 277
 278     base64_macaroons = [
 279         c['Value'] for c in cookies
 280         if c['Name'].startswith('macaroon-') and c['Value']
 281     ]
 282
 283     return [
 284         json.loads(base64.b64decode(value).decode('utf-8'))
 285         for value in base64_macaroons
 286     ]