ns.py

# A prototype of a library to aid in the development and operation of
# OSM Network Service charms

# This class handles the heavy lifting associated with asyncio.
from juju import controller
import asyncio
import logging
import os
import time
import yaml

# Quiet the debug logging
logging.getLogger("websockets.protocol").setLevel(logging.INFO)
logging.getLogger("juju.client.connection").setLevel(logging.WARN)
logging.getLogger("juju.model").setLevel(logging.WARN)
logging.getLogger("juju.machine").setLevel(logging.WARN)

logger = logging.getLogger(__name__)


class NetworkService:
    """A lightweight interface to the Juju controller.

    This NetworkService client is specifically designed to allow a higher-level
    "NS" charm to interoperate with "VNF" charms, allowing for the execution of
    Primitives across other charms within the same model.
    """

    endpoint = None
    user = "admin"
    secret = None
    port = 17070
    loop = None
    client = None
    model = None
    cacert = None

    def __init__(self, user, secret, endpoint=None):

        self.user = user
        self.secret = secret
        if endpoint is None:
            addresses = os.environ["JUJU_API_ADDRESSES"]
            for address in addresses.split(" "):
                self.endpoint = address
        else:
            self.endpoint = endpoint

        # Stash the name of the model
        self.model = os.environ["JUJU_MODEL_NAME"]

        # Load the ca-cert from agent.conf
        AGENT_PATH = os.path.dirname(os.environ["JUJU_CHARM_DIR"])
        with open("{}/agent.conf".format(AGENT_PATH), "r") as f:
            try:
                y = yaml.safe_load(f)
                self.cacert = y["cacert"]
            except yaml.YAMLError as exc:
                logger.error("Unable to find Juju ca-cert.")
                raise exc

        # Create our event loop
        self.loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.loop)

    async def connect(self):
        """Connect to the Juju controller."""
        controller = controller.Controller()

        logger.debug(
            "Connecting to controller... ws://{}:{} as {}/{}".format(
                self.endpoint,
                self.port,
                self.user,
                self.secret[-4:].rjust(len(self.secret), "*"),
            )
        )
        await controller.connect(
            endpoint=self.endpoint,
            username=self.user,
            password=self.secret,
            cacert=self.cacert,
        )

        return controller

    def __del__(self):
        self.logout()

    async def disconnect(self):
        """Disconnect from the Juju controller."""
        if self.client:
            logger.debug("Disconnecting Juju controller")
            await self.client.disconnect()

    def login(self):
        """Login to the Juju controller."""
        if not self.client:
            # Connect to the Juju API server
            self.client = self.loop.run_until_complete(self.connect())
        return self.client

    def logout(self):
        """Logout of the Juju controller."""

        if self.loop:
            logger.debug("Disconnecting from API")
            self.loop.run_until_complete(self.disconnect())

    def ExecutePrimitiveGetOutput(self, application, primitive, params={}, timeout=600):
        """Execute a single primitive and return it's output.

        This is a blocking method that will execute a single primitive and wait
        for its completion before return it's output.

        :param application str: The application name provided by `GetApplicationName`.
        :param primitive str: The name of the primitive to execute.
        :param params list: A list of parameters.
        :param timeout int: A timeout, in seconds, to wait for the primitive to finish. Defaults to 600 seconds.
        """
        uuid = self.ExecutePrimitive(application, primitive, params)

        status = None
        output = None

        starttime = time.time()
        while time.time() < starttime + timeout:
            status = self.GetPrimitiveStatus(uuid)
            if status in ["completed", "failed"]:
                break
            time.sleep(10)

        # When the primitive is done, get the output
        if status in ["completed", "failed"]:
            output = self.GetPrimitiveOutput(uuid)

        return output

    def ExecutePrimitive(self, application, primitive, params={}):
        """Execute a primitive.

        This is a non-blocking method to execute a primitive. It will return
        the UUID of the queued primitive execution, which you can use
        for subsequent calls to `GetPrimitiveStatus` and `GetPrimitiveOutput`.

        :param application string: The name of the application
        :param primitive string: The name of the Primitive.
        :param params list: A list of parameters.

        :returns uuid string: The UUID of the executed Primitive
        """
        uuid = None

        if not self.client:
            self.login()

        model = self.loop.run_until_complete(self.client.get_model(self.model))

        # Get the application
        if application in model.applications:
            app = model.applications[application]

            # Execute the primitive
            unit = app.units[0]
            if unit:
                action = self.loop.run_until_complete(
                    unit.run_action(primitive, **params)
                )
                uuid = action.id
                logger.debug("Executing action: {}".format(uuid))
            self.loop.run_until_complete(model.disconnect())
        else:
            # Invalid mapping: application not found. Raise exception
            raise Exception("Application not found: {}".format(application))

        return uuid

    def GetPrimitiveStatus(self, uuid):
        """Get the status of a Primitive execution.

        This will return one of the following strings:
        - pending
        - running
        - completed
        - failed

        :param uuid string: The UUID of the executed Primitive.
        :returns: The status of the executed Primitive
        """
        status = None

        if not self.client:
            self.login()

        model = self.loop.run_until_complete(self.client.get_model(self.model))

        status = self.loop.run_until_complete(model.get_action_status(uuid))

        self.loop.run_until_complete(model.disconnect())

        return status[uuid]

    def GetPrimitiveOutput(self, uuid):
        """Get the output of a completed Primitive execution.


        :param uuid string: The UUID of the executed Primitive.
        :returns: The output of the execution, or None if it's still running.
        """
        result = None
        if not self.client:
            self.login()

        model = self.loop.run_until_complete(self.client.get_model(self.model))

        result = self.loop.run_until_complete(model.get_action_output(uuid))

        self.loop.run_until_complete(model.disconnect())

        return result