12 from concurrent
.futures
import CancelledError
13 from http
.client
import HTTPSConnection
14 from pathlib
import Path
19 from juju
import tag
, utils
20 from juju
.client
import client
21 from juju
.errors
import JujuError
, JujuAPIError
, JujuConnectionError
22 from juju
.utils
import IdQueue
24 log
= logging
.getLogger("websocket")
29 Monitor helper class for our Connection class.
31 Contains a reference to an instantiated Connection, along with a
32 reference to the Connection.receiver Future. Upon inspecttion of
33 these objects, this class determines whether the connection is in
34 an 'error', 'connected' or 'disconnected' state.
36 Use this class to stay up to date on the health of a connection,
37 and take appropriate action if the connection errors out due to
38 network issues or other unexpected circumstances.
42 CONNECTED
= 'connected'
43 DISCONNECTED
= 'disconnected'
46 def __init__(self
, connection
):
47 self
.connection
= connection
48 self
.close_called
= asyncio
.Event(loop
=self
.connection
.loop
)
49 self
.receiver_stopped
= asyncio
.Event(loop
=self
.connection
.loop
)
50 self
.pinger_stopped
= asyncio
.Event(loop
=self
.connection
.loop
)
51 self
.receiver_stopped
.set()
52 self
.pinger_stopped
.set()
57 Determine the status of the connection and receiver, and return
58 ERROR, CONNECTED, or DISCONNECTED as appropriate.
60 For simplicity, we only consider ourselves to be connected
61 after the Connection class has setup a receiver task. This
62 only happens after the websocket is open, and the connection
63 isn't usable until that receiver has been started.
67 # DISCONNECTED: connection not yet open
68 if not self
.connection
.ws
:
69 return self
.DISCONNECTED
70 if self
.receiver_stopped
.is_set():
71 return self
.DISCONNECTED
73 # ERROR: Connection closed (or errored), but we didn't call
75 if not self
.close_called
.is_set() and self
.receiver_stopped
.is_set():
77 if not self
.close_called
.is_set() and not self
.connection
.ws
.open:
78 # The check for self.receiver_stopped existing above guards
79 # against the case where we're not open because we simply
80 # haven't setup the connection yet.
83 # DISCONNECTED: cleanly disconnected.
84 if self
.close_called
.is_set() and not self
.connection
.ws
.open:
85 return self
.DISCONNECTED
87 # CONNECTED: everything is fine!
88 if self
.connection
.ws
.open:
91 # UNKNOWN: We should never hit this state -- if we do,
92 # something went wrong with the logic above, and we do not
93 # know what state the connection is in.
101 # Connect to an arbitrary api server
102 client = await Connection.connect(
103 api_endpoint, model_uuid, username, password, cacert)
105 # Connect using a controller/model name
106 client = await Connection.connect_model('local.local:default')
108 # Connect to the currently active model
109 client = await Connection.connect_current()
111 Note: Any connection method or constructor can accept an optional `loop`
112 argument to override the default event loop from `asyncio.get_event_loop`.
115 self
, endpoint
, uuid
, username
, password
, cacert
=None,
116 macaroons
=None, loop
=None):
117 self
.endpoint
= endpoint
120 self
.macaroons
= macaroons
125 self
.username
= username
126 self
.password
= password
128 self
.loop
= loop
or asyncio
.get_event_loop()
130 self
.__request
_id
__ = 0
134 self
.messages
= IdQueue(loop
=self
.loop
)
135 self
.monitor
= Monitor(connection
=self
)
143 def _get_ssl(self
, cert
=None):
144 return ssl
.create_default_context(
145 purpose
=ssl
.Purpose
.CLIENT_AUTH
, cadata
=cert
)
147 async def open(self
):
149 url
= "wss://{}/model/{}/api".format(self
.endpoint
, self
.uuid
)
151 url
= "wss://{}/api".format(self
.endpoint
)
154 kw
['ssl'] = self
._get
_ssl
(self
.cacert
)
155 kw
['loop'] = self
.loop
157 self
.ws
= await websockets
.connect(url
, **kw
)
158 self
.loop
.create_task(self
.receiver())
159 self
.monitor
.receiver_stopped
.clear()
160 log
.info("Driver connected to juju %s", url
)
161 self
.monitor
.close_called
.clear()
164 async def close(self
):
167 self
.monitor
.close_called
.set()
168 await self
.monitor
.pinger_stopped
.wait()
169 await self
.monitor
.receiver_stopped
.wait()
170 await self
.ws
.close()
172 async def recv(self
, request_id
):
174 raise websockets
.exceptions
.ConnectionClosed(0, 'websocket closed')
175 return await self
.messages
.get(request_id
)
177 async def receiver(self
):
180 result
= await utils
.run_with_interrupt(
182 self
.monitor
.close_called
,
184 if self
.monitor
.close_called
.is_set():
186 if result
is not None:
187 result
= json
.loads(result
)
188 await self
.messages
.put(result
['request-id'], result
)
189 except CancelledError
:
191 except Exception as e
:
192 await self
.messages
.put_all(e
)
193 if isinstance(e
, websockets
.ConnectionClosed
):
194 # ConnectionClosed is not really exceptional for us,
195 # but it may be for any pending message listeners
197 log
.exception("Error in receiver")
200 self
.monitor
.receiver_stopped
.set()
202 async def pinger(self
):
204 A Controller can time us out if we are silent for too long. This
205 is especially true in JaaS, which has a fairly strict timeout.
207 To prevent timing out, we send a ping every ten seconds.
210 async def _do_ping():
212 await pinger_facade
.Ping()
213 await asyncio
.sleep(10, loop
=self
.loop
)
214 except CancelledError
:
217 pinger_facade
= client
.PingerFacade
.from_connection(self
)
220 await utils
.run_with_interrupt(
222 self
.monitor
.close_called
,
224 if self
.monitor
.close_called
.is_set():
227 self
.monitor
.pinger_stopped
.set()
229 async def rpc(self
, msg
, encoder
=None):
230 self
.__request
_id
__ += 1
231 msg
['request-id'] = self
.__request
_id
__
232 if'params' not in msg
:
234 if "version" not in msg
:
235 msg
['version'] = self
.facades
[msg
['type']]
236 outgoing
= json
.dumps(msg
, indent
=2, cls
=encoder
)
237 await self
.ws
.send(outgoing
)
238 result
= await self
.recv(msg
['request-id'])
243 if 'error' in result
:
245 raise JujuAPIError(result
)
247 if 'response' not in result
:
248 # This may never happen
251 if 'results' in result
['response']:
252 # Check for errors in a result list.
254 for res
in result
['response']['results']:
255 if res
.get('error', {}).get('message'):
256 errors
.append(res
['error']['message'])
258 raise JujuError(errors
)
260 elif result
['response'].get('error', {}).get('message'):
261 raise JujuError(result
['response']['error']['message'])
265 def http_headers(self
):
266 """Return dictionary of http headers necessary for making an http
267 connection to the endpoint of this Connection.
269 :return: Dictionary of headers
272 if not self
.username
:
275 creds
= u
'{}:{}'.format(
276 tag
.user(self
.username
),
279 token
= base64
.b64encode(creds
.encode())
281 'Authorization': 'Basic {}'.format(token
.decode())
284 def https_connection(self
):
285 """Return an https connection to this Connection's endpoint.
287 Returns a 3-tuple containing::
289 1. The :class:`HTTPSConnection` instance
290 2. Dictionary of auth headers to be used with the connection
291 3. The root url path (str) to be used for requests.
294 endpoint
= self
.endpoint
295 host
, remainder
= endpoint
.split(':', 1)
298 port
, _
= remainder
.split('/', 1)
300 conn
= HTTPSConnection(
302 context
=self
._get
_ssl
(self
.cacert
),
306 "/model/{}".format(self
.uuid
)
309 return conn
, self
.http_headers(), path
311 async def clone(self
):
312 """Return a new Connection, connected to the same websocket endpoint
316 return await Connection
.connect(
326 async def controller(self
):
327 """Return a Connection to the controller at self.endpoint
330 return await Connection
.connect(
340 async def _try_endpoint(self
, endpoint
, cacert
):
345 self
.endpoint
= endpoint
349 result
= await self
.login()
350 if 'discharge-required-error' in result
['response']:
351 log
.info('Macaroon discharge required, disconnecting')
354 log
.info('Authenticated')
356 except JujuAPIError
as e
:
357 if e
.error_code
!= 'redirection required':
359 log
.info('Controller requested redirect')
360 redirect_info
= await self
.redirect_info()
361 redir_cacert
= redirect_info
['ca-cert']
363 ("{value}:{port}".format(**s
), redir_cacert
)
364 for servers
in redirect_info
['servers']
365 for s
in servers
if s
["scope"] == 'public'
370 return success
, result
, new_endpoints
374 cls
, endpoint
, uuid
, username
, password
, cacert
=None,
375 macaroons
=None, loop
=None):
376 """Connect to the websocket.
378 If uuid is None, the connection will be to the controller. Otherwise it
379 will be to the model.
382 client
= cls(endpoint
, uuid
, username
, password
, cacert
, macaroons
,
384 endpoints
= [(endpoint
, cacert
)]
386 _endpoint
, _cacert
= endpoints
.pop(0)
387 success
, result
, new_endpoints
= await client
._try
_endpoint
(
391 endpoints
.extend(new_endpoints
)
393 # ran out of endpoints without a successful login
394 raise Exception("Couldn't authenticate to {}".format(endpoint
))
396 response
= result
['response']
397 client
.info
= response
.copy()
398 client
.build_facades(response
.get('facades', {}))
399 client
.loop
.create_task(client
.pinger())
400 client
.monitor
.pinger_stopped
.clear()
405 async def connect_current(cls
, loop
=None):
406 """Connect to the currently active model.
409 jujudata
= JujuData()
411 controller_name
= jujudata
.current_controller()
412 if not controller_name
:
413 raise JujuConnectionError('No current controller')
415 model_name
= jujudata
.current_model()
417 return await cls
.connect_model(
418 '{}:{}'.format(controller_name
, model_name
), loop
)
421 async def connect_current_controller(cls
, loop
=None):
422 """Connect to the currently active controller.
425 jujudata
= JujuData()
426 controller_name
= jujudata
.current_controller()
427 if not controller_name
:
428 raise JujuConnectionError('No current controller')
430 return await cls
.connect_controller(controller_name
, loop
)
433 async def connect_controller(cls
, controller_name
, loop
=None):
434 """Connect to a controller by name.
437 jujudata
= JujuData()
438 controller
= jujudata
.controllers()[controller_name
]
439 endpoint
= controller
['api-endpoints'][0]
440 cacert
= controller
.get('ca-cert')
441 accounts
= jujudata
.accounts()[controller_name
]
442 username
= accounts
['user']
443 password
= accounts
.get('password')
444 macaroons
= get_macaroons(controller_name
) if not password
else None
446 return await cls
.connect(
447 endpoint
, None, username
, password
, cacert
, macaroons
, loop
)
450 async def connect_model(cls
, model
, loop
=None):
451 """Connect to a model by name.
453 :param str model: [<controller>:]<model>
456 jujudata
= JujuData()
459 # explicit controller given
460 controller_name
, model_name
= model
.split(':')
462 # use the current controller if one isn't explicitly given
463 controller_name
= jujudata
.current_controller()
466 accounts
= jujudata
.accounts()[controller_name
]
467 username
= accounts
['user']
468 # model name must include a user prefix, so add it if it doesn't
469 if '/' not in model_name
:
470 model_name
= '{}/{}'.format(username
, model_name
)
472 controller
= jujudata
.controllers()[controller_name
]
473 endpoint
= controller
['api-endpoints'][0]
474 cacert
= controller
.get('ca-cert')
475 password
= accounts
.get('password')
476 models
= jujudata
.models()[controller_name
]
477 model_uuid
= models
['models'][model_name
]['uuid']
478 macaroons
= get_macaroons(controller_name
) if not password
else None
480 return await cls
.connect(
481 endpoint
, model_uuid
, username
, password
, cacert
, macaroons
, loop
)
483 def build_facades(self
, facades
):
485 for facade
in facades
:
486 self
.facades
[facade
['name']] = facade
['versions'][-1]
488 async def login(self
):
489 username
= self
.username
490 if username
and not username
.startswith('user-'):
491 username
= 'user-{}'.format(username
)
493 result
= await self
.rpc({
498 "auth-tag": username
,
499 "credentials": self
.password
,
500 "nonce": "".join(random
.sample(string
.printable
, 12)),
501 "macaroons": self
.macaroons
505 async def redirect_info(self
):
507 result
= await self
.rpc({
509 "request": "RedirectInfo",
512 except JujuAPIError
as e
:
513 if e
.message
== 'not redirected':
516 return result
['response']
521 self
.path
= os
.environ
.get('JUJU_DATA') or '~/.local/share/juju'
522 self
.path
= os
.path
.abspath(os
.path
.expanduser(self
.path
))
524 def current_controller(self
):
525 cmd
= shlex
.split('juju list-controllers --format yaml')
526 output
= subprocess
.check_output(cmd
)
527 output
= yaml
.safe_load(output
)
528 return output
.get('current-controller', '')
530 def current_model(self
, controller_name
=None):
531 if not controller_name
:
532 controller_name
= self
.current_controller()
533 models
= self
.models()[controller_name
]
534 if 'current-model' not in models
:
535 raise JujuError('No current model')
536 return models
['current-model']
538 def controllers(self
):
539 return self
._load
_yaml
('controllers.yaml', 'controllers')
542 return self
._load
_yaml
('models.yaml', 'controllers')
545 return self
._load
_yaml
('accounts.yaml', 'controllers')
547 def _load_yaml(self
, filename
, key
):
548 filepath
= os
.path
.join(self
.path
, filename
)
549 with io
.open(filepath
, 'rt') as f
:
550 return yaml
.safe_load(f
)[key
]
553 def get_macaroons(controller_name
=None):
554 """Decode and return macaroons from default ~/.go-cookies
559 cookie_files
.append('~/.local/share/juju/cookies/{}.json'.format(
561 cookie_files
.append('~/.go-cookies')
562 for cookie_file
in cookie_files
:
563 cookie_file
= Path(cookie_file
).expanduser()
564 if cookie_file
.exists():
566 cookies
= json
.loads(cookie_file
.read_text())
568 except (OSError, ValueError):
569 log
.warn("Couldn't load macaroons from %s", cookie_file
)
572 log
.warn("Couldn't load macaroons from %s", ' or '.join(cookie_files
))
576 c
['Value'] for c
in cookies
577 if c
['Name'].startswith('macaroon-') and c
['Value']
581 json
.loads(base64
.b64decode(value
).decode('utf-8'))
582 for value
in base64_macaroons