docker_helpers.py

import logging
import os
import re
import json
import copy
import base64
import tarfile
import requests
import time
from io import BytesIO
from pathlib import Path
from contextlib import contextmanager, asynccontextmanager

import aiodocker
from aiodocker.utils import clean_map
from aiodocker.exceptions import DockerError
import docker
from docker.models.services import _get_create_service_kwargs
from docker.types.services import ServiceMode, Resources, EndpointSpec, RestartPolicy, SecretReference

from common.config import config, static
from common.helpers import sint, sfloat

logger = logging.getLogger("UMPIRE")


class DockerBuildError(Exception):
    pass


# TODO: Clean a lot of this up and rectify the inconsistencies between the different docker libraries
class ServiceKwargs:
    @classmethod
    def configure(cls, image, service, secrets=None, mounts=None, **kwargs):
        from compose.utils import timeparse, parse_bytes

        self = ServiceKwargs()
        options = service.options
        deploy_opts = options.get("deploy", {})
        prefs = deploy_opts.get("placement", {}).get("preferences", {})

        # Map compose options to service options
        self.image = image
        self.constraints = deploy_opts.get("placement", {}).get("constraints")
        self.preferences = [kv for pref in prefs for kv in pref.items()]
        self.container_labels = options.get("labels")

        self.endpoint_spec = EndpointSpec(deploy_opts.get("endpoint_mode"),
                                          {p.published: p.target for p in options.get("ports", [])})

        self.env = options.get("environment", None)
        self.hostname = options.get("hostname")
        self.isolation = options.get("isolation")
        self.labels = {k: v for k, v in (kv.split('=') for kv in deploy_opts.get("labels", []))}
        self.log_driver = options.get("logging", {}).get("driver")
        self.log_driver_options = options.get("logging", {}).get("options")
        self.mode = ServiceMode(deploy_opts.get("mode", "replicated"), deploy_opts.get("replicas", 1))
        self.networks = [config.SWARM_NETWORK]  # Similar to mounts. I don't see the use case but see the issues

        resource_opts = deploy_opts.get("resources", {})
        if resource_opts:
            # Unpack any generic_resources defined i.e. gpus and such
            reservation_opts = resource_opts.get("reservations", {})
            generic_resources = {}
            for generic_resource in reservation_opts.get("generic_resources", {}):
                discrete_resource_spec = generic_resource["discrete_resource_spec"]
                generic_resources[discrete_resource_spec["kind"]] = discrete_resource_spec["value"]
            cpu_limit = sfloat(resource_opts.get("limits", {}).get("cpus"), 0)
            cpu_reservation = sfloat(reservation_opts.get("cpus"), 0)
            nano_cpu_limit = sint(cpu_limit * 1e9, 0) if cpu_limit is not None else None
            nano_cpu_reservation = sint(cpu_reservation * 1e9, 0) if cpu_reservation is not None else None
            self.resources = Resources(cpu_limit=nano_cpu_limit,
                                       mem_limit=parse_bytes(resource_opts.get("limits", {}).get("memory", '')),
                                       cpu_reservation=nano_cpu_reservation,
                                       mem_reservation=parse_bytes(reservation_opts.get("memory", '')),
                                       generic_resources=generic_resources)

        restart_opts = deploy_opts.get("restart_policy", {})
        if restart_opts:
            # Parse the restart policy
            delay = timeparse(restart_opts.get("delay", "0s"))
            window = timeparse(restart_opts.get("restart_opts", "0s"))
            self.restart_policy = RestartPolicy(condition=restart_opts.get("condition", ),
                                                delay=delay,
                                                max_attempts=sint(restart_opts.get("max_attempts", 0), 0),
                                                window=window)

        self.secrets = secrets
        self.mounts = mounts

        # Grab any key word arguments that may have been given
        [setattr(self, k, v) for k, v in kwargs.items() if hasattr(self, k)]

        service_kwargs = _get_create_service_kwargs('create', copy.copy(self.__dict__))

        # This is needed because aiodocker assumes the Env is a dictionary for some reason...
        if self.env is not None:
            service_kwargs["task_template"]["ContainerSpec"]["Env"] = self.env

        return service_kwargs


async def create_secret(client, name, data):
    data = base64.b64encode(data)
    data = data.decode("ascii")
    body = {"Data": data, "Name": name}
    headers = {"Content-Type": "application/json"}
    resp = await client._query("secrets/create", "POST", data=json.dumps(body), headers=headers)
    return await resp.json()


async def update_service(client, service_id, *, version=None, image=None, rollback=None, mode=None, force=False):
    if image is None and rollback is False and not force:
        raise ValueError("You need to specify an image.")

    inspect_service = await client.services.inspect(service_id)
    spec = inspect_service["Spec"]

    if mode is not None:
        spec["Mode"] = mode

    if image is not None:
        spec["TaskTemplate"]["ContainerSpec"]["Image"] = image

    if force:
        spec["TaskTemplate"]["ForceUpdate"] = 1

    if version is None:
        params = {"version": inspect_service["Version"]["Index"]}
    else:
        params = {"version": version}

    if rollback is True:
        params["rollback"] = "previous"

    data = json.dumps(clean_map(spec))

    resp = await client._query_json(
        "services/{service_id}/update".format(service_id=service_id),  # Todo: fstring
        method="POST",
        data=data,
        params=params,
    )
    return resp


async def get_secret(client: aiodocker.Docker, secret_id):
    resp = await client._query(f"secrets/{secret_id}")
    return await resp.json()


async def delete_secret(client: aiodocker.Docker, secret_id):
    await client._query(f"secrets/{secret_id}", "DELETE")


async def disconnect_from_network(client: aiodocker.Docker):
    await client._query(f"networks/walkoff_network/disconnect", "POST")


async def get_network(client: aiodocker.Docker, network_id):
    resp = await client._query(f"networks/{network_id}")
    return await resp.json()


async def get_nodes(client: aiodocker.Docker):
    resp = await client._query("nodes")
    return await resp.json()


async def get_tasks(client: aiodocker.Docker, params):
    resp = await client._query("tasks" + '?' + params)
    return await resp.json()


def normalize_name(name, delimiter=''):
    """ Super arbitrary naming convention for docker images/services... """
    return re.sub(r'[^-_a-z0-9]', delimiter, name.lower())


def get_project(path):
    from compose.cli.command import get_project as get_compose_project

    project = get_compose_project(path, environment=load_docker_env(), project_name=static.APP_PREFIX)
    project.path = path  # we'll add this in to refresh the project later
    return project


def load_docker_env():
    # TODO: remove this since it is likely no longer needed
    from compose.config.environment import Environment

    environment = os.environ
    # environment.update({key: val for key, val in config["DOCKER_ENV"].items()})
    return Environment(environment)


async def get_service(docker_client, service_id):
    try:
        s = await docker_client.services.inspect(service_id)
        return {'id': s["ID"], 'version': s['Version']['Index'], 'image': s['Spec']['Labels']['com.docker.stack.image']}
    except DockerError:
        return {}


async def remove_service(docker_client, service):
    try:
        return await docker_client.services.delete(service)
    except DockerError:
        logger.error(f"Could not delete {service}.")
        return False


async def get_replicas(docker_client, service):
    """
    Gets the running and desired replica counts for the given service ID
    :param service: The docker id of the service
    :return: a dictionary giving the number of "running" and "desired" replicas
    """
    tasks = await docker_client.tasks.list(filters={"service": [service]})
    desired = sum([t["DesiredState"] == "running" for t in tasks])
    running = sum([t["Status"]["State"] == "running" for t in tasks])
    return {"running": running, "desired": desired}


async def get_containers(docker_client, service, short_ids=False):
    """
    Gets the running containers the given service ID
    :param service: The docker id of the service
    :return: a set of the running containers
    """

    def get_container_id(task_spec):
        return task_spec["Status"]["ContainerStatus"]["ContainerID"]

    def get_state(task_spec):
        return task_spec["Status"]["State"]

    def has_container(task_spec):
        return task_spec["Status"].get("ContainerStatus") is not None

    tasks = await docker_client.tasks.list(filters={"service": [service]})

    if short_ids:
        return set(get_container_id(t)[:12] for t in tasks if get_state(t) == "running" and has_container(t))
    return set(get_container_id(t) for t in tasks if get_state(t) == "running" and has_container(t))


async def load_volumes(project):
    service = project.services[0]
    volume_references = []
    for service_volume in service.options["volumes"]:
        host_path = service_volume.external
        container_path = service_volume.internal
        mode = service_volume.mode

        mount = host_path + ":" + container_path + ":" + mode
        volume_references.append(mount)

    return volume_references


async def load_secrets(docker_client, project):
    service = project.services[0]
    secret_references = []
    for service_secret in service.secrets:
        secret = service_secret["secret"]
        filename = service_secret.get("file", secret.source)

        # Compose doesn't parse external secrets so we'll assume there is one and build if it doesn't exist
        try:
            secret_id = await get_secret(docker_client, secret.source)

        except (AttributeError, DockerError):
            with open(filename, 'rb') as fp:
                data = fp.read()
            secret_id = (await create_secret(docker_client, name=secret.source, data=data)).get("ID")

        if secret_id is not None:
            secret_references.append(SecretReference(secret_id=secret_id, secret_name=secret.source,
                                                     uid=secret.uid, gid=secret.gid, mode=secret.mode))
    return secret_references


def connect_to_docker():
    client = docker.from_env(environment=load_docker_env())
    try:
        if client.ping():
            logger.debug(f"Connected to Docker Engine: v{client.version()['Version']}")
            return client
    except docker.errors.APIError as e:
        logger.error(f"Docker API error during connect: {e}")


@asynccontextmanager
async def connect_to_aiodocker():
    client = aiodocker.Docker()
    try:

        if (await client._query("_ping")).status == 200:
            resp = await client._query("version")
            version = (await resp.json())["Version"]
            logger.debug(f"Connected to Docker Engine: v{version}")
            yield client
    finally:
        await client.close()
        logger.info("Docker connection closed.")


@contextmanager
def docker_context(path, dirs=None, dockerignore=None):
    """
    Tars and compresses the given docker context in memory. Useful for sending contexts to `docker build` commands.
    :param path: str or pathlib.Path object representing the path of the context
    :param dirs: whitelist of directories under path to grab
    :param dockerignore: blacklist of directories under path to ignore
    :return: an in memory tar of the context
    """

    dockerignore = [] if not dockerignore else dockerignore

    if not isinstance(path, Path):
        try:
            path = Path(path)
        except (ValueError, NotImplementedError):
            logger.exception(f"Error accessing path: \"{path}\"")
            return

    fileobj = BytesIO()
    tar = tarfile.open(fileobj=fileobj, mode="w")

    # If a list of subdirectories is listed, only grab them
    if dirs is not None:
        for d in dirs:
            tar.add(path / d, arcname=d)
    else:
        tar.add(path, arcname='', filter=(lambda ti: None if ti.isdir() and ti.name in dockerignore
                                                             or "__pycache__" in ti.name else ti))
    # tar.list()
    tar.close()
    try:
        fileobj.seek(0)  # must go back to start of file after tarfile writes to it
        yield fileobj
    finally:
        fileobj.close()


async def stream_docker_log(log_stream):
    async for line in log_stream:
        if "stream" in line and line["stream"].strip():
            logger.debug(line["stream"].strip())
        elif "status" in line:
            logger.debug(line["status"].strip())
        elif "error" in line:
            logger.error(line["error"].strip())
            raise DockerBuildError


async def stream_umpire_build_log(log_stream, build_id):
    async for line in log_stream:
        url = f"{config.API_GATEWAY_URI}/walkoff/api/streams/build/build_logger"
        if "stream" in line and line["stream"].strip():
            data = line["stream"].strip()
            body = {"stream": data, "build_status": "building"}
            params = {"build_id": build_id}
            requests.post(url, json=body, params=params)
        elif "status" in line:
            data = line["status"].strip()
            body = {"stream": data, "build_status": "building"}
            params = {"build_id": build_id}
            requests.post(url, json=body, params=params)
        elif "error" in line:
            data = line["error"].strip()
            body = {"stream": data, "build_status": "failure"}
            params = {"build_id": build_id}
            requests.post(url, json=body, params=params)
            raise DockerBuildError
    body = {"stream": " ", "build_status": "Success"}
    params = {"build_id": build_id}
    requests.post(url, json=body, params=params)