+++ /dev/null
-import asyncio
-import base64
-import json
-import logging
-import ssl
-import urllib.request
-import weakref
-from concurrent.futures import CancelledError
-from http.client import HTTPSConnection
-
-import macaroonbakery.httpbakery as httpbakery
-import macaroonbakery.bakery as bakery
-import websockets
-from juju import errors, tag, utils
-from juju.client import client
-from juju.utils import IdQueue
-
-log = logging.getLogger('juju.client.connection')
-
-
-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 inspection 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)
-
- @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)
- stopped = connection._receiver_task.stopped.is_set()
- if stopped 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)
-
- Note: Any connection method or constructor can accept an optional `loop`
- argument to override the default event loop from `asyncio.get_event_loop`.
- """
-
- MAX_FRAME_SIZE = 2**22
- "Maximum size for a single frame. Defaults to 4MB."
-
- @classmethod
- async def connect(
- cls,
- endpoint=None,
- uuid=None,
- username=None,
- password=None,
- cacert=None,
- bakery_client=None,
- loop=None,
- max_frame_size=None,
- retries=3,
- retry_backoff=10,
- ):
- """Connect to the websocket.
-
- If uuid is None, the connection will be to the controller. Otherwise it
- will be to the model.
-
- :param str endpoint: The hostname:port of the controller to connect to.
- :param str uuid: The model UUID to connect to (None for a
- controller-only connection).
- :param str username: The username for controller-local users (or None
- to use macaroon-based login.)
- :param str password: The password for controller-local users.
- :param str cacert: The CA certificate of the controller
- (PEM formatted).
- :param httpbakery.Client bakery_client: The macaroon bakery client to
- to use when performing macaroon-based login. Macaroon tokens
- acquired when logging will be saved to bakery_client.cookies.
- If this is None, a default bakery_client will be used.
- :param asyncio.BaseEventLoop loop: The event loop to use for async
- operations.
- :param int max_frame_size: The maximum websocket frame size to allow.
- :param int retries: When connecting or reconnecting, and all endpoints
- fail, how many times to retry the connection before giving up.
- :param int retry_backoff: Number of seconds to increase the wait
- between connection retry attempts (a backoff of 10 with 3 retries
- would wait 10s, 20s, and 30s).
- """
- self = cls()
- if endpoint is None:
- raise ValueError('no endpoint provided')
- self.uuid = uuid
- if bakery_client is None:
- bakery_client = httpbakery.Client()
- self.bakery_client = bakery_client
- if username and '@' in username and not username.endswith('@local'):
- # We're trying to log in as an external user - we need to use
- # macaroon authentication with no username or password.
- if password is not None:
- raise errors.JujuAuthError('cannot log in as external '
- 'user with a password')
- username = None
- self.usertag = tag.user(username)
- self.password = password
- self.loop = loop or asyncio.get_event_loop()
-
- self.__request_id__ = 0
-
- # The following instance variables are initialized by the
- # _connect_with_redirect method, but create them here
- # as a reminder that they will exist.
- self.addr = None
- self.ws = None
- self.endpoint = None
- self.cacert = None
- self.info = None
-
- # Create that _Task objects but don't start the tasks yet.
- self._pinger_task = _Task(self._pinger, self.loop)
- self._receiver_task = _Task(self._receiver, self.loop)
-
- self._retries = retries
- self._retry_backoff = retry_backoff
-
- self.facades = {}
- self.messages = IdQueue(loop=self.loop)
- self.monitor = Monitor(connection=self)
- if max_frame_size is None:
- max_frame_size = self.MAX_FRAME_SIZE
- self.max_frame_size = max_frame_size
- await self._connect_with_redirect([(endpoint, cacert)])
- return self
-
- @property
- def username(self):
- if not self.usertag:
- return None
- return self.usertag[len('user-'):]
-
- @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, endpoint, cacert):
- if self.uuid:
- url = "wss://{}/model/{}/api".format(endpoint, self.uuid)
- else:
- url = "wss://{}/api".format(endpoint)
-
- return (await websockets.connect(
- url,
- ssl=self._get_ssl(cacert),
- loop=self.loop,
- max_size=self.max_frame_size,
- ), url, endpoint, cacert)
-
- async def close(self):
- if not self.ws:
- return
- self.monitor.close_called.set()
- await self._pinger_task.stopped.wait()
- await self._receiver_task.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
-
- 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
- except websockets.exceptions.ConnectionClosed:
- # The connection has closed - we can't do anything
- # more until the connection is restarted.
- log.debug('ping failed because of closed connection')
- pass
-
- async def rpc(self, msg, encoder=None):
- '''Make an RPC to the API. The message is encoded as JSON
- using the given encoder if any.
- :param msg: Parameters for the call (will be encoded as JSON).
- :param encoder: Encoder to be used when encoding the message.
- :return: The result of the call.
- :raises JujuAPIError: When there's an error returned.
- :raises JujuError:
- '''
- 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)
- log.debug('connection {} -> {}'.format(id(self), outgoing))
- for attempt in range(3):
- if self.monitor.status == Monitor.DISCONNECTED:
- # closed cleanly; shouldn't try to reconnect
- raise websockets.exceptions.ConnectionClosed(
- 0, 'websocket closed')
- 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)
- if self.monitor.status != Monitor.CONNECTED:
- # reconnect failed; abort and shutdown
- log.error('RPC: Automatic reconnect failed')
- raise
- result = await self._recv(msg['request-id'])
- log.debug('connection {} <- {}'.format(id(self), result))
-
- if not result:
- return result
-
- if 'error' in result:
- # API Error Response
- raise errors.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.
- # TODO This loses the results that might have succeeded.
- # Perhaps JujuError should return all the results including
- # errors, or perhaps a keyword parameter to the rpc method
- # could be added to trigger this behaviour.
- err_results = []
- for res in result['response']['results']:
- if res.get('error', {}).get('message'):
- err_results.append(res['error']['message'])
- if err_results:
- raise errors.JujuError(err_results)
-
- elif result['response'].get('error', {}).get('message'):
- raise errors.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.usertag:
- return {}
-
- creds = u'{}:{}'.format(
- self.usertag,
- 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.connect_params())
-
- def connect_params(self):
- """Return a tuple of parameters suitable for passing to
- Connection.connect that can be used to make a new connection
- to the same controller (and model if specified. The first
- element in the returned tuple holds the endpoint argument;
- the other holds a dict of the keyword args.
- """
- return {
- 'endpoint': self.endpoint,
- 'uuid': self.uuid,
- 'username': self.username,
- 'password': self.password,
- 'cacert': self.cacert,
- 'bakery_client': self.bakery_client,
- 'loop': self.loop,
- 'max_frame_size': self.max_frame_size,
- }
-
- async def controller(self):
- """Return a Connection to the controller at self.endpoint
- """
- return await Connection.connect(
- self.endpoint,
- username=self.username,
- password=self.password,
- cacert=self.cacert,
- bakery_client=self.bakery_client,
- loop=self.loop,
- max_frame_size=self.max_frame_size,
- )
-
- 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_with_login([(self.endpoint, self.cacert)])
-
- async def _connect(self, endpoints):
- if len(endpoints) == 0:
- raise errors.JujuConnectionError('no endpoints to connect to')
-
- async def _try_endpoint(endpoint, cacert, delay):
- if delay:
- await asyncio.sleep(delay)
- return await self._open(endpoint, cacert)
-
- # Try all endpoints in parallel, with slight increasing delay (+100ms
- # for each subsequent endpoint); the delay allows us to prefer the
- # earlier endpoints over the latter. Use first successful connection.
- tasks = [self.loop.create_task(_try_endpoint(endpoint, cacert,
- 0.1 * i))
- for i, (endpoint, cacert) in enumerate(endpoints)]
- for attempt in range(self._retries + 1):
- for task in asyncio.as_completed(tasks, loop=self.loop):
- try:
- result = await task
- break
- except ConnectionError:
- continue # ignore; try another endpoint
- else:
- _endpoints_str = ', '.join([endpoint
- for endpoint, cacert in endpoints])
- if attempt < self._retries:
- log.debug('Retrying connection to endpoints: {}; '
- 'attempt {} of {}'.format(_endpoints_str,
- attempt + 1,
- self._retries + 1))
- await asyncio.sleep((attempt + 1) * self._retry_backoff)
- continue
- else:
- raise errors.JujuConnectionError(
- 'Unable to connect to any endpoint: '
- '{}'.format(_endpoints_str))
- # only executed if inner loop's else did not continue
- # (i.e., inner loop did break due to successful connection)
- break
- for task in tasks:
- task.cancel()
- self.ws = result[0]
- self.addr = result[1]
- self.endpoint = result[2]
- self.cacert = result[3]
- self._receiver_task.start()
- log.debug("Driver connected to juju %s", self.addr)
- self.monitor.close_called.clear()
-
- async def _connect_with_login(self, endpoints):
- """Connect to the websocket.
-
- If uuid is None, the connection will be to the controller. Otherwise it
- will be to the model.
- :return: The response field of login response JSON object.
- """
- success = False
- try:
- await self._connect(endpoints)
- # It's possible that we may get several discharge-required errors,
- # corresponding to different levels of authentication, so retry
- # a few times.
- for i in range(0, 2):
- result = (await self.login())['response']
- macaroonJSON = result.get('discharge-required')
- if macaroonJSON is None:
- self.info = result
- success = True
- return result
- macaroon = bakery.Macaroon.from_dict(macaroonJSON)
- self.bakery_client.handle_error(
- httpbakery.Error(
- code=httpbakery.ERR_DISCHARGE_REQUIRED,
- message=result.get('discharge-required-error'),
- version=macaroon.version,
- info=httpbakery.ErrorInfo(
- macaroon=macaroon,
- macaroon_path=result.get('macaroon-path'),
- ),
- ),
- # note: remove the port number.
- 'https://' + self.endpoint + '/',
- )
- raise errors.JujuAuthError('failed to authenticate '
- 'after several attempts')
- finally:
- if not success:
- await self.close()
-
- async def _connect_with_redirect(self, endpoints):
- try:
- login_result = await self._connect_with_login(endpoints)
- except errors.JujuRedirectException as e:
- login_result = await self._connect_with_login(e.endpoints)
- self._build_facades(login_result.get('facades', {}))
- self._pinger_task.start()
-
- def _build_facades(self, facades):
- self.facades.clear()
- for facade in facades:
- self.facades[facade['name']] = facade['versions'][-1]
-
- async def login(self):
- params = {}
- params['auth-tag'] = self.usertag
- if self.password:
- params['credentials'] = self.password
- else:
- macaroons = _macaroons_for_domain(self.bakery_client.cookies,
- self.endpoint)
- params['macaroons'] = [[bakery.macaroon_to_dict(m) for m in ms]
- for ms in macaroons]
-
- try:
- return await self.rpc({
- "type": "Admin",
- "request": "Login",
- "version": 3,
- "params": params,
- })
- except errors.JujuAPIError as e:
- if e.error_code != 'redirection required':
- raise
- log.info('Controller requested redirect')
- # Fetch additional redirection information now so that
- # we can safely close the connection after login
- # fails.
- redirect_info = (await self.rpc({
- "type": "Admin",
- "request": "RedirectInfo",
- "version": 3,
- }))['response']
- raise errors.JujuRedirectException(redirect_info) from e
-
-
-class _Task:
- def __init__(self, task, loop):
- self.stopped = asyncio.Event(loop=loop)
- self.stopped.set()
- self.task = task
- self.loop = loop
-
- def start(self):
- async def run():
- try:
- return await self.task()
- finally:
- self.stopped.set()
- self.stopped.clear()
- self.loop.create_task(run())
-
-
-def _macaroons_for_domain(cookies, domain):
- '''Return any macaroons from the given cookie jar that
- apply to the given domain name.'''
- req = urllib.request.Request('https://' + domain + '/')
- cookies.add_cookie_header(req)
- return httpbakery.extract_macaroons(req)
projects / osm / N2VC.git / blobdiff