diff --git a/zaza/charm_tests/vault/setup.py b/zaza/charm_tests/vault/setup.py index 5866d65..769c760 100644 --- a/zaza/charm_tests/vault/setup.py +++ b/zaza/charm_tests/vault/setup.py @@ -1,18 +1,8 @@ #!/usr/bin/env python3 -import hvac -import logging -import os -import requests -import time -import unittest -import urllib3 -import uuid -import yaml - -import zaza.charm_tests.test_utils as test_utils import zaza.charm_tests.vault.utils as vault_utils + def basic_setup(): clients = vault_utils.get_clients() unseal_client = clients[0] @@ -20,8 +10,7 @@ def basic_setup(): # The credentials are written to a file to allow the tests to be re-run # this is mainly useful for manually working on the tests. if initialized: - vault_creds = vault_utils.get_credentails() + vault_creds = vault_utils.get_credentails() else: - vault_creds = vault_utils.init_vault(unseal_client[1]) - vault_utils.store_credentails(vault_creds) - print(vault_creds) + vault_creds = vault_utils.init_vault(unseal_client[1]) + vault_utils.store_credentails(vault_creds) diff --git a/zaza/charm_tests/vault/tests.py b/zaza/charm_tests/vault/tests.py index 6995013..33fae16 100644 --- a/zaza/charm_tests/vault/tests.py +++ b/zaza/charm_tests/vault/tests.py @@ -1,14 +1,9 @@ #!/usr/bin/env python3 import hvac -import logging -import os -import requests import time import unittest -import urllib3 import uuid -import yaml import zaza.charm_tests.test_utils as test_utils import zaza.charm_tests.vault.utils as vault_utils @@ -18,18 +13,6 @@ class VaultTest(unittest.TestCase): @classmethod def setUpClass(cls): -# vault_utils = VaultUtils() -# cls.clients = vault_utils.get_clients() -# unseal_client = cls.clients[0] -# initialized = vault_utils.is_initialized(unseal_client) -# # The credentials are written to a file to allow the tests to be re-run -# # this is mainly useful for manually working on the tests. -# auth_file = "/tmp/vault_tests.yaml" -# if initialized: -# vault_creds = vault_utils.get_credentails_from_file(auth_file) -# else: -# vault_creds = vault_utils.init_vault(unseal_client[1]) -# vault_utils.write_credentails(auth_file, vault_creds) cls.clients = vault_utils.get_clients() vault_creds = vault_utils.get_credentails() vault_utils.unseal_all(cls.clients, vault_creds['keys'][0]) diff --git a/zaza/charm_tests/vault/utils.py b/zaza/charm_tests/vault/utils.py index eba9739..8c0859c 100644 --- a/zaza/charm_tests/vault/utils.py +++ b/zaza/charm_tests/vault/utils.py @@ -1,18 +1,12 @@ #!/usr/bin/env python3 -import asyncio import hvac -import logging -import os import requests -import time import tempfile -import unittest +import time import urllib3 -import uuid import yaml -import zaza.charm_tests.test_utils as test_utils import zaza.model AUTH_FILE = "vault_tests.yaml" @@ -26,6 +20,7 @@ def get_client(vault_url): """ return hvac.Client(url=vault_url) + def init_vault(client, shares=1, threshold=1): """Initialise vault @@ -36,6 +31,7 @@ def init_vault(client, shares=1, threshold=1): """ return client.initialize(shares, threshold) + def get_clients(units=None): """Create a list of clients, one per vault server @@ -50,6 +46,7 @@ def get_clients(units=None): clients.append((unit, get_client(vault_url))) return clients + def is_initialized(client): """Check if vault is initialized @@ -72,6 +69,7 @@ def is_initialized(client): raise Exception("Cannot connect") return initialized + def get_credentails(): unit = zaza.model.get_first_unit('vault') with tempfile.TemporaryDirectory() as tmpdirname: @@ -81,13 +79,15 @@ def get_credentails(): creds = yaml.load(stream) return creds + def store_credentails(creds): unit = zaza.model.get_first_unit('vault') with tempfile.NamedTemporaryFile(mode='w') as fp: fp.write(yaml.dump(creds)) fp.flush() zaza.model.scp_to_unit(unit, fp.name, '~/{}'.format(AUTH_FILE)) - + + def get_credentails_from_file(auth_file): """Read the vault credentials from the auth_file @@ -98,6 +98,7 @@ def get_credentails_from_file(auth_file): vault_creds = yaml.load(stream) return vault_creds + def write_credentails(auth_file, vault_creds): """Write the vault credentials to the auth_file @@ -106,6 +107,7 @@ def write_credentails(auth_file, vault_creds): with open(auth_file, 'w') as outfile: yaml.dump(vault_creds, outfile, default_flow_style=False) + def unseal_all(clients, key): """Unseal all the vaults with the given clients with the provided key @@ -116,6 +118,7 @@ def unseal_all(clients, key): if client.is_sealed(): client.unseal(key) + def auth_all(clients, token): """Authenticate all the given clients with the provided token diff --git a/zaza/model.py b/zaza/model.py index 9b0f802..9bca14c 100644 --- a/zaza/model.py +++ b/zaza/model.py @@ -1,13 +1,10 @@ -import functools - from juju import loop import subprocess -import logging -import sys import yaml from juju.model import Model + async def deployed(filter=None): # Create a Model instance. We need to connect our Model to a Juju api # server before we can use it. @@ -21,6 +18,7 @@ async def deployed(filter=None): # Disconnect from the api server and cleanup. await model.disconnect() + def get_unit_from_name(unit_name, model): app = unit_name.split('/')[0] for u in model.applications[app].units: @@ -31,6 +29,7 @@ def get_unit_from_name(unit_name, model): raise Exception return unit + async def _scp_to_unit(unit_name, source_file, target_file): model = Model() await model.connect_current() @@ -41,6 +40,7 @@ async def _scp_to_unit(unit_name, source_file, target_file): # Disconnect from the api server and cleanup. await model.disconnect() + async def _get_first_unit(app_name): model = Model() await model.connect_current() @@ -51,22 +51,27 @@ async def _get_first_unit(app_name): await model.disconnect() return unit + def scp_from_unit(unit_name, source_file, target_file): cmd = ['juju', 'scp', '{}:{}'.format(unit_name, source_file), target_file] subprocess.check_call(cmd) + def scp_to_unit(unit_name, source_file, target_file): cmd = ['juju', 'scp', source_file, '{}:{}'.format(unit_name, target_file)] subprocess.check_call(cmd) + def get_status(): status = subprocess.check_output(['juju', 'status', '--format', 'yaml']) return yaml.load(status) + def get_first_unit(app_name): status = get_status() return sorted(status['applications'][app_name]['units'].keys())[0] + def get_app_ips(app_name): status = get_status() addresses = [] @@ -74,235 +79,11 @@ def get_app_ips(app_name): addresses.append(unit['public-address']) return addresses -def get_unit_from_name(unit_name, model): - """Return the units that corresponds to the name in the given model - - :param unit_name: Name of unit to match - :type unit_name: str - :param model: Model to perform lookup in - :type model: juju.model.Model - :returns: Unit matching given name - :rtype: juju.unit.Unit or None - """ - app = unit_name.split('/')[0] - unit = None - for u in model.applications[app].units: - if u.entity_id == unit_name: - unit = u - break - else: - raise Exception - return unit - - -async def run_in_model(model_name, f, add_model_arg=False, awaitable=True): - """Run the given function in the model matching the model_name - - :param model_name: Name of model to run function in - :type model_name: str - :param f: Function to run with given moel in focus - :type f: functools.partial - :param add_model_arg: Whether to add kwarg pointing at model to the given - function before running it - :type add_model_arg: boolean - :param awaitable: Whether f is awaitable - :type awaitable: boolean - :returns: Output of f - :rtype: Unknown, depends on the passed in function - """ - model = Model() - await model.connect_model(model_name) - output = None - try: - if add_model_arg: - f.keywords.update(model=model) - if awaitable: - output = await f() - else: - output = f() - finally: - # Disconnect from the api server and cleanup. - await model.disconnect() - return output - - -def scp_to_unit(unit_name, model_name, source, destination, user='ubuntu', - proxy=False, scp_opts=''): - """Transfer files to unit_name in model_name. - - :param unit_name: Name of unit to scp to - :type unit_name: str - :param model_name: Name of model unit is in - :type model_name: str - :param source: Local path of file(s) to transfer - :type source: str - :param destination: Remote destination of transferred files - :type source: str - :param user: Remote username - :type source: str - :param proxy: Proxy through the Juju API server - :type proxy: bool - :param scp_opts: Additional options to the scp command - :type scp_opts: str - """ - async def _scp_to_unit(unit_name, source, destination, user, proxy, - scp_opts, model): - unit = get_unit_from_name(unit_name, model) - await unit.scp_to(source, destination, user=user, proxy=proxy, - scp_opts=scp_opts) - scp_func = functools.partial( - _scp_to_unit, - unit_name, - source, - destination, - user=user, - proxy=proxy, - scp_opts=scp_opts) - loop.run( - run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True)) - - -def scp_to_all_units(application_name, model_name, source, destination, - user='ubuntu', proxy=False, scp_opts=''): - """Transfer files from to all units of an application - - :param application_name: Name of application to scp file to - :type unit_name: str - :param model_name: Name of model unit is in - :type model_name: str - :param source: Local path of file(s) to transfer - :type source: str - :param destination: Remote destination of transferred files - :type source: str - :param user: Remote username - :type source: str - :param proxy: Proxy through the Juju API server - :type proxy: bool - :param scp_opts: Additional options to the scp command - :type scp_opts: str - """ - async def _scp_to_all_units(application_name, source, destination, user, - proxy, scp_opts, model): - for unit in model.applications[application_name].units: - await unit.scp_to(source, destination, user=user, proxy=proxy, - scp_opts=scp_opts) - scp_func = functools.partial( - _scp_to_all_units, - application_name, - source, - destination, - user=user, - proxy=proxy, - scp_opts=scp_opts) - loop.run( - run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True)) - - -def scp_from_unit(unit_name, model_name, source, destination, user='ubuntu', - proxy=False, scp_opts=''): - """Transfer files from to unit_name in model_name. - - :param unit_name: Name of unit to scp from - :type unit_name: str - :param model_name: Name of model unit is in - :type model_name: str - :param source: Remote path of file(s) to transfer - :type source: str - :param destination: Local destination of transferred files - :type source: str - :param user: Remote username - :type source: str - :param proxy: Proxy through the Juju API server - :type proxy: bool - :param scp_opts: Additional options to the scp command - :type scp_opts: str - """ - async def _scp_from_unit(unit_name, source, destination, user, proxy, - scp_opts, model): - unit = get_unit_from_name(unit_name, model) - await unit.scp_from(source, destination, user=user, proxy=proxy, - scp_opts=scp_opts) - scp_func = functools.partial( - _scp_from_unit, - unit_name, - source, - destination, - user=user, - proxy=proxy, - scp_opts=scp_opts) - loop.run( - run_in_model(model_name, scp_func, add_model_arg=True, awaitable=True)) - - -def get_units(model_name, application_name): - """Return all the units of a given application - - :param model_name: Name of model to query. - :type model_name: str - :param application_name: Name of application to retrieve units for - :type application_name: str - - :returns: List of juju units - :rtype: [juju.unit.Unit, juju.unit.Unit,...] - """ - async def _get_units(application_name, model): - return model.applications[application_name].units - f = functools.partial(_get_units, application_name) - return loop.run(run_in_model(model_name, f, add_model_arg=True)) - - -def get_machines(model_name, application_name): - """Return all the machines of a given application - - :param model_name: Name of model to query. - :type model_name: str - :param application_name: Name of application to retrieve units for - :type application_name: str - - :returns: List of juju machines - :rtype: [juju.machine.Machine, juju.machine.Machine,...] - """ - async def _get_machines(application_name, model): - machines = [] - for unit in model.applications[application_name].units: - machines.append(unit.machine) - return machines - f = functools.partial(_get_machines, application_name) - return loop.run(run_in_model(model_name, f, add_model_arg=True)) - - -def get_first_unit_name(model_name, application_name): - """Return name of lowest numbered unit of given application - - :param model_name: Name of model to query. - :type model_name: str - :param application_name: Name of application - :type application_name: str - - :returns: Name of lowest numbered unit - :rtype: str - """ - return get_units(model_name, application_name)[0].name - - -def get_app_ips(model_name, application_name): - """Return public address of all units of an application - - :param model_name: Name of model to query. - :type model_name: str - :param application_name: Name of application - :type application_name: str - - :returns: List of ip addresses - :rtype: [str, str,...] - """ - return [u.public_address for u in get_units(model_name, application_name)] - def main(): # Run the deploy coroutine in an asyncio event loop, using a helper # that abstracts loop creation and teardown. - print("Current applications: {}".format( ", ".join(loop.run(deployed())))) + print("Current applications: {}".format(", ".join(loop.run(deployed())))) if __name__ == '__main__':