juju/client/connection.py

   1 import asyncio
   2 import io
   3 import json
   4 import logging
   5 import os
   6 import random
   7 import ssl
   8 import string
   9 import websockets
  10
  11 import yaml
  12
  13 log = logging.getLogger("websocket")
  14
  15
  16 class Connection:
  17     """
  18     Usage::
  19
  20         # Connect to an arbitrary api server
  21         client = await Connection.connect(
  22             api_endpoint, model_uuid, username, password, cacert)
  23
  24         # Connect using a controller/model name
  25         client = await Connection.connect_model('local.local:default')
  26
  27         # Connect to the currently active model
  28         client = await Connection.connect_current()
  29
  30     """
  31     def __init__(self):
  32         self.__request_id__ = 0
  33         self.addr = None
  34         self.ws = None
  35         self.facades = {}
  36
  37     def _get_ssl(self, cert):
  38         return ssl.create_default_context(
  39             purpose=ssl.Purpose.CLIENT_AUTH, cadata=cert)
  40
  41     async def open(self, addr, cert=None):
  42         kw = dict()
  43         if cert:
  44             kw['ssl'] = self._get_ssl(cert)
  45         self.addr = addr
  46         self.ws = await websockets.connect(addr, **kw)
  47         return self
  48
  49     async def close(self):
  50         await self.ws.close()
  51
  52     async def recv(self):
  53         result = await self.ws.recv()
  54         if result is not None:
  55             result = json.loads(result)
  56         return result
  57
  58     async def rpc(self, msg, encoder=None):
  59         self.__request_id__ += 1
  60         msg['RequestId'] = self.__request_id__
  61         if'Params' not in msg:
  62             msg['Params'] = {}
  63         if "Version" not in msg:
  64             msg['Version'] = self.facades[msg['Type']]
  65         outgoing = json.dumps(msg, indent=2, cls=encoder)
  66         await self.ws.send(outgoing)
  67         result = await self.recv()
  68         log.debug("send %s got %s", msg, result)
  69         if result and 'Error' in result:
  70             raise RuntimeError(result)
  71         return result
  72
  73     @classmethod
  74     async def connect(cls, endpoint, uuid, username, password, cacert=None):
  75         url = "wss://{}/model/{}/api".format(endpoint, uuid)
  76         client = cls()
  77         await client.open(url, cacert)
  78         server_info = await client.login(username, password)
  79         client.build_facades(server_info['facades'])
  80         log.info("Driver connected to juju %s", endpoint)
  81         return client
  82
  83     @classmethod
  84     async def connect_current(cls):
  85         """Connect to the currently active model.
  86
  87         """
  88         jujudata = JujuData()
  89         controller_name = jujudata.current_controller()
  90         controller = jujudata.controllers()[controller_name]
  91         endpoint = controller['api-endpoints'][0]
  92         cacert = controller.get('ca-cert')
  93         accounts = jujudata.accounts()[controller_name]
  94         username = accounts['current-account']
  95         password = accounts['accounts'][username]['password']
  96         models = jujudata.models()[controller_name]['accounts'][username]
  97         model_name = models['current-model']
  98         model_uuid = models['models'][model_name]['uuid']
  99
 100         return await cls.connect(
 101             endpoint, model_uuid, username, password, cacert)
 102
 103     @classmethod
 104     async def connect_model(cls, model):
 105         """Connect to a model by name.
 106
 107         :param str model: <controller>:<model>
 108
 109         """
 110         controller_name, model_name = model.split(':')
 111
 112         jujudata = JujuData()
 113         controller = jujudata.controllers()[controller_name]
 114         endpoint = controller['api-endpoints'][0]
 115         cacert = controller.get('ca-cert')
 116         accounts = jujudata.accounts()[controller_name]
 117         username = accounts['current-account']
 118         password = accounts['accounts'][username]['password']
 119         models = jujudata.models()[controller_name]['accounts'][username]
 120         model_uuid = models['models'][model_name]['uuid']
 121
 122         return await cls.connect(
 123             endpoint, model_uuid, username, password, cacert)
 124
 125     def build_facades(self, info):
 126         self.facades.clear()
 127         for facade in info:
 128             self.facades[facade['Name']] = facade['Versions'][-1]
 129
 130     async def login(self, username, password):
 131         if not username.startswith('user-'):
 132             username = 'user-{}'.format(username)
 133
 134         result = await self.rpc({
 135             "Type": "Admin",
 136             "Request": "Login",
 137             "Version": 3,
 138             "Params": {
 139                 "auth-tag": username,
 140                 "credentials": password,
 141                 "Nonce": "".join(random.sample(string.printable, 12)),
 142             }})
 143         return result['Response']
 144
 145
 146 class JujuData:
 147     def __init__(self):
 148         self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
 149         self.path = os.path.abspath(os.path.expanduser(self.path))
 150
 151     def current_controller(self):
 152         try:
 153             filepath = os.path.join(self.path, 'current-controller')
 154             with io.open(filepath, 'rt') as f:
 155                 return f.read().strip()
 156         except OSError as e:
 157             log.exception(e)
 158             return None
 159
 160     def controllers(self):
 161         return self._load_yaml('controllers.yaml', 'controllers')
 162
 163     def models(self):
 164         return self._load_yaml('models.yaml', 'controllers')
 165
 166     def accounts(self):
 167         return self._load_yaml('accounts.yaml', 'controllers')
 168
 169     def _load_yaml(self, filename, key):
 170         filepath = os.path.join(self.path, filename)
 171         with io.open(filepath, 'rt') as f:
 172             return yaml.safe_load(f)[key]