<commit_before>import base64
import errno
from fcntl import LOCK_EX, LOCK_NB, LOCK_SH, LOCK_UN
import inspect
import logging
import os
import random
import re
import signal
import subprocess
import sys
import time

try:                  # Prefer system installation of Mesos protos if available
    import mesos_pb2 as protos
    import containerizer_pb2 as protos
except:
    import deimos.mesos_pb2 as protos
    import deimos.containerizer_pb2 as protos

import deimos.cgroups
from deimos.cmd import Run
import deimos.config
import deimos.containerizer
import deimos.docker
from deimos.err import Err
import deimos.logger
from deimos.logger import log
import deimos.mesos
import deimos.path
from deimos.proto import recordio
from deimos._struct import _Struct
import deimos.state
import deimos.sig

class Containerizer(object):
    def __init__(self): pass
    def launch(self, *args): pass
    def update(self, *args): pass
    def usage(self, *args): pass
    def wait(self, *args): pass
    def destroy(self, *args): pass
    def __call__(self, *args):
        try:
            name   = args[0]
            method = { "launch"  : self.launch,
                       "update"  : self.update,
                       "usage"   : self.usage,
                       "wait"    : self.wait,
                       "destroy" : self.destroy }[name]
        except IndexError:
            raise Err("Please choose a subcommand")
        except KeyError:
            raise Err("Subcommand %s is not valid for containerizers" % name)
        return method(*args[1:])

def methods():
    "Names of operations provided by containerizers, as a set."
    pairs = inspect.getmembers(Containerizer, predicate=inspect.ismethod)
    return set( k for k, _ in pairs if k[0:1] != "_" )

class Docker(Containerizer, _Struct):
    def __init__(self, workdir="/tmp/mesos-sandbox",
                       state_root="/tmp/deimos",
                       shared_dir="fs",
                       optimistic_unpack=True,
                       container_settings=deimos.config.Containers(),
                       index_settings=deimos.config.DockerIndex()):
        _Struct.__init__(self, workdir=workdir,
                               state_root=state_root,
                               shared_dir=shared_dir,
                               optimistic_unpack=optimistic_unpack,
                               container_settings=container_settings,
                               index_settings=index_settings,
                               runner=None,
                               state=None)
    def launch(self, *args):
        log.info(" ".join(args))
        fork = False if "--no-fork" in args else True
        deimos.sig.install(self.log_signal)
        run_options = []
        proto = recordio.read(protos.Launch)
        launchy = deimos.mesos.Launch(proto)
        state = deimos.state.State(self.state_root,
                                   mesos_id=launchy.container_id)
        state.push()
        lk_l = state.lock("launch", LOCK_EX)
        state.executor_id = launchy.executor_id
        state.push()
        state.ids()
        mesos_directory() # Redundant?
        if launchy.directory:
            os.chdir(launchy.directory)
        # TODO: if launchy.user:
        #           os.seteuid(launchy.user)
        url, options = self.container_settings.override(*launchy.container)
        pre, image = url.split("docker:///")
        if pre != "":
            raise Err("URL '%s' is not a valid docker:// URL!" % url)
        if image == "":
            image = self.default_image(launchy)
        log.info("image  = %s", image)
        run_options += [ "--sig-proxy" ]
        run_options += [ "--rm" ]     # This is how we ensure container cleanup
        run_options += [ "--cidfile", state.resolve("cid") ]

        place_uris(launchy, self.shared_dir, self.optimistic_unpack)
        run_options += [ "-w", self.workdir ]

        # Docker requires an absolute path to a source filesystem, separated
        # from the bind path in the container with a colon, but the absolute
        # path to the Mesos sandbox might have colons in it (TaskIDs with
        # timestamps can cause this situation). So we create a soft link to it
        # and mount that.
        shared_full = os.path.abspath(self.shared_dir)
        sandbox_symlink = state.sandbox_symlink(shared_full)
        run_options += [ "-v", "%s:%s" % (sandbox_symlink, self.workdir) ]

        cpus, mems = launchy.cpu_and_mem
        env = launchy.env
        run_options += options

        # We need to wrap the call to Docker in a call to the Mesos executor
        # if no executor is passed as part of the task. We need to pass the
        # MESOS_* environment variables in to the container if we're going to
        # start an executor.
        observer_argv = None
        if launchy.needs_observer:
            observer_argv = [ mesos_executor(),
                              deimos.path.me(), "wait", "--docker" ]
        else:
            env += mesos_env() + [("MESOS_DIRECTORY", self.workdir)]

        runner_argv = deimos.docker.run(run_options, image, launchy.argv,
                                        env=env, ports=launchy.ports,
                                        cpus=cpus, mems=mems)

        log_mesos_env(logging.DEBUG)

        observer = None
        with open("stdout", "w") as o:        # This awkward multi 'with' is a
            with open("stderr", "w") as e:    # concession to 2.6 compatibility
                with open(os.devnull) as devnull:
                    log.info(deimos.cmd.present(runner_argv))
                    self.runner = subprocess.Popen(runner_argv, stdin=devnull,
                                                                stdout=o,
                                                                stderr=e)
                    state.pid(self.runner.pid)
                    state.await_cid()
                    state.push()
                    lk_w = state.lock("wait", LOCK_EX)
                    lk_l.unlock()
                    state.ids()
                    if observer_argv is not None:
                        observer_argv += [state.cid()]
                        log.info(deimos.cmd.present(observer_argv))
                        call = deimos.cmd.in_sh(observer_argv)
                        # TODO: Collect these leaking file handles.
                        obs_out = open(state.resolve("observer.out"), "w+")
                        obs_err = open(state.resolve("observer.err"), "w+")
                        # If the Mesos executor sees LIBPROCESS_PORT=0 (which
                        # is passed by the slave) there are problems when it
                        # attempts to bind. ("Address already in use").
                        # Purging both LIBPROCESS_* net variables, to be safe.
                        for v in ["LIBPROCESS_PORT", "LIBPROCESS_IP"]:
                            if v in os.environ:
                                del os.environ[v]
                        observer = subprocess.Popen(call, stdin=devnull,
                                                          stdout=obs_out,
                                                          stderr=obs_err,
                                                          close_fds=True)
#                   else:
#                       if fork:
#                           pid = os.fork()
#                           if pid is not 0:
#                               state.ids()
#                               log.info("Forking watcher into child...")
#                               return 1
        data = Run(data=True)(deimos.docker.wait(state.cid()))
        state.exit(data)
        lk_w.unlock()
        for p, arr in [(self.runner, runner_argv), (observer, observer_argv)]:
            if p is None or p.wait() == 0:
                continue
            log.warning(deimos.cmd.present(arr, p.wait()))
        return state.exit()
    def update(self, *args):
        log.info(" ".join(args))
        log.info("Update is a no-op for Docker...")
    def usage(self, *args):
        log.info(" ".join(args))
        message = recordio.read(protos.Usage)
        container_id = message.container_id.value
        state = deimos.state.State(self.state_root, mesos_id=container_id)
        state.await_launch()
        state.ids()
        if state.cid() is None:
            log.info("Container not started?")
            return 0
        if state.exit() is not None:
            log.info("Container is stopped")
            return 0
        cg = deimos.cgroups.CGroups(**deimos.docker.cgroups(state.cid()))
        if len(cg.keys()) == 0:
            log.info("Container has no CGroups...already stopped?")
            return 0
        try:
            recordio.write(protos.ResourceStatistics,
                           timestamp             = time.time(),
                           mem_limit_bytes       = cg.memory.limit(),
                           cpus_limit            = cg.cpu.limit(),
                         # cpus_user_time_secs   = cg.cpuacct.user_time(),
                         # cpus_system_time_secs = cg.cpuacct.system_time(),
                           mem_rss_bytes         = cg.memory.rss())
        except AttributeError as e:
            log.error("Missing CGroup!")
            raise e
        return 0
    def wait(self, *args):
        log.info(" ".join(args))
        if list(args[0:1]) == ["--docker"]:
            # In Docker mode, we use Docker wait to wait for the container
            # and then exit with the returned exit code. The Docker CID is
            # passed on the command line.
            state = deimos.state.State(self.state_root, docker_id=args[1])
        else:
            message = recordio.read(protos.Wait)
            container_id = message.container_id.value
            state = deimos.state.State(self.state_root, mesos_id=container_id)
        self.state = state
        deimos.sig.install(self.stop_docker_and_resume)
        state.await_launch()
        try:
            state.lock("wait", LOCK_SH, seconds=None)
        except IOError as e:                       # Allows for signal recovery
            if e.errno != errno.EINTR:
                raise e
            state.lock("wait", LOCK_SH, 1)
        recordio.write(protos.Termination,
                       killed  = False,
                       message = "",
                       status  = (state.exit() or 255))
        if state.exit() is not None:
            return state.exit()
        raise Err("Wait lock is not held nor is exit file present")
    def destroy(self, *args):
        log.info(" ".join(args))
        message = recordio.read(protos.Destroy)
        container_id = message.container_id.value
        state = deimos.state.State(self.state_root, mesos_id=container_id)
        state.await_launch()
        lk_d = state.lock("destroy", LOCK_EX)
        if state.exit() is None:
            Run()(deimos.docker.stop(state.cid()))
        else:
            log.info("Container is stopped")
        return 0
    def log_signal(self, signum):
        pass
    def stop_docker_and_resume(self, signum):
        if self.state is not None and self.state.cid() is not None:
            cid = self.state.cid()
            log.info("Trying to stop Docker container: %s", cid)
            try:
                Run()(deimos.docker.stop(cid))
            except subprocess.CalledProcessError:
                pass
            return deimos.sig.Resume()
    def default_image(self, launchy):
        opts = dict(self.index_settings.items(onlyset=True))
        if "account_libmesos" in opts:
            if not launchy.needs_observer:
                opts["account"] = opts["account_libmesos"]
            del opts["account_libmesos"]
        return deimos.docker.matching_image_for_host(**opts)

####################################################### Mesos interface helpers

MESOS_ESSENTIAL_ENV = [ "MESOS_SLAVE_ID",     "MESOS_SLAVE_PID",
                        "MESOS_FRAMEWORK_ID", "MESOS_EXECUTOR_ID" ]

def mesos_env():
    env = os.environ.get
    return [ (k, env(k)) for k in MESOS_ESSENTIAL_ENV if env(k) ]

def log_mesos_env(level=logging.INFO):
    for k, v in os.environ.items():
        if k.startswith("MESOS_") or k.startswith("LIBPROCESS_"):
            log.log(level, "%s=%s" % (k, v))

def mesos_directory():
    if not "MESOS_DIRECTORY" in os.environ: return
    work_dir = os.path.abspath(os.getcwd())
    task_dir = os.path.abspath(os.environ["MESOS_DIRECTORY"])
    if task_dir != work_dir:
        log.info("Changing directory to MESOS_DIRECTORY=%s", task_dir)
        os.chdir(task_dir)

def mesos_executor():
    return os.path.join(os.environ["MESOS_LIBEXEC_DIRECTORY"],
                        "mesos-executor")

def mesos_default_image():
    return os.environ.get("MESOS_DEFAULT_CONTAINER_IMAGE")

def place_uris(launchy, directory, optimistic_unpack=False):
    cmd = deimos.cmd.Run()
    cmd(["mkdir", "-p", directory])
    for item in launchy.uris:
        uri = item.value
        gen_unpack_cmd = unpacker(uri) if optimistic_unpack else None
        log.info("Retrieving URI: %s", deimos.cmd.escape([uri]))
        try:
            basename = uri.split("/")[-1]
            f = os.path.join(directory, basename)
            if basename == "":
                raise IndexError
        except IndexError:
            log.info("Not able to determine basename: %r", uri)
            continue
        try:
            cmd(["curl", "-sSfL", uri, "--output", f])
        except subprocess.CalledProcessError as e:
            log.warning("Failed while processing URI: %s",
                        deimos.cmd.escape(uri))
            continue
        if item.executable:
            os.chmod(f, 0755)
        if gen_unpack_cmd is not None:
            log.info("Unpacking %s" % f)
            cmd(gen_unpack_cmd(f, directory))
            cmd(["rm", "-f", f])

def unpacker(uri):
    if re.search(r"[.](t|tar[.])(bz2|xz|gz)$", uri):
        return lambda f, directory: ["tar", "-C", directory, "-xf", f]
    if re.search(r"[.]zip$", uri):
        return lambda f, directory: ["unzip", "-d", directory, f]

<commit_msg>Support xcon6 executor convention.<commit_after>import base64
import errno
from fcntl import LOCK_EX, LOCK_NB, LOCK_SH, LOCK_UN
import inspect
import logging
import os
import random
import re
import signal
import subprocess
import sys
import time

try:                  # Prefer system installation of Mesos protos if available
    import mesos_pb2 as protos
    import containerizer_pb2 as protos
except:
    import deimos.mesos_pb2 as protos
    import deimos.containerizer_pb2 as protos

import deimos.cgroups
from deimos.cmd import Run
import deimos.config
import deimos.containerizer
import deimos.docker
from deimos.err import Err
import deimos.logger
from deimos.logger import log
import deimos.mesos
import deimos.path
from deimos.proto import recordio
from deimos._struct import _Struct
import deimos.state
import deimos.sig

class Containerizer(object):
    def __init__(self): pass
    def launch(self, *args): pass
    def update(self, *args): pass
    def usage(self, *args): pass
    def wait(self, *args): pass
    def destroy(self, *args): pass
    def __call__(self, *args):
        try:
            name   = args[0]
            method = { "launch"  : self.launch,
                       "update"  : self.update,
                       "usage"   : self.usage,
                       "wait"    : self.wait,
                       "destroy" : self.destroy }[name]
        except IndexError:
            raise Err("Please choose a subcommand")
        except KeyError:
            raise Err("Subcommand %s is not valid for containerizers" % name)
        return method(*args[1:])

def methods():
    "Names of operations provided by containerizers, as a set."
    pairs = inspect.getmembers(Containerizer, predicate=inspect.ismethod)
    return set( k for k, _ in pairs if k[0:1] != "_" )

class Docker(Containerizer, _Struct):
    def __init__(self, workdir="/tmp/mesos-sandbox",
                       state_root="/tmp/deimos",
                       shared_dir="fs",
                       optimistic_unpack=True,
                       container_settings=deimos.config.Containers(),
                       index_settings=deimos.config.DockerIndex()):
        _Struct.__init__(self, workdir=workdir,
                               state_root=state_root,
                               shared_dir=shared_dir,
                               optimistic_unpack=optimistic_unpack,
                               container_settings=container_settings,
                               index_settings=index_settings,
                               runner=None,
                               state=None)
    def launch(self, *args):
        log.info(" ".join(args))
        fork = False if "--no-fork" in args else True
        deimos.sig.install(self.log_signal)
        run_options = []
        proto = recordio.read(protos.Launch)
        launchy = deimos.mesos.Launch(proto)
        state = deimos.state.State(self.state_root,
                                   mesos_id=launchy.container_id)
        state.push()
        lk_l = state.lock("launch", LOCK_EX)
        state.executor_id = launchy.executor_id
        state.push()
        state.ids()
        mesos_directory() # Redundant?
        if launchy.directory:
            os.chdir(launchy.directory)
        # TODO: if launchy.user:
        #           os.seteuid(launchy.user)
        url, options = self.container_settings.override(*launchy.container)
        pre, image = url.split("docker:///")
        if pre != "":
            raise Err("URL '%s' is not a valid docker:// URL!" % url)
        if image == "":
            image = self.default_image(launchy)
        log.info("image  = %s", image)
        run_options += [ "--sig-proxy" ]
        run_options += [ "--rm" ]     # This is how we ensure container cleanup
        run_options += [ "--cidfile", state.resolve("cid") ]

        place_uris(launchy, self.shared_dir, self.optimistic_unpack)
        run_options += [ "-w", self.workdir ]

        # Docker requires an absolute path to a source filesystem, separated
        # from the bind path in the container with a colon, but the absolute
        # path to the Mesos sandbox might have colons in it (TaskIDs with
        # timestamps can cause this situation). So we create a soft link to it
        # and mount that.
        shared_full = os.path.abspath(self.shared_dir)
        sandbox_symlink = state.sandbox_symlink(shared_full)
        run_options += [ "-v", "%s:%s" % (sandbox_symlink, self.workdir) ]

        cpus, mems = launchy.cpu_and_mem
        env = launchy.env
        run_options += options

        # We need to wrap the call to Docker in a call to the Mesos executor
        # if no executor is passed as part of the task. We need to pass the
        # MESOS_* environment variables in to the container if we're going to
        # start an executor.
        observer_argv = None
        if launchy.needs_observer:
            # NB: The "@@docker@@" variant is a work around for Mesos's option
            # parser. There is a fix in the pipeline.
            observer_argv = [ mesos_executor(), "--override",
                              deimos.path.me(), "wait", "@@docker@@" ]
        else:
            env += mesos_env() + [("MESOS_DIRECTORY", self.workdir)]

        runner_argv = deimos.docker.run(run_options, image, launchy.argv,
                                        env=env, ports=launchy.ports,
                                        cpus=cpus, mems=mems)

        log_mesos_env(logging.DEBUG)

        observer = None
        with open("stdout", "w") as o:        # This awkward multi 'with' is a
            with open("stderr", "w") as e:    # concession to 2.6 compatibility
                with open(os.devnull) as devnull:
                    log.info(deimos.cmd.present(runner_argv))
                    self.runner = subprocess.Popen(runner_argv, stdin=devnull,
                                                                stdout=o,
                                                                stderr=e)
                    state.pid(self.runner.pid)
                    state.await_cid()
                    state.push()
                    lk_w = state.lock("wait", LOCK_EX)
                    lk_l.unlock()
                    state.ids()
                    if observer_argv is not None:
                        observer_argv += [state.cid()]
                        log.info(deimos.cmd.present(observer_argv))
                        call = deimos.cmd.in_sh(observer_argv)
                        # TODO: Collect these leaking file handles.
                        obs_out = open(state.resolve("observer.out"), "w+")
                        obs_err = open(state.resolve("observer.err"), "w+")
                        # If the Mesos executor sees LIBPROCESS_PORT=0 (which
                        # is passed by the slave) there are problems when it
                        # attempts to bind. ("Address already in use").
                        # Purging both LIBPROCESS_* net variables, to be safe.
                        for v in ["LIBPROCESS_PORT", "LIBPROCESS_IP"]:
                            if v in os.environ:
                                del os.environ[v]
                        observer = subprocess.Popen(call, stdin=devnull,
                                                          stdout=obs_out,
                                                          stderr=obs_err,
                                                          close_fds=True)
#                   else:
#                       if fork:
#                           pid = os.fork()
#                           if pid is not 0:
#                               state.ids()
#                               log.info("Forking watcher into child...")
#                               return 1
        data = Run(data=True)(deimos.docker.wait(state.cid()))
        state.exit(data)
        lk_w.unlock()
        for p, arr in [(self.runner, runner_argv), (observer, observer_argv)]:
            if p is None or p.wait() == 0:
                continue
            log.warning(deimos.cmd.present(arr, p.wait()))
        return state.exit()
    def update(self, *args):
        log.info(" ".join(args))
        log.info("Update is a no-op for Docker...")
    def usage(self, *args):
        log.info(" ".join(args))
        message = recordio.read(protos.Usage)
        container_id = message.container_id.value
        state = deimos.state.State(self.state_root, mesos_id=container_id)
        state.await_launch()
        state.ids()
        if state.cid() is None:
            log.info("Container not started?")
            return 0
        if state.exit() is not None:
            log.info("Container is stopped")
            return 0
        cg = deimos.cgroups.CGroups(**deimos.docker.cgroups(state.cid()))
        if len(cg.keys()) == 0:
            log.info("Container has no CGroups...already stopped?")
            return 0
        try:
            recordio.write(protos.ResourceStatistics,
                           timestamp             = time.time(),
                           mem_limit_bytes       = cg.memory.limit(),
                           cpus_limit            = cg.cpu.limit(),
                         # cpus_user_time_secs   = cg.cpuacct.user_time(),
                         # cpus_system_time_secs = cg.cpuacct.system_time(),
                           mem_rss_bytes         = cg.memory.rss())
        except AttributeError as e:
            log.error("Missing CGroup!")
            raise e
        return 0
    def wait(self, *args):
        log.info(" ".join(args))
        # NB: The "@@docker@@" variant is a work around for Mesos's option
        # parser. There is a fix in the pipeline.
        if list(args[0:1]) in [ ["--docker"], ["@@docker@@"] ]:
            # In Docker mode, we use Docker wait to wait for the container
            # and then exit with the returned exit code. The Docker CID is
            # passed on the command line.
            state = deimos.state.State(self.state_root, docker_id=args[1])
        else:
            message = recordio.read(protos.Wait)
            container_id = message.container_id.value
            state = deimos.state.State(self.state_root, mesos_id=container_id)
        self.state = state
        deimos.sig.install(self.stop_docker_and_resume)
        state.await_launch()
        try:
            state.lock("wait", LOCK_SH, seconds=None)
        except IOError as e:                       # Allows for signal recovery
            if e.errno != errno.EINTR:
                raise e
            state.lock("wait", LOCK_SH, 1)
        recordio.write(protos.Termination,
                       killed  = False,
                       message = "",
                       status  = (state.exit() or 255))
        if state.exit() is not None:
            return state.exit()
        raise Err("Wait lock is not held nor is exit file present")
    def destroy(self, *args):
        log.info(" ".join(args))
        message = recordio.read(protos.Destroy)
        container_id = message.container_id.value
        state = deimos.state.State(self.state_root, mesos_id=container_id)
        state.await_launch()
        lk_d = state.lock("destroy", LOCK_EX)
        if state.exit() is None:
            Run()(deimos.docker.stop(state.cid()))
        else:
            log.info("Container is stopped")
        return 0
    def log_signal(self, signum):
        pass
    def stop_docker_and_resume(self, signum):
        if self.state is not None and self.state.cid() is not None:
            cid = self.state.cid()
            log.info("Trying to stop Docker container: %s", cid)
            try:
                Run()(deimos.docker.stop(cid))
            except subprocess.CalledProcessError:
                pass
            return deimos.sig.Resume()
    def default_image(self, launchy):
        opts = dict(self.index_settings.items(onlyset=True))
        if "account_libmesos" in opts:
            if not launchy.needs_observer:
                opts["account"] = opts["account_libmesos"]
            del opts["account_libmesos"]
        return deimos.docker.matching_image_for_host(**opts)

####################################################### Mesos interface helpers

MESOS_ESSENTIAL_ENV = [ "MESOS_SLAVE_ID",     "MESOS_SLAVE_PID",
                        "MESOS_FRAMEWORK_ID", "MESOS_EXECUTOR_ID" ]

def mesos_env():
    env = os.environ.get
    return [ (k, env(k)) for k in MESOS_ESSENTIAL_ENV if env(k) ]

def log_mesos_env(level=logging.INFO):
    for k, v in os.environ.items():
        if k.startswith("MESOS_") or k.startswith("LIBPROCESS_"):
            log.log(level, "%s=%s" % (k, v))

def mesos_directory():
    if not "MESOS_DIRECTORY" in os.environ: return
    work_dir = os.path.abspath(os.getcwd())
    task_dir = os.path.abspath(os.environ["MESOS_DIRECTORY"])
    if task_dir != work_dir:
        log.info("Changing directory to MESOS_DIRECTORY=%s", task_dir)
        os.chdir(task_dir)

def mesos_executor():
    return os.path.join(os.environ["MESOS_LIBEXEC_DIRECTORY"],
                        "mesos-executor")

def mesos_default_image():
    return os.environ.get("MESOS_DEFAULT_CONTAINER_IMAGE")

def place_uris(launchy, directory, optimistic_unpack=False):
    cmd = deimos.cmd.Run()
    cmd(["mkdir", "-p", directory])
    for item in launchy.uris:
        uri = item.value
        gen_unpack_cmd = unpacker(uri) if optimistic_unpack else None
        log.info("Retrieving URI: %s", deimos.cmd.escape([uri]))
        try:
            basename = uri.split("/")[-1]
            f = os.path.join(directory, basename)
            if basename == "":
                raise IndexError
        except IndexError:
            log.info("Not able to determine basename: %r", uri)
            continue
        try:
            cmd(["curl", "-sSfL", uri, "--output", f])
        except subprocess.CalledProcessError as e:
            log.warning("Failed while processing URI: %s",
                        deimos.cmd.escape(uri))
            continue
        if item.executable:
            os.chmod(f, 0755)
        if gen_unpack_cmd is not None:
            log.info("Unpacking %s" % f)
            cmd(gen_unpack_cmd(f, directory))
            cmd(["rm", "-f", f])

def unpacker(uri):
    if re.search(r"[.](t|tar[.])(bz2|xz|gz)$", uri):
        return lambda f, directory: ["tar", "-C", directory, "-xf", f]
    if re.search(r"[.]zip$", uri):
        return lambda f, directory: ["unzip", "-d", directory, f]

<|endoftext|>


<commit_before>"""
Unit tests for format checking
"""

from __future__ import print_function

from nose.plugins.skip import SkipTest

import os
import subprocess

import cleverhans
from cleverhans.devtools.tests.docscrape import docstring_errors
from cleverhans.devtools.list_files import list_files
from cleverhans.utils import shell_call

# Enter a manual list of files that are allowed to violate PEP8 here
whitelist_pep8 = [
    # This file is broken but could be fixed
    "../examples/multigpu_advtrain/test_attack_multigpu.py"
]

# We don't want to test RL-attack because it has so many dependencies
# not used elsewhere, and pylint wants to import them all
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py") if "RL-attack" in path])
# Similarly, we don't want to require robust_vision_benchmark installed
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "robust_vision_benchmark" in path])
# Similarly, we don't want to require that cloud be installed
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "cloud_client" in path])
# This example has more dependencies too
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "facenet_adversarial_faces" in path])
# This too
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "madry_lab_challenges" in path])

whitelist_docstrings = [
]

def test_format_pep8():
  """
  Test if pep8 is respected.
  """
  files_to_check = []
  module_dir = cleverhans.__path__[0]
  for path in list_files(".py"):
    rel_path = os.path.relpath(path, module_dir)
    if rel_path in whitelist_pep8:
      continue
    else:
      files_to_check.append(path)
  repo_dir = os.path.join(module_dir, os.pardir)
  rcpath = os.path.join(repo_dir, '.pylintrc')
  assert os.path.exists(rcpath)

  # We must run pylint via the command line and subprocess because of
  # problems with the pylint module.
  # The documentation claims you can run it as a python module, but
  # the documentation is wrong: https://github.com/PyCQA/pylint/issues/1870
  # If you run the version described in the linked issue, pylint
  # calls sys.exit once it is done, so it kills the test.

  try:
    shell_call(['pylint', '--rcfile', rcpath] + files_to_check)
  except subprocess.CalledProcessError as e:
    raise ValueError(e.output.decode("utf-8"))

def print_files_information_pep8():
  """
  Print the list of files which can be removed from the whitelist and the
  list of files which do not respect PEP8 formatting that aren't in the
  whitelist
  """

  raise NotImplementedError("Broken by move to TensorFlow style")

  infracting_files = []
  non_infracting_files = []
  # We must replace StyleGuide with pylint, as is done in the tests
  pep8_checker = StyleGuide(quiet=True)
  for path in list_files(".py"):
    number_of_infractions = pep8_checker.input_file(path)
    rel_path = os.path.relpath(path, cleverhans.__path__[0])
    if number_of_infractions > 0:
      if rel_path not in whitelist_pep8:
        infracting_files.append(path)
    else:
      if rel_path in whitelist_pep8:
        non_infracting_files.append(path)
  print("Files that must be corrected or added to whitelist:")
  for file in infracting_files:
    print(file)
  print("Files that can be removed from whitelist:")
  for file in non_infracting_files:
    print(file)

def test_format_docstrings():
  """
  Test if docstrings are well formatted.
  """
  # Disabled for now
  return True

  try:
    verify_format_docstrings()
  except SkipTest as e:
    import traceback
    traceback.print_exc(e)
    raise AssertionError(
        "Some file raised SkipTest on import, and inadvertently"
        " canceled the documentation testing."
    )

def verify_format_docstrings():
  """
  Implementation of `test_format_docstrings`. The implementation is
  factored out so it can be placed inside a guard against SkipTest.
  """
  format_infractions = []

  for path in list_files(".py"):
    rel_path = os.path.relpath(path, cleverhans.__path__[0])
    if rel_path in whitelist_docstrings:
      continue
    try:
      format_infractions.extend(docstring_errors(path))
    except Exception as e:
      format_infractions.append(["%s failed to run so format cannot "
                                 "be checked. Error message:\n %s" %
                                 (rel_path, e)])

  if len(format_infractions) > 0:
    msg = "\n".join(':'.join(line) for line in format_infractions)
    raise AssertionError("Docstring format not respected:\n%s" % msg)

if __name__ == "__main__":
  print_files_information_pep8()
<commit_msg>use multiple threads<commit_after>"""
Unit tests for format checking
"""

from __future__ import print_function

from nose.plugins.skip import SkipTest

import os
import subprocess

import cleverhans
from cleverhans.devtools.tests.docscrape import docstring_errors
from cleverhans.devtools.list_files import list_files
from cleverhans.utils import shell_call

# Enter a manual list of files that are allowed to violate PEP8 here
whitelist_pep8 = [
    # This file is broken but could be fixed
    "../examples/multigpu_advtrain/test_attack_multigpu.py"
]

# We don't want to test RL-attack because it has so many dependencies
# not used elsewhere, and pylint wants to import them all
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py") if "RL-attack" in path])
# Similarly, we don't want to require robust_vision_benchmark installed
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "robust_vision_benchmark" in path])
# Similarly, we don't want to require that cloud be installed
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "cloud_client" in path])
# This example has more dependencies too
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "facenet_adversarial_faces" in path])
# This too
whitelist_pep8.extend([os.path.relpath(path, cleverhans.__path__[0])
                       for path in list_files(".py")
                       if "madry_lab_challenges" in path])

whitelist_docstrings = [
]

def test_format_pep8():
  """
  Test if pep8 is respected.
  """
  files_to_check = []
  module_dir = cleverhans.__path__[0]
  for path in list_files(".py"):
    rel_path = os.path.relpath(path, module_dir)
    if rel_path in whitelist_pep8:
      continue
    else:
      files_to_check.append(path)
  repo_dir = os.path.join(module_dir, os.pardir)
  rcpath = os.path.join(repo_dir, '.pylintrc')
  assert os.path.exists(rcpath)

  # We must run pylint via the command line and subprocess because of
  # problems with the pylint module.
  # The documentation claims you can run it as a python module, but
  # the documentation is wrong: https://github.com/PyCQA/pylint/issues/1870
  # If you run the version described in the linked issue, pylint
  # calls sys.exit once it is done, so it kills the test.

  # Running all files in one pylint command is important for 2 reasons:
  # 1) Correctness: pylint can detect issues that require access to multiple
  #    files, such as cyclic imports
  # 2) Speed: pylint imports modules for deep analysis, so if you run
  #    multiple subprocesses each needs to re-import tensorflow.
  # On Ian's laptop, pylint takes about 10s per file to run on the repo,
  # and there are about 90 files as of the writing of this comment.
  # Running pylint on all files simultaneously takes about 70s, so it
  # is a little better than a 10X speedup.

  # Running multiple jobs in parallel helps but far less than linearly.
  # On Ian's 4-core laptop, running 4 jobs drops the runtime from 70s
  # to 45s.
  # Some of the work is I/O, so it actually makes some sense to run
  # more jobs than cores. On Ian's 4-core laptop, running 8 jobs drops
  # the runtime to 40s.
  # There's a further complication though: I think each job needs to
  # redo imports, so the total amount of work to do increases with
  # the number of jobs. On Ian's laptop, using 64 jobs causes the
  # runtime to increase to 220s. There is not an obvious simple
  # formula like "use one job per CPU core" or "use way more jobs
  # than cores to saturate I/O". For now I'm hoping that 8 will be
  # a reasonable default: it gets good performance on my laptop,
  # and on machines with fewer than 4 cores there should still be
  # a benefit to not being blocked on I/O.

  try:
    shell_call(['pylint', '--rcfile', rcpath, '--jobs', '8'] + files_to_check)
  except subprocess.CalledProcessError as e:
    raise ValueError(e.output.decode("utf-8"))

def print_files_information_pep8():
  """
  Print the list of files which can be removed from the whitelist and the
  list of files which do not respect PEP8 formatting that aren't in the
  whitelist
  """

  raise NotImplementedError("Broken by move to TensorFlow style")

  infracting_files = []
  non_infracting_files = []
  # We must replace StyleGuide with pylint, as is done in the tests
  pep8_checker = StyleGuide(quiet=True)
  for path in list_files(".py"):
    number_of_infractions = pep8_checker.input_file(path)
    rel_path = os.path.relpath(path, cleverhans.__path__[0])
    if number_of_infractions > 0:
      if rel_path not in whitelist_pep8:
        infracting_files.append(path)
    else:
      if rel_path in whitelist_pep8:
        non_infracting_files.append(path)
  print("Files that must be corrected or added to whitelist:")
  for file in infracting_files:
    print(file)
  print("Files that can be removed from whitelist:")
  for file in non_infracting_files:
    print(file)

def test_format_docstrings():
  """
  Test if docstrings are well formatted.
  """
  # Disabled for now
  return True

  try:
    verify_format_docstrings()
  except SkipTest as e:
    import traceback
    traceback.print_exc(e)
    raise AssertionError(
        "Some file raised SkipTest on import, and inadvertently"
        " canceled the documentation testing."
    )

def verify_format_docstrings():
  """
  Implementation of `test_format_docstrings`. The implementation is
  factored out so it can be placed inside a guard against SkipTest.
  """
  format_infractions = []

  for path in list_files(".py"):
    rel_path = os.path.relpath(path, cleverhans.__path__[0])
    if rel_path in whitelist_docstrings:
      continue
    try:
      format_infractions.extend(docstring_errors(path))
    except Exception as e:
      format_infractions.append(["%s failed to run so format cannot "
                                 "be checked. Error message:\n %s" %
                                 (rel_path, e)])

  if len(format_infractions) > 0:
    msg = "\n".join(':'.join(line) for line in format_infractions)
    raise AssertionError("Docstring format not respected:\n%s" % msg)

if __name__ == "__main__":
  print_files_information_pep8()
<|endoftext|>


<commit_before># Copyright 2018 PerfKitBenchmarker Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runs Glibc Microbenchmark.

The glibc microbenchmark suite automatically generates code for specified
functions, builds and calls them repeatedly for given inputs to give some
basic performance properties of the function.

Homepage: https://fossies.org/linux/glibc/benchtests/README

Installs glibc-2.27. The benchmark needs python 2.7 or later in addition to the
dependencies required to build the GNU C Library.
"""

import json
import logging

from perfkitbenchmarker import configs
from perfkitbenchmarker import flags
from perfkitbenchmarker import linux_packages
from perfkitbenchmarker import sample

FLAGS = flags.FLAGS

BENCHMARK_NAME = 'glibc'
BENCHMARK_CONFIG = """
glibc:
  description: Runs Glibc Microbenchmark.
  vm_groups:
    default:
      vm_spec: *default_dual_core
      vm_count: null
"""

glibc_default_benchset = ['bench-math',
                          'bench-pthread',
                          'bench-string',
                          'string-benchset',
                          'wcsmbs-benchset',
                          'stdlib-benchset',
                          'stdio-common-benchset',
                          'math-benchset',
                          'malloc-thread',]
flags.DEFINE_multi_enum(
    'glibc_benchset', glibc_default_benchset, glibc_default_benchset,
    'By default, it will run the whole set of benchmarks. To run only a subset '
    'of benchmarks, one may set "glibc_benchset = bench-math bench-pthread" by '
    'using the flag on the command line multiple times.')

GLIBC_BENCH = ['bench-math', 'bench-pthread', 'bench-string']
GLIBC_BENCH_MALLOC = ['malloc-thread']
# TODO(yanfeiren): Parse other *-benchset benchmarks.
GLIBC_MATH_BENCHSET = ['math-benchset']

RESULTS_DIR = '%s/glibc/glibc-build/benchtests' % linux_packages.INSTALL_DIR

def GetConfig(user_config):
  return configs.LoadConfig(BENCHMARK_CONFIG, user_config, BENCHMARK_NAME)

def PrepareGlibc(vm):
  """Building glibc on a single vm."""
  logging.info('Installing Glibc on %s', vm)
  vm.Install('glibc')

def Prepare(benchmark_spec):
  """Install glibc on the target vms.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  PrepareGlibc(vm)

def UpdateMetadata(metadata):
  """Update metadata with glibc-related flag values."""
  metadata['gcc'] = '5.5.0-12'
  metadata['glibc_benchset'] = FLAGS.glibc_benchset

# The reason to write this helper function is that glibc_benchset_output.json
# contains duplicate keys.
def HelperParseOutput(lst):
  """Self defined parameter function for 'object_pairs_hook' of json.loads().

  Purpose:
      This function helps to parse json text with duplicate keys, for instance:
      json text: {"key":1, "key":2, "key2":3, "key2":4}
      With this function, we will have the following dict:
      {"key":[1, 2],
       "key2":[3, 4]
      }

  Args:
    lst: A list of tuples contains the output of benchmark tests.
  Returns:
    A dict contains the output of benchmark tests.
  """

  result = {}
  for key, val in lst:
    if key in result:
      if isinstance(result[key], list):
        result[key].append(val)
      else:
        result[key] = [result[key], val]
    else:
      result[key] = val
  return result

def ParseOutput(glibc_output, benchmark_spec, upper_key, results):
  """Parses the output from glibc.

  Args:
    glibc_output: A json format string containing the output of benchmark tests.
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
    upper_key: The first dimension key of the glibc_output dict.
    results:
      A list to which this function will append new samples based on the glibc
      output. (in the same format as Run() returns).
  """

  metadata = dict()
  metadata['num_machines'] = len(benchmark_spec.vms)
  UpdateMetadata(metadata)

  jsondata = json.loads(glibc_output, object_pairs_hook=HelperParseOutput)
  for function, items in jsondata[upper_key].iteritems():
    # handle the jsondata with duplicate keys
    if isinstance(items, list):
      for item in items:
        current_metadata = metadata.copy()
        for key, val in item.iteritems():
          metric = '{0}:{1}'.format(function, key)
          for subitem, value in val.iteritems():
            current_metadata[subitem] = value
          results.append(sample.Sample(metric, -1, '', current_metadata))
    # handle the jsondata with unique keys
    else:
      for item in items:
        current_metadata = metadata.copy()
        metric = '{0}:{1}'.format(function, item)
        for subitem, value in items[item].iteritems():
          current_metadata[subitem] = value
        results.append(sample.Sample(metric, -1, '', current_metadata))

def Run(benchmark_spec):
  """Run Glibc on the cluster.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  Returns:
    A list of sample.Sample objects.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  results = []

  glibc_user_benchset = ' '.join(FLAGS.glibc_benchset)
  vm.RobustRemoteCommand('cd %s/glibc/glibc-build && '
                         'make bench BENCHSET="%s"' % (
                             linux_packages.INSTALL_DIR, glibc_user_benchset))

  logging.info('Glibc Benchmark Tests Results:')
  # Parse the output for "bench-math", "bench-string" and "bench-pthread".
  if any(i in GLIBC_BENCH for i in FLAGS.glibc_benchset):
    stdout, _ = vm.RemoteCommand(
        'cat {0}/bench.out'.format(RESULTS_DIR), should_log=True)
    ParseOutput(stdout, benchmark_spec, 'functions', results)
  # Parse the output for "malloc-thread".
  if any(i in GLIBC_BENCH_MALLOC for i in FLAGS.glibc_benchset):
    thread_num = ['1', '8', '16', '32']
    for num in thread_num:
      stdout, _ = vm.RemoteCommand(
          'cat {0}/bench-malloc-thread-{1}.out'.format(RESULTS_DIR, num),
          should_log=True)
      ParseOutput(stdout, benchmark_spec, 'functions', results)
  # Parse the output for "math-benchset".
  if any(i in GLIBC_MATH_BENCHSET for i in FLAGS.glibc_benchset):
    stdout, _ = vm.RemoteCommand(
        'cat {0}/bench-math-inlines.out'.format(RESULTS_DIR), should_log=True)
    ParseOutput('{%s}' % stdout, benchmark_spec, 'math-inlines', results)
  return results

def Cleanup(unused_benchmark_spec):
  """Cleanup Glibc on the cluster."""
  pass

<commit_msg>Fix tox issues.<commit_after># Copyright 2018 PerfKitBenchmarker Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Runs Glibc Microbenchmark.

The glibc microbenchmark suite automatically generates code for specified
functions, builds and calls them repeatedly for given inputs to give some
basic performance properties of the function.

Homepage: https://fossies.org/linux/glibc/benchtests/README

Installs glibc-2.27. The benchmark needs python 2.7 or later in addition to the
dependencies required to build the GNU C Library.
"""

import json
import logging

from perfkitbenchmarker import configs
from perfkitbenchmarker import flags
from perfkitbenchmarker import linux_packages
from perfkitbenchmarker import sample

FLAGS = flags.FLAGS

BENCHMARK_NAME = 'glibc'
BENCHMARK_CONFIG = """
glibc:
  description: Runs Glibc Microbenchmark.
  vm_groups:
    default:
      vm_spec: *default_dual_core
      vm_count: null
"""

glibc_default_benchset = ['bench-math',
                          'bench-pthread',
                          'bench-string',
                          'string-benchset',
                          'wcsmbs-benchset',
                          'stdlib-benchset',
                          'stdio-common-benchset',
                          'math-benchset',
                          'malloc-thread']
flags.DEFINE_multi_enum(
    'glibc_benchset', glibc_default_benchset, glibc_default_benchset,
    'By default, it will run the whole set of benchmarks. To run only a subset '
    'of benchmarks, one may set "glibc_benchset = bench-math bench-pthread" by '
    'using the flag on the command line multiple times.')

GLIBC_BENCH = ['bench-math', 'bench-pthread', 'bench-string']
GLIBC_BENCH_MALLOC = ['malloc-thread']
# TODO(yanfeiren): Parse other *-benchset benchmarks.
GLIBC_MATH_BENCHSET = ['math-benchset']

RESULTS_DIR = '%s/glibc/glibc-build/benchtests' % linux_packages.INSTALL_DIR

def GetConfig(user_config):
  return configs.LoadConfig(BENCHMARK_CONFIG, user_config, BENCHMARK_NAME)

def PrepareGlibc(vm):
  """Building glibc on a single vm."""
  logging.info('Installing Glibc on %s', vm)
  vm.Install('glibc')

def Prepare(benchmark_spec):
  """Install glibc on the target vms.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  PrepareGlibc(vm)

def UpdateMetadata(metadata):
  """Update metadata with glibc-related flag values."""
  metadata['gcc'] = '5.5.0-12'
  metadata['glibc_benchset'] = FLAGS.glibc_benchset

# The reason to write this helper function is that glibc_benchset_output.json
# contains duplicate keys.
def HelperParseOutput(lst):
  """Self defined parameter function for 'object_pairs_hook' of json.loads().

  Purpose:
      This function helps to parse json text with duplicate keys, for instance:
      json text: {"key":1, "key":2, "key2":3, "key2":4}
      With this function, we will have the following dict:
      {"key":[1, 2],
       "key2":[3, 4]
      }

  Args:
    lst: A list of tuples contains the output of benchmark tests.
  Returns:
    A dict contains the output of benchmark tests.
  """

  result = {}
  for key, val in lst:
    if key in result:
      if isinstance(result[key], list):
        result[key].append(val)
      else:
        result[key] = [result[key], val]
    else:
      result[key] = val
  return result

def ParseOutput(glibc_output, benchmark_spec, upper_key, results):
  """Parses the output from glibc.

  Args:
    glibc_output: A json format string containing the output of benchmark tests.
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.
    upper_key: The first dimension key of the glibc_output dict.
    results:
      A list to which this function will append new samples based on the glibc
      output. (in the same format as Run() returns).
  """

  metadata = dict()
  metadata['num_machines'] = len(benchmark_spec.vms)
  UpdateMetadata(metadata)

  jsondata = json.loads(glibc_output, object_pairs_hook=HelperParseOutput)
  for function, items in jsondata[upper_key].iteritems():
    # handle the jsondata with duplicate keys
    if isinstance(items, list):
      for item in items:
        current_metadata = metadata.copy()
        for key, val in item.iteritems():
          metric = '{0}:{1}'.format(function, key)
          for subitem, value in val.iteritems():
            current_metadata[subitem] = value
          results.append(sample.Sample(metric, -1, '', current_metadata))
    # handle the jsondata with unique keys
    else:
      for item in items:
        current_metadata = metadata.copy()
        metric = '{0}:{1}'.format(function, item)
        for subitem, value in items[item].iteritems():
          current_metadata[subitem] = value
        results.append(sample.Sample(metric, -1, '', current_metadata))

def Run(benchmark_spec):
  """Run Glibc on the cluster.

  Args:
    benchmark_spec: The benchmark specification. Contains all data that is
        required to run the benchmark.

  Returns:
    A list of sample.Sample objects.
  """
  vms = benchmark_spec.vms
  vm = vms[0]
  results = []

  glibc_user_benchset = ' '.join(FLAGS.glibc_benchset)
  vm.RobustRemoteCommand('cd %s/glibc/glibc-build && '
                         'make bench BENCHSET="%s"' % (
                             linux_packages.INSTALL_DIR, glibc_user_benchset))

  logging.info('Glibc Benchmark Tests Results:')
  # Parse the output for "bench-math", "bench-string" and "bench-pthread".
  if any(i in GLIBC_BENCH for i in FLAGS.glibc_benchset):
    stdout, _ = vm.RemoteCommand(
        'cat {0}/bench.out'.format(RESULTS_DIR), should_log=True)
    ParseOutput(stdout, benchmark_spec, 'functions', results)
  # Parse the output for "malloc-thread".
  if any(i in GLIBC_BENCH_MALLOC for i in FLAGS.glibc_benchset):
    thread_num = ['1', '8', '16', '32']
    for num in thread_num:
      stdout, _ = vm.RemoteCommand(
          'cat {0}/bench-malloc-thread-{1}.out'.format(RESULTS_DIR, num),
          should_log=True)
      ParseOutput(stdout, benchmark_spec, 'functions', results)
  # Parse the output for "math-benchset".
  if any(i in GLIBC_MATH_BENCHSET for i in FLAGS.glibc_benchset):
    stdout, _ = vm.RemoteCommand(
        'cat {0}/bench-math-inlines.out'.format(RESULTS_DIR), should_log=True)
    ParseOutput('{%s}' % stdout, benchmark_spec, 'math-inlines', results)
  return results

def Cleanup(unused_benchmark_spec):
  """Cleanup Glibc on the cluster."""
  pass
<|endoftext|>


<commit_before># Copyright (C) 2014 Universidad Politecnica de Madrid
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

import logging

from django.conf import settings
from django.core.urlresolvers import reverse
from django.utils.translation import ugettext_lazy as _

from horizon import exceptions
from horizon import forms
from horizon import messages
from horizon import workflows

from openstack_dashboard import api
from openstack_dashboard import fiware_api
from openstack_dashboard.dashboards.idm import utils as idm_utils

INDEX_URL = "horizon:idm:organizations:index"
PROJECT_USER_MEMBER_SLUG = "update_members"
LOG = logging.getLogger('idm_logger')

class RolesMixin(object):
    """Simple container for the roles logic."""
    # TODO(garcianavalon) maybe move to fiware_api
    def list_all_roles(self, request, project_id):
        role_list = {}
        # NOTE(garcianavalon) list all roles grouped by application
        # on which current user has the right to get and assign
        role_list['applications'] = fiware_api.keystone.list_allowed_roles_to_assign(request,
                                                user=request.user.id,
                                                organization=request.user.project_id)
        # NOTE(garcianavalon) we also need the organization (keystone)
        # roles here to add members
        role_list['organization'] = api.keystone.role_list(request)
        return role_list

    def list_users_with_roles(self, request, project_id, available_roles):
        # First, load all users from the organization
        # with their roles
        project_users_roles = api.keystone.get_project_users_roles(request,
                                                              project_id)
        # Second, load all the organization-scoped application roles for every user
        # but only the ones the user can assign (and only scoped for the current
        # organization)
        for user_id in project_users_roles:
            project_users_roles[user_id] += [
                r.id for r in fiware_api.keystone.role_list(
                                                    request,
                                                    user=user_id,
                                                    organization=project_id)
                    if r in available_roles
            ]
        return project_users_roles

class UpdateProjectMembersAction(workflows.MembershipAction, RolesMixin):
    def __init__(self, request, *args, **kwargs):
        super(UpdateProjectMembersAction, self).__init__(request,
                                                         *args,
                                                         **kwargs)
        err_msg = _('Unable to retrieve user list. Please try again later.')
        project_id = self.initial['project_id']

        # Get the default role
        try:
            default_role = api.keystone.get_default_role(self.request)
            # Default role is necessary to add members to a project
            if default_role is None:
                default = getattr(settings,
                                  "OPENSTACK_KEYSTONE_DEFAULT_ROLE", None)
                msg = (_('Could not find default role "%s" in Keystone') %
                       default)
                raise exceptions.NotFound(msg)
        except Exception:
            exceptions.handle(self.request,
                              err_msg,
                              redirect=reverse(INDEX_URL))
        # TODO(garcianavalon) hide the default role, show the 'owner'
        # role under the organization ('owner' is 'admin' for now)
        default_role_name = self.get_default_role_field_name()
        self.fields[default_role_name] = forms.CharField(required=False)
        self.fields[default_role_name].initial = default_role.id

        # Get list of available users
        try:
            all_users = api.keystone.user_list(request)
        except Exception:
            exceptions.handle(request, err_msg)
        users_list = [(user.id, user.name) for user in all_users]

        # Get list of roles
        try:
            role_list = self.list_all_roles(request, project_id)
        except Exception:
            exceptions.handle(request,
                              err_msg,
                              redirect=reverse(INDEX_URL))
        for k in role_list:
            # TODO(garcianavalon) different behaviour
            # TODO(garcianavalon) hide the default role, show the 'owner'
            # role under the organization ('owner' is 'admin' for now)
            for role in role_list[k]:
                field_name = self.get_member_field_name(role.id)
                label = role.name
                widget = forms.widgets.SelectMultiple(
                                attrs={'data-application-name':'patata'})
                self.fields[field_name] = forms.MultipleChoiceField(
                                                        required=False,
                                                        label=label,
                                                        widget=widget)
                self.fields[field_name].choices = users_list
                self.fields[field_name].initial = []

        # Figure out users & roles
        # NOTE(garcianavalon) logic for this part:
        # find all the (fiware)roles from the role_list(the ones the current user
        # can get and assign) that are already assigned for each user and flag them
        # ALSO find out the organization members and assign them the
        # default (keystone)role
        try:
            project_users_roles = self.list_users_with_roles(request,
                                                            project_id,
                                                            role_list['applications'])
        except Exception:
            exceptions.handle(request,
                              err_msg,
                              redirect=reverse(INDEX_URL))

        # Flag the alredy owned ones, both organization and application
        for user_id in project_users_roles:
            roles_ids = project_users_roles[user_id]
            for role_id in roles_ids:
                field_name = self.get_member_field_name(role_id)
                self.fields[field_name].initial.append(user_id)

    class Meta:
        name = _("Organization Members")
        slug = PROJECT_USER_MEMBER_SLUG

class UpdateProjectMembers(workflows.UpdateMembersStep, RolesMixin):
    action_class = UpdateProjectMembersAction
    available_list_title = _("All Users")
    members_list_title = _("Organization Members")
    no_available_text = _("No users found.")
    no_members_text = _("No users.")
    contributes = ("project_id",)

    def contribute(self, data, context):
        project_id = context['project_id']
        if data:
            try:
                role_list = self.list_all_roles(self.workflow.request, project_id)
            except Exception:
                exceptions.handle(self.workflow.request,
                                  _('Unable to retrieve role list.'))

            post = self.workflow.request.POST
            for k in role_list:
                for role in role_list[k]:
                    field = self.get_member_field_name(role.id)
                    context[field] = post.getlist(field)
        return context

class ManageOrganizationMembers(workflows.Workflow, RolesMixin):
    slug = "manage_organization_users"
    name = _("Manage Members")
    finalize_button_name = _("Save")
    success_message = _('Modified users.')
    failure_message = _('Unable to modify users.')
    success_url = "horizon:idm:organizations:detail"
    default_steps = (UpdateProjectMembers,)

    def get_success_url(self):
        # Overwrite to allow passing kwargs
        return reverse(self.success_url,
                    kwargs={'organization_id':self.context['project_id']})

    def handle(self, request, data):
        project_id = data['project_id']
        # Project-user member step
        member_step = self.get_step(PROJECT_USER_MEMBER_SLUG)

        try:
            role_list = self.list_all_roles(request, project_id)
            project_users_roles = self.list_users_with_roles(request,
                                                            project_id,
                                                            role_list['applications'])
            # re-index by role with a user list for easier processing in later steps
            current_roles = idm_utils.swap_dict(project_users_roles)

            # Parse the form data
            modified_roles = {}
            for k in role_list:
                for role in role_list[k]:
                    field_name = member_step.get_member_field_name(role.id)
                    modified_roles[role.id] = data[field_name]

            # Create the delete and add sets
            roles_to_add, roles_to_delete = self._create_add_and_delete_sets(
                                                                modified_roles,
                                                                current_roles)
            application_roles = [r.id for r in role_list['applications']]
            organization_roles = [r.id for r in role_list['organization']]

            # Add the roles
            add_methods = [
                (application_roles, fiware_api.keystone.add_role_to_user),
                (organization_roles, api.keystone.add_tenant_user_role)
            ]
            self._apply_method(roles_to_add, add_methods, project_id)

            # Remove the roles
            delete_methods = [
                (application_roles, fiware_api.keystone.remove_role_from_user),
                (organization_roles, api.keystone.remove_tenant_user_role)
            ]
            self._apply_method(roles_to_delete, delete_methods, project_id)

            return True
        except Exception:
            exceptions.handle(request,
                          _('Failed to modify organization\'s members.'))
            return False

    def _create_add_and_delete_sets(self, modified_roles, current_roles):
        roles_to_add = {}
        roles_to_delete = {}
        for role_id in modified_roles:
            new_users = set(modified_roles.get(role_id, []))
            current_users = set(current_roles.get(role_id, []))
            # users to add-> users in N and not in C -> N-C
            users_to_add = new_users - current_users
            if users_to_add:
                roles_to_add[role_id] = users_to_add
            # users to delete -> users in C and not in N -> C-N
            users_to_delete = current_users - new_users
            if users_to_delete:
                roles_to_delete[role_id] = users_to_delete
        return roles_to_add, roles_to_delete

    def _apply_method(self, roles_users, methods, project_id):
        """Utility method to iterate all elements in the dictionary and
        apply a method based on a condition.

            :param methods: a list with the methods to apply to each group
            :type methods: list of tuples (container, function)
        """
        for role_id in roles_users:
            function = self._decide_function(role_id, methods)
            if not function:
                # We should never get to this point
                LOG.error('Unexpected role {0} not in possible roles'.format(role_id))
                raise ValueError('Invalid role')
            for user_id in roles_users[role_id]:
                function(self.request,
                        project=project_id,
                        user=user_id,
                        role=role_id)

    def _decide_function(self, role_id, methods):
        for (container, function) in methods:
            if role_id not in container:
                continue
            return function
        return None
<commit_msg>showing application id<commit_after># Copyright (C) 2014 Universidad Politecnica de Madrid
#
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

import logging

from django.conf import settings
from django.core.urlresolvers import reverse
from django.utils.translation import ugettext_lazy as _

from horizon import exceptions
from horizon import forms
from horizon import messages
from horizon import workflows

from openstack_dashboard import api
from openstack_dashboard import fiware_api
from openstack_dashboard.dashboards.idm import utils as idm_utils

INDEX_URL = "horizon:idm:organizations:index"
PROJECT_USER_MEMBER_SLUG = "update_members"
LOG = logging.getLogger('idm_logger')

class RolesMixin(object):
    """Simple container for the roles logic."""
    # TODO(garcianavalon) maybe move to fiware_api
    def list_all_roles(self, request, project_id):
        role_list = {}
        # NOTE(garcianavalon) list all roles grouped by application
        # on which current user has the right to get and assign
        role_list['applications'] = fiware_api.keystone.list_allowed_roles_to_assign(request,
                                                user=request.user.id,
                                                organization=project_id)
        # NOTE(garcianavalon) we also need the organization (keystone)
        # roles here to add members.
        role_list['organization'] = api.keystone.role_list(request)
        # Little trick to show them together, we simulate the organization to be
        # an application
        for role in role_list['organization']:
            role.application = project_id
        return role_list

    def list_users_with_roles(self, request, project_id, available_roles):
        # First, load all users from the organization
        # with their roles
        project_users_roles = api.keystone.get_project_users_roles(request,
                                                              project_id)
        # Second, load all the organization-scoped application roles for every user
        # but only the ones the user can assign (and only scoped for the current
        # organization)
        for user_id in project_users_roles:
            project_users_roles[user_id] += [
                r.id for r in fiware_api.keystone.role_list(
                                                    request,
                                                    user=user_id,
                                                    organization=project_id)
                    if r in available_roles
            ]
        return project_users_roles

class UpdateProjectMembersAction(workflows.MembershipAction, RolesMixin):
    def __init__(self, request, *args, **kwargs):
        super(UpdateProjectMembersAction, self).__init__(request,
                                                         *args,
                                                         **kwargs)
        err_msg = _('Unable to retrieve user list. Please try again later.')
        project_id = self.initial['project_id']

        # Get the default role
        try:
            default_role = api.keystone.get_default_role(self.request)
            # Default role is necessary to add members to a project
            if default_role is None:
                default = getattr(settings,
                                  "OPENSTACK_KEYSTONE_DEFAULT_ROLE", None)
                msg = (_('Could not find default role "%s" in Keystone') %
                       default)
                raise exceptions.NotFound(msg)
        except Exception:
            exceptions.handle(self.request,
                              err_msg,
                              redirect=reverse(INDEX_URL))
        # TODO(garcianavalon) hide the default role, show the 'owner'
        # role under the organization ('owner' is 'admin' for now)
        default_role_name = self.get_default_role_field_name()
        self.fields[default_role_name] = forms.CharField(required=False)
        self.fields[default_role_name].initial = default_role.id

        # Get list of available users
        try:
            all_users = api.keystone.user_list(request)
        except Exception:
            exceptions.handle(request, err_msg)
        users_list = [(user.id, user.name) for user in all_users]

        # Get list of roles
        try:
            role_list = self.list_all_roles(request, project_id)
        except Exception:
            exceptions.handle(request,
                              err_msg,
                              redirect=reverse(INDEX_URL))
        for k in role_list:
            # TODO(garcianavalon) different behaviour
            # TODO(garcianavalon) hide the default role, show the 'owner'
            # role under the organization ('owner' is 'admin' for now)
            for role in role_list[k]:
                field_name = self.get_member_field_name(role.id)
                label = role.name
                widget = forms.widgets.SelectMultiple(
                                attrs={'data-application-name':role.application})
                self.fields[field_name] = forms.MultipleChoiceField(
                                                        required=False,
                                                        label=label,
                                                        widget=widget)
                self.fields[field_name].choices = users_list
                self.fields[field_name].initial = []

        # Figure out users & roles
        # NOTE(garcianavalon) logic for this part:
        # find all the (fiware)roles from the role_list(the ones the current user
        # can get and assign) that are already assigned for each user and flag them
        # ALSO find out the organization members and assign them the
        # default (keystone)role
        try:
            project_users_roles = self.list_users_with_roles(request,
                                                            project_id,
                                                            role_list['applications'])
        except Exception:
            exceptions.handle(request,
                              err_msg,
                              redirect=reverse(INDEX_URL))

        # Flag the alredy owned ones, both organization and application
        for user_id in project_users_roles:
            roles_ids = project_users_roles[user_id]
            for role_id in roles_ids:
                field_name = self.get_member_field_name(role_id)
                self.fields[field_name].initial.append(user_id)

    class Meta:
        name = _("Organization Members")
        slug = PROJECT_USER_MEMBER_SLUG

class UpdateProjectMembers(workflows.UpdateMembersStep, RolesMixin):
    action_class = UpdateProjectMembersAction
    available_list_title = _("All Users")
    members_list_title = _("Organization Members")
    no_available_text = _("No users found.")
    no_members_text = _("No users.")
    contributes = ("project_id",)

    def contribute(self, data, context):
        project_id = context['project_id']
        if data:
            try:
                role_list = self.list_all_roles(self.workflow.request, project_id)
            except Exception:
                exceptions.handle(self.workflow.request,
                                  _('Unable to retrieve role list.'))

            post = self.workflow.request.POST
            for k in role_list:
                for role in role_list[k]:
                    field = self.get_member_field_name(role.id)
                    context[field] = post.getlist(field)
        return context

class ManageOrganizationMembers(workflows.Workflow, RolesMixin):
    slug = "manage_organization_users"
    name = _("Manage Members")
    finalize_button_name = _("Save")
    success_message = _('Modified users.')
    failure_message = _('Unable to modify users.')
    success_url = "horizon:idm:organizations:detail"
    default_steps = (UpdateProjectMembers,)

    def get_success_url(self):
        # Overwrite to allow passing kwargs
        return reverse(self.success_url,
                    kwargs={'organization_id':self.context['project_id']})

    def handle(self, request, data):
        project_id = data['project_id']
        # Project-user member step
        member_step = self.get_step(PROJECT_USER_MEMBER_SLUG)

        try:
            role_list = self.list_all_roles(request, project_id)
            project_users_roles = self.list_users_with_roles(request,
                                                            project_id,
                                                            role_list['applications'])
            # re-index by role with a user list for easier processing in later steps
            current_roles = idm_utils.swap_dict(project_users_roles)

            # Parse the form data
            modified_roles = {}
            for k in role_list:
                for role in role_list[k]:
                    field_name = member_step.get_member_field_name(role.id)
                    modified_roles[role.id] = data[field_name]

            # Create the delete and add sets
            roles_to_add, roles_to_delete = self._create_add_and_delete_sets(
                                                                modified_roles,
                                                                current_roles)
            application_roles = [r.id for r in role_list['applications']]
            organization_roles = [r.id for r in role_list['organization']]

            # Add the roles
            add_methods = [
                (application_roles, fiware_api.keystone.add_role_to_user),
                (organization_roles, api.keystone.add_tenant_user_role)
            ]
            self._apply_method(roles_to_add, add_methods, project_id)

            # Remove the roles
            delete_methods = [
                (application_roles, fiware_api.keystone.remove_role_from_user),
                (organization_roles, api.keystone.remove_tenant_user_role)
            ]
            self._apply_method(roles_to_delete, delete_methods, project_id)

            return True
        except Exception:
            exceptions.handle(request,
                          _('Failed to modify organization\'s members.'))
            return False

    def _create_add_and_delete_sets(self, modified_roles, current_roles):
        roles_to_add = {}
        roles_to_delete = {}
        for role_id in modified_roles:
            new_users = set(modified_roles.get(role_id, []))
            current_users = set(current_roles.get(role_id, []))
            # users to add-> users in N and not in C -> N-C
            users_to_add = new_users - current_users
            if users_to_add:
                roles_to_add[role_id] = users_to_add
            # users to delete -> users in C and not in N -> C-N
            users_to_delete = current_users - new_users
            if users_to_delete:
                roles_to_delete[role_id] = users_to_delete
        return roles_to_add, roles_to_delete

    def _apply_method(self, roles_users, methods, project_id):
        """Utility method to iterate all elements in the dictionary and
        apply a method based on a condition.

            :param methods: a list with the methods to apply to each group
            :type methods: list of tuples (container, function)
        """
        for role_id in roles_users:
            function = self._decide_function(role_id, methods)
            if not function:
                # We should never get to this point
                LOG.error('Unexpected role {0} not in possible roles'.format(role_id))
                raise ValueError('Invalid role')
            for user_id in roles_users[role_id]:
                function(self.request,
                        project=project_id,
                        user=user_id,
                        role=role_id)

    def _decide_function(self, role_id, methods):
        for (container, function) in methods:
            if role_id not in container:
                continue
            return function
        return None
<|endoftext|>