+import base64
+import io
+import json
+import logging
+import os
+import random
+import shlex
+import ssl
+import string
+import subprocess
+import weakref
+import websockets
+from concurrent.futures import CancelledError
+from http.client import HTTPSConnection
+from pathlib import Path
+
+import asyncio
+import yaml
+
+from juju import tag, utils
+from juju.client import client
+from juju.errors import JujuError, JujuAPIError, JujuConnectionError
+from juju.utils import IdQueue
+
+log = logging.getLogger("websocket")
+
+
+class Monitor:
+ """
+ Monitor helper class for our Connection class.
+
+ Contains a reference to an instantiated Connection, along with a
+ reference to the Connection.receiver Future. Upon inspecttion of
+ these objects, this class determines whether the connection is in
+ an 'error', 'connected' or 'disconnected' state.
+
+ Use this class to stay up to date on the health of a connection,
+ and take appropriate action if the connection errors out due to
+ network issues or other unexpected circumstances.
+
+ """
+ ERROR = 'error'
+ CONNECTED = 'connected'
+ DISCONNECTING = 'disconnecting'
+ DISCONNECTED = 'disconnected'
+
+ def __init__(self, connection):
+ self.connection = weakref.ref(connection)
+ self.reconnecting = asyncio.Lock(loop=connection.loop)
+ self.close_called = asyncio.Event(loop=connection.loop)
+ self.receiver_stopped = asyncio.Event(loop=connection.loop)
+ self.pinger_stopped = asyncio.Event(loop=connection.loop)
+ self.receiver_stopped.set()
+ self.pinger_stopped.set()
+
+ @property
+ def status(self):
+ """
+ Determine the status of the connection and receiver, and return
+ ERROR, CONNECTED, or DISCONNECTED as appropriate.
+
+ For simplicity, we only consider ourselves to be connected
+ after the Connection class has setup a receiver task. This
+ only happens after the websocket is open, and the connection
+ isn't usable until that receiver has been started.
+
+ """
+ connection = self.connection()
+
+ # the connection instance was destroyed but someone kept
+ # a separate reference to the monitor for some reason
+ if not connection:
+ return self.DISCONNECTED
+
+ # connection cleanly disconnected or not yet opened
+ if not connection.ws:
+ return self.DISCONNECTED
+
+ # close called but not yet complete
+ if self.close_called.is_set():
+ return self.DISCONNECTING
+
+ # connection closed uncleanly (we didn't call connection.close)
+ if self.receiver_stopped.is_set() or not connection.ws.open:
+ return self.ERROR
+
+ # everything is fine!
+ return self.CONNECTED
+
+
+class Connection:
+ """
+ Usage::
+
+ # Connect to an arbitrary api server
+ client = await Connection.connect(
+ api_endpoint, model_uuid, username, password, cacert)
+
+ # Connect using a controller/model name
+ client = await Connection.connect_model('local.local:default')
+
+ # Connect to the currently active model
+ client = await Connection.connect_current()
+
+ Note: Any connection method or constructor can accept an optional `loop`
+ argument to override the default event loop from `asyncio.get_event_loop`.
+ """
+
+ DEFAULT_FRAME_SIZE = 'default_frame_size'
+ MAX_FRAME_SIZE = 2**22
+ "Maximum size for a single frame. Defaults to 4MB."
+
+ def __init__(
+ self, endpoint, uuid, username, password, cacert=None,
+ macaroons=None, loop=None, max_frame_size=DEFAULT_FRAME_SIZE):
+ self.endpoint = endpoint
+ self._endpoint = endpoint
+ self.uuid = uuid
+ if macaroons:
+ self.macaroons = macaroons
+ self.username = ''
+ self.password = ''
+ else:
+ self.macaroons = []
+ self.username = username
+ self.password = password
+ self.cacert = cacert
+ self._cacert = cacert
+ self.loop = loop or asyncio.get_event_loop()
+
+ self.__request_id__ = 0
+ self.addr = None
+ self.ws = None
+ self.facades = {}
+ self.messages = IdQueue(loop=self.loop)
+ self.monitor = Monitor(connection=self)
+ if max_frame_size is self.DEFAULT_FRAME_SIZE:
+ max_frame_size = self.MAX_FRAME_SIZE
+ self.max_frame_size = max_frame_size
+
+ @property
+ def is_open(self):
+ return self.monitor.status == Monitor.CONNECTED
+
+ def _get_ssl(self, cert=None):
+ return ssl.create_default_context(
+ purpose=ssl.Purpose.CLIENT_AUTH, cadata=cert)
+
+ async def open(self):
+ if self.uuid:
+ url = "wss://{}/model/{}/api".format(self.endpoint, self.uuid)
+ else:
+ url = "wss://{}/api".format(self.endpoint)
+
+ kw = dict()
+ kw['ssl'] = self._get_ssl(self.cacert)
+ kw['loop'] = self.loop
+ kw['max_size'] = self.max_frame_size
+ self.addr = url
+ self.ws = await websockets.connect(url, **kw)
+ self.loop.create_task(self.receiver())
+ self.monitor.receiver_stopped.clear()
+ log.info("Driver connected to juju %s", url)
+ self.monitor.close_called.clear()
+ return self
+
+ async def close(self):
+ if not self.ws:
+ return
+ self.monitor.close_called.set()
+ await self.monitor.pinger_stopped.wait()
+ await self.monitor.receiver_stopped.wait()
+ await self.ws.close()
+ self.ws = None
+
+ async def recv(self, request_id):
+ if not self.is_open:
+ raise websockets.exceptions.ConnectionClosed(0, 'websocket closed')
+ return await self.messages.get(request_id)
+
+ async def receiver(self):
+ try:
+ while self.is_open:
+ result = await utils.run_with_interrupt(
+ self.ws.recv(),
+ self.monitor.close_called,
+ loop=self.loop)
+ if self.monitor.close_called.is_set():
+ break
+ if result is not None:
+ result = json.loads(result)
+ await self.messages.put(result['request-id'], result)
+ except CancelledError:
+ pass
+ except websockets.ConnectionClosed as e:
+ log.warning('Receiver: Connection closed, reconnecting')
+ await self.messages.put_all(e)
+ # the reconnect has to be done as a task because the receiver will
+ # be cancelled by the reconnect and we don't want the reconnect
+ # to be aborted half-way through
+ self.loop.create_task(self.reconnect())
+ return
+ except Exception as e:
+ log.exception("Error in receiver")
+ # make pending listeners aware of the error
+ await self.messages.put_all(e)
+ raise
+ finally:
+ self.monitor.receiver_stopped.set()
+
+ async def pinger(self):
+ '''
+ A Controller can time us out if we are silent for too long. This
+ is especially true in JaaS, which has a fairly strict timeout.
+
+ To prevent timing out, we send a ping every ten seconds.
+
+ '''
+ async def _do_ping():
+ try:
+ await pinger_facade.Ping()
+ await asyncio.sleep(10, loop=self.loop)
+ except CancelledError:
+ pass
+
+ pinger_facade = client.PingerFacade.from_connection(self)
+ try:
+ while True:
+ await utils.run_with_interrupt(
+ _do_ping(),
+ self.monitor.close_called,
+ loop=self.loop)
+ if self.monitor.close_called.is_set():
+ break
+ finally:
+ self.monitor.pinger_stopped.set()
+ return
+
+ async def rpc(self, msg, encoder=None):
+ self.__request_id__ += 1
+ msg['request-id'] = self.__request_id__
+ if'params' not in msg:
+ msg['params'] = {}
+ if "version" not in msg:
+ msg['version'] = self.facades[msg['type']]
+ outgoing = json.dumps(msg, indent=2, cls=encoder)
+ for attempt in range(3):
+ try:
+ await self.ws.send(outgoing)
+ break
+ except websockets.ConnectionClosed:
+ if attempt == 2:
+ raise
+ log.warning('RPC: Connection closed, reconnecting')
+ # the reconnect has to be done in a separate task because,
+ # if it is triggered by the pinger, then this RPC call will
+ # be cancelled when the pinger is cancelled by the reconnect,
+ # and we don't want the reconnect to be aborted halfway through
+ await asyncio.wait([self.reconnect()], loop=self.loop)
+ result = await self.recv(msg['request-id'])
+
+ if not result:
+ return result
+
+ if 'error' in result:
+ # API Error Response
+ raise JujuAPIError(result)
+
+ if 'response' not in result:
+ # This may never happen
+ return result
+
+ if 'results' in result['response']:
+ # Check for errors in a result list.
+ errors = []
+ for res in result['response']['results']:
+ if res.get('error', {}).get('message'):
+ errors.append(res['error']['message'])
+ if errors:
+ raise JujuError(errors)
+
+ elif result['response'].get('error', {}).get('message'):
+ raise JujuError(result['response']['error']['message'])
+
+ return result
+
+ def http_headers(self):
+ """Return dictionary of http headers necessary for making an http
+ connection to the endpoint of this Connection.
+
+ :return: Dictionary of headers
+
+ """
+ if not self.username:
+ return {}
+
+ creds = u'{}:{}'.format(
+ tag.user(self.username),
+ self.password or ''
+ )
+ token = base64.b64encode(creds.encode())
+ return {
+ 'Authorization': 'Basic {}'.format(token.decode())
+ }
+
+ def https_connection(self):
+ """Return an https connection to this Connection's endpoint.
+
+ Returns a 3-tuple containing::
+
+ 1. The :class:`HTTPSConnection` instance
+ 2. Dictionary of auth headers to be used with the connection
+ 3. The root url path (str) to be used for requests.
+
+ """
+ endpoint = self.endpoint
+ host, remainder = endpoint.split(':', 1)
+ port = remainder
+ if '/' in remainder:
+ port, _ = remainder.split('/', 1)
+
+ conn = HTTPSConnection(
+ host, int(port),
+ context=self._get_ssl(self.cacert),
+ )
+
+ path = (
+ "/model/{}".format(self.uuid)
+ if self.uuid else ""
+ )
+ return conn, self.http_headers(), path
+
+ async def clone(self):
+ """Return a new Connection, connected to the same websocket endpoint
+ as this one.
+
+ """
+ return await Connection.connect(
+ self.endpoint,
+ self.uuid,
+ self.username,
+ self.password,
+ self.cacert,
+ self.macaroons,
+ self.loop,
+ self.max_frame_size,
+ )
+
+ async def controller(self):
+ """Return a Connection to the controller at self.endpoint
+
+ """
+ return await Connection.connect(
+ self.endpoint,
+ None,
+ self.username,
+ self.password,
+ self.cacert,
+ self.macaroons,
+ self.loop,
+ )
+
+ async def _try_endpoint(self, endpoint, cacert):
+ success = False
+ result = None
+ new_endpoints = []
+
+ self.endpoint = endpoint
+ self.cacert = cacert
+ await self.open()
+ try:
+ result = await self.login()
+ if 'discharge-required-error' in result['response']:
+ log.info('Macaroon discharge required, disconnecting')
+ else:
+ # successful login!
+ log.info('Authenticated')
+ success = True
+ except JujuAPIError as e:
+ if e.error_code != 'redirection required':
+ raise
+ log.info('Controller requested redirect')
+ redirect_info = await self.redirect_info()
+ redir_cacert = redirect_info['ca-cert']
+ new_endpoints = [
+ ("{value}:{port}".format(**s), redir_cacert)
+ for servers in redirect_info['servers']
+ for s in servers if s["scope"] == 'public'
+ ]
+ finally:
+ if not success:
+ await self.close()
+ return success, result, new_endpoints
+
+ async def reconnect(self):
+ """ Force a reconnection.
+ """
+ monitor = self.monitor
+ if monitor.reconnecting.locked() or monitor.close_called.is_set():
+ return
+ async with monitor.reconnecting:
+ await self.close()
+ await self._connect()
+
+ async def _connect(self):
+ endpoints = [(self._endpoint, self._cacert)]
+ while endpoints:
+ _endpoint, _cacert = endpoints.pop(0)
+ success, result, new_endpoints = await self._try_endpoint(
+ _endpoint, _cacert)
+ if success:
+ break
+ endpoints.extend(new_endpoints)
+ else:
+ # ran out of endpoints without a successful login
+ raise Exception("Couldn't authenticate to {}".format(
+ self._endpoint))
+
+ response = result['response']
+ self.info = response.copy()
+ self.build_facades(response.get('facades', {}))
+ self.loop.create_task(self.pinger())
+ self.monitor.pinger_stopped.clear()
+
+ @classmethod
+ async def connect(
+ cls, endpoint, uuid, username, password, cacert=None,
+ macaroons=None, loop=None, max_frame_size=None):
+ """Connect to the websocket.
+
+ If uuid is None, the connection will be to the controller. Otherwise it
+ will be to the model.
+
+ """
+ client = cls(endpoint, uuid, username, password, cacert, macaroons,
+ loop, max_frame_size)
+ await client._connect()
+ return client
+
+ @classmethod
+ async def connect_current(cls, loop=None, max_frame_size=None):
+ """Connect to the currently active model.
+
+ """
+ jujudata = JujuData()
+
+ controller_name = jujudata.current_controller()
+ if not controller_name:
+ raise JujuConnectionError('No current controller')
+
+ model_name = jujudata.current_model()
+
+ return await cls.connect_model(
+ '{}:{}'.format(controller_name, model_name), loop, max_frame_size)
+
+ @classmethod
+ async def connect_current_controller(cls, loop=None, max_frame_size=None):
+ """Connect to the currently active controller.
+
+ """
+ jujudata = JujuData()
+ controller_name = jujudata.current_controller()
+ if not controller_name:
+ raise JujuConnectionError('No current controller')
+
+ return await cls.connect_controller(controller_name, loop,
+ max_frame_size)
+
+ @classmethod
+ async def connect_controller(cls, controller_name, loop=None,
+ max_frame_size=None):
+ """Connect to a controller by name.
+
+ """
+ jujudata = JujuData()
+ controller = jujudata.controllers()[controller_name]
+ endpoint = controller['api-endpoints'][0]
+ cacert = controller.get('ca-cert')
+ accounts = jujudata.accounts()[controller_name]
+ username = accounts['user']
+ password = accounts.get('password')
+ macaroons = get_macaroons(controller_name) if not password else None
+
+ return await cls.connect(
+ endpoint, None, username, password, cacert, macaroons, loop,
+ max_frame_size)
+
+ @classmethod
+ async def connect_model(cls, model, loop=None, max_frame_size=None):
+ """Connect to a model by name.
+
+ :param str model: [<controller>:]<model>
+
+ """
+ jujudata = JujuData()
+
+ if ':' in model:
+ # explicit controller given
+ controller_name, model_name = model.split(':')
+ else:
+ # use the current controller if one isn't explicitly given
+ controller_name = jujudata.current_controller()
+ model_name = model
+
+ accounts = jujudata.accounts()[controller_name]
+ username = accounts['user']
+ # model name must include a user prefix, so add it if it doesn't
+ if '/' not in model_name:
+ model_name = '{}/{}'.format(username, model_name)
+
+ controller = jujudata.controllers()[controller_name]
+ endpoint = controller['api-endpoints'][0]
+ cacert = controller.get('ca-cert')
+ password = accounts.get('password')
+ models = jujudata.models()[controller_name]
+ model_uuid = models['models'][model_name]['uuid']
+ macaroons = get_macaroons(controller_name) if not password else None
+
+ return await cls.connect(
+ endpoint, model_uuid, username, password, cacert, macaroons, loop,
+ max_frame_size)
+
+ def build_facades(self, facades):
+ self.facades.clear()
+ for facade in facades:
+ self.facades[facade['name']] = facade['versions'][-1]
+
+ async def login(self):
+ username = self.username
+ if username and not username.startswith('user-'):
+ username = 'user-{}'.format(username)
+
+ result = await self.rpc({
+ "type": "Admin",
+ "request": "Login",
+ "version": 3,
+ "params": {
+ "auth-tag": username,
+ "credentials": self.password,
+ "nonce": "".join(random.sample(string.printable, 12)),
+ "macaroons": self.macaroons
+ }})
+ return result
+
+ async def redirect_info(self):
+ try:
+ result = await self.rpc({
+ "type": "Admin",
+ "request": "RedirectInfo",
+ "version": 3,
+ })
+ except JujuAPIError as e:
+ if e.message == 'not redirected':
+ return None
+ raise
+ return result['response']
+
+
+class JujuData:
+ def __init__(self):
+ self.path = os.environ.get('JUJU_DATA') or '~/.local/share/juju'
+ self.path = os.path.abspath(os.path.expanduser(self.path))
+
+ def current_controller(self):
+ cmd = shlex.split('juju list-controllers --format yaml')
+ output = subprocess.check_output(cmd)
+ output = yaml.safe_load(output)
+ return output.get('current-controller', '')
+
+ def current_model(self, controller_name=None):
+ if not controller_name:
+ controller_name = self.current_controller()
+ models = self.models()[controller_name]
+ if 'current-model' not in models:
+ raise JujuError('No current model')
+ return models['current-model']
+
+ def controllers(self):
+ return self._load_yaml('controllers.yaml', 'controllers')
+
+ def models(self):
+ return self._load_yaml('models.yaml', 'controllers')
+
+ def accounts(self):
+ return self._load_yaml('accounts.yaml', 'controllers')
+
+ def _load_yaml(self, filename, key):
+ filepath = os.path.join(self.path, filename)
+ with io.open(filepath, 'rt') as f:
+ return yaml.safe_load(f)[key]
+
+
+def get_macaroons(controller_name=None):
+ """Decode and return macaroons from default ~/.go-cookies
+
+ """
+ cookie_files = []
+ if controller_name:
+ cookie_files.append('~/.local/share/juju/cookies/{}.json'.format(
+ controller_name))
+ cookie_files.append('~/.go-cookies')
+ for cookie_file in cookie_files:
+ cookie_file = Path(cookie_file).expanduser()
+ if cookie_file.exists():
+ try:
+ cookies = json.loads(cookie_file.read_text())
+ break
+ except (OSError, ValueError):
+ log.warn("Couldn't load macaroons from %s", cookie_file)
+ return []
+ else:
+ log.warn("Couldn't load macaroons from %s", ' or '.join(cookie_files))
+ return []
+
+ base64_macaroons = [
+ c['Value'] for c in cookies
+ if c['Name'].startswith('macaroon-') and c['Value']
+ ]
+
+ return [
+ json.loads(base64.b64decode(value).decode('utf-8'))
+ for value in base64_macaroons
+ ]
projects / osm / N2VC.git / blobdiff