Exemplo n.º 1
0
def copper_node(workdir):
    logpath = os.path.join(workdir, 'copper.log')
    unixpath = os.path.join(workdir, 'copper.sock')
    httppath = os.path.join(workdir, 'copper.http')
    confpath = os.path.join(workdir, 'copper.conf')
    config = {
        "listen": [
            {
                "net": "unix",
                "type": "http",
                "addr": httppath,
            },
            {
                "net": "unix",
                "addr": unixpath,
                "allow-changes": True,
            },
        ],
    }
    with open(confpath, 'w') as f:
        # YAML parses valid JSON data
        json.dump(config, f)
    with open(logpath, 'wb') as logfile:
        p = Popen(['copper-node', '-config=' + confpath], shell=False, cwd=workdir, stdout=logfile, stderr=logfile)
    try:
        while not os.path.exists(unixpath):
            time.sleep(0.001)
            rc = p.poll()
            if rc is not None:
                with open(logpath, 'rb') as logfile:
                    sys.stderr.write(logfile.read())
                raise RuntimeError('copper-node exited with status %r' % (rc,))
        yield {
            'unix': unixpath,
            'http': httppath,
        }
    finally:
        if p.poll() is None:
            p.terminate()
            p.wait()
Exemplo n.º 2
0
class PlatformWrapper:
    def __init__(self):
        """ Initializes a new VOLTTRON instance

        Creates a temporary VOLTTRON_HOME directory with a packaged directory
        for agents that are built.
        """

        # This is hopefully going to keep us from attempting to shutdown
        # multiple times.  For example if a fixture calls shutdown and a
        # lower level fixture calls shutdown, this won't hang.
        self._instance_shutdown = False

        self.volttron_home = tempfile.mkdtemp()
        self.packaged_dir = os.path.join(self.volttron_home, "packaged")
        os.makedirs(self.packaged_dir)

        # in the context of this platform it is very important not to
        # use the main os.environ for anything.
        self.env = {
            'VOLTTRON_HOME': self.volttron_home,
            'PACKAGED_DIR': self.packaged_dir,
            'DEBUG_MODE': os.environ.get('DEBUG_MODE', ''),
            'DEBUG': os.environ.get('DEBUG', ''),
            'PATH': VOLTTRON_ROOT + ':' + os.environ['PATH']
        }
        self.volttron_root = VOLTTRON_ROOT

        volttron_exe = subprocess.check_output(['which', 'volttron']).strip()

        assert os.path.exists(volttron_exe)
        self.python = os.path.join(os.path.dirname(volttron_exe), 'python')
        assert os.path.exists(self.python)

        # By default no web server should be started.
        self.bind_web_address = None
        self.discovery_address = None
        self.jsonrpc_endpoint = None
        self.volttron_central_address = None
        self.instance_name = None
        self.serverkey = None

        self.p_process = None
        self.t_process = None

        self.started_agent_pids = []
        self.local_vip_address = None
        self.vip_address = None
        self.logit('Creating platform wrapper')

        # This was used when we are testing the SMAP historian.
        self.use_twistd = False

        # Added restricted code properties
        self.certsobj = None

        # Control whether the instance directory is cleaned up when shutdown.
        # if the environment variable DEBUG is set to a True value then the
        # instance is not cleaned up.
        self.skip_cleanup = False

        # This is used as command line entry replacement.  Especially working
        # with older 2.0 agents.
        self.opts = None

        keystorefile = os.path.join(self.volttron_home, 'keystore')
        self.keystore = KeyStore(keystorefile)
        self.keystore.generate()

    def logit(self, message):
        print('{}: {}'.format(self.volttron_home, message))

    def allow_all_connections(self):
        """ Add a /.*/ entry to the auth.json file.
        """
        entry = AuthEntry(credentials="/.*/")
        authfile = AuthFile(self.volttron_home + "/auth.json")
        try:
            authfile.add(entry)
        except AuthFileEntryAlreadyExists:
            pass

    def build_connection(self,
                         peer=None,
                         address=None,
                         identity=None,
                         publickey=None,
                         secretkey=None,
                         serverkey=None,
                         capabilities=[],
                         **kwargs):
        self.logit('Building connection to {}'.format(peer))
        self.allow_all_connections()

        if address is None:
            self.logit(
                'Default address was None so setting to current instances')
            address = self.vip_address
            serverkey = self.serverkey
        if serverkey is None:
            self.logit("serverkey wasn't set but the address was.")
            raise Exception("Invalid state.")

        if publickey is None or secretkey is None:
            self.logit('generating new public secret key pair')
            keyfile = tempfile.mktemp(".keys", "agent", self.volttron_home)
            keys = KeyStore(keyfile)
            keys.generate()
            publickey = keys.public
            secretkey = keys.secret
            entry = AuthEntry(capabilities=capabilities,
                              comments="Added by test",
                              credentials=keys.public)
            file = AuthFile(self.volttron_home + "/auth.json")
            file.add(entry)

        conn = Connection(address=address,
                          peer=peer,
                          publickey=publickey,
                          secretkey=secretkey,
                          serverkey=serverkey,
                          volttron_home=self.volttron_home)

        return conn

    def build_agent(self,
                    address=None,
                    should_spawn=True,
                    identity=None,
                    publickey=None,
                    secretkey=None,
                    serverkey=None,
                    agent_class=Agent,
                    **kwargs):
        """ Build an agent connnected to the passed bus.

        By default the current instance that this class wraps will be the
        vip address of the agent.

        :param address:
        :param should_spawn:
        :param identity:
        :param publickey:
        :param secretkey:
        :param serverkey:
        :param agent_class: Agent class to build
        :return:
        """
        self.logit("Building generic agent.")

        use_ipc = kwargs.pop('use_ipc', False)

        if serverkey is None:
            serverkey = self.serverkey
        if publickey is None:
            self.logit('generating new public secret key pair')
            keyfile = tempfile.mktemp(".keys", "agent", self.volttron_home)
            keys = KeyStore(keyfile)
            keys.generate()
            publickey = keys.public
            secretkey = keys.secret

        if address is None:
            self.logit('Using vip-address ' + self.vip_address)
            address = self.vip_address

        if publickey and not serverkey:
            self.logit('using instance serverkey: {}'.format(self.publickey))
            serverkey = self.publickey

        agent = agent_class(address=address,
                            identity=identity,
                            publickey=publickey,
                            secretkey=secretkey,
                            serverkey=serverkey,
                            volttron_home=self.volttron_home,
                            **kwargs)
        self.logit('platformwrapper.build_agent.address: {}'.format(address))

        # Automatically add agent's credentials to auth.json file
        if publickey:
            self.logit('Adding publickey to auth.json')
            gevent.spawn(self._append_allow_curve_key, publickey)
            gevent.sleep(0.1)

        if should_spawn:
            self.logit('platformwrapper.build_agent spawning')
            event = gevent.event.Event()
            gevent.spawn(agent.core.run, event)  # .join(0)
            event.wait(timeout=2)

            hello = agent.vip.hello().get(timeout=.3)
            self.logit('Got hello response {}'.format(hello))
        agent.publickey = publickey
        return agent

    def _read_auth_file(self):
        auth_path = os.path.join(self.volttron_home, 'auth.json')
        try:
            with open(auth_path, 'r') as fd:
                data = strip_comments(FileObject(fd, close=False).read())
                if data:
                    auth = jsonapi.loads(data)
                else:
                    auth = {}
        except IOError:
            auth = {}
        if 'allow' not in auth:
            auth['allow'] = []
        return auth, auth_path

    def _append_allow_curve_key(self, publickey):
        entry = AuthEntry(credentials=publickey)
        authfile = AuthFile(self.volttron_home + "/auth.json")
        try:
            authfile.add(entry)
        except AuthFileEntryAlreadyExists:
            pass

    def add_vc(self):
        return add_vc_to_instance(self)

    def add_capabilities(self, publickey, capabilities):
        if isinstance(capabilities, basestring):
            capabilities = [capabilities]
        auth, auth_path = self._read_auth_file()
        cred = publickey
        allow = auth['allow']
        entry = next((item for item in allow if item['credentials'] == cred),
                     {})
        caps = entry.get('capabilities', [])
        entry['capabilities'] = list(set(caps + capabilities))

        with open(auth_path, 'w+') as fd:
            json.dump(auth, fd)

    def set_auth_dict(self, auth_dict):
        if auth_dict:
            with open(os.path.join(self.volttron_home, 'auth.json'),
                      'w') as fd:
                fd.write(json.dumps(auth_dict))

    def startup_platform(self,
                         vip_address,
                         auth_dict=None,
                         use_twistd=False,
                         mode=UNRESTRICTED,
                         bind_web_address=None,
                         volttron_central_address=None,
                         volttron_central_serverkey=None):

        # if not isinstance(vip_address, list):
        #     self.vip_address = [vip_address]
        # else:
        #     self.vip_address = vip_address

        self.vip_address = vip_address
        self.mode = mode
        self.bind_web_address = bind_web_address
        if self.bind_web_address:
            self.discovery_address = "{}/discovery/".format(
                self.bind_web_address)

            # Only available if vc is installed!
            self.jsonrpc_endpoint = "{}/jsonrpc".format(self.bind_web_address)

        enable_logging = self.env.get('ENABLE_LOGGING', False)
        debug_mode = self.env.get('DEBUG_MODE', False)
        if not debug_mode:
            debug_mode = self.env.get('DEBUG', False)
        self.skip_cleanup = self.env.get('SKIP_CLEANUP', False)
        if debug_mode:
            self.skip_cleanup = True
            enable_logging = True
        self.logit("In start up platform enable_logging is {} ".format(
            enable_logging))
        assert self.mode in MODES, 'Invalid platform mode set: ' + str(mode)
        opts = None

        # see main.py for how we handle pub sub addresses.
        ipc = 'ipc://{}{}/run/'.format(
            '@' if sys.platform.startswith('linux') else '',
            self.volttron_home)
        self.local_vip_address = ipc + 'vip.socket'
        self.set_auth_dict(auth_dict)

        self.opts = {
            'verify_agents': False,
            'volttron_home': self.volttron_home,
            'vip_address': vip_address,
            'vip_local_address': ipc + 'vip.socket',
            'publish_address': ipc + 'publish',
            'subscribe_address': ipc + 'subscribe',
            'bind_web_address': bind_web_address,
            'volttron_central_address': volttron_central_address,
            'volttron_central_serverkey': volttron_central_serverkey,
            'platform_name': None,
            'log': os.path.join(self.volttron_home, 'volttron.log'),
            'log_config': None,
            'monitor': True,
            'autostart': True,
            'log_level': logging.DEBUG,
            'verboseness': logging.DEBUG
        }

        pconfig = os.path.join(self.volttron_home, 'config')
        config = {}

        # Add platform's public key to known hosts file
        publickey = self.keystore.public
        known_hosts_file = os.path.join(self.volttron_home, 'known_hosts')
        known_hosts = KnownHostsStore(known_hosts_file)
        known_hosts.add(self.opts['vip_local_address'], publickey)
        known_hosts.add(self.opts['vip_address'], publickey)

        # Set up the configuration file based upon the passed parameters.
        parser = configparser.ConfigParser()
        parser.add_section('volttron')
        parser.set('volttron', 'vip-address', vip_address)
        if bind_web_address:
            parser.set('volttron', 'bind-web-address', bind_web_address)
        if volttron_central_address:
            parser.set('volttron', 'volttron-central-address',
                       volttron_central_address)
        if volttron_central_serverkey:
            parser.set('volttron', 'volttron-central-serverkey',
                       volttron_central_serverkey)
        if self.mode == UNRESTRICTED:
            # TODO Restricted code should set with volttron as contianer
            # if RESTRICTED_AVAILABLE:
            #     config['mobility'] = False
            #     config['resource-monitor'] = False
            #     config['verify'] = False
            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
                parser.write(cfg)

        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(self.volttron_home, 'certificates')

            print("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)

            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
        else:
            raise PlatformWrapperError(
                "Invalid platform mode specified: {}".format(mode))

        log = os.path.join(self.volttron_home, 'volttron.log')
        if enable_logging:
            cmd = ['volttron', '-vv', '-l{}'.format(log)]
        else:
            cmd = ['volttron', '-l{}'.format(log)]

        print('process environment: {}'.format(self.env))
        print('popen params: {}'.format(cmd))
        self.p_process = Popen(cmd,
                               env=self.env,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)

        assert self.p_process is not None
        # A None value means that the process is still running.
        # A negative means that the process exited with an error.
        assert self.p_process.poll() is None

        self.serverkey = self.keystore.public
        assert self.serverkey
        agent = self.build_agent()

        has_control = False
        times = 0
        while not has_control and times < 10:
            times += 1
            try:
                has_control = agent.vip.peerlist().get(timeout=.2)
            except gevent.Timeout:
                pass

        if not has_control:
            self.shutdown_platform()
            raise "Couldn't connect to core platform!"

        if bind_web_address:
            times = 0
            has_discovery = False
            while times < 10:
                times += 1
                try:
                    resp = requests.get(self.discovery_address)
                    if resp.ok:
                        has_discovery = True
                        break
                except Exception as e:
                    gevent.sleep(0.1)
                    self.logit("Connection error found {}".format(e))
            if not has_discovery:
                raise "Couldn't connect to discovery platform."

        self.use_twistd = use_twistd

        # TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.volttron_home,
                                   TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = [TWISTED_START, "-n", "smap", tconfig]
            self.t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)

    def is_running(self):
        self.logit("PROCESS IS RUNNING: {}".format(self.p_process))
        return self.p_process is not None and self.p_process.poll() is None

    def twistd_is_running(self):
        return self.t_process is not None

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print("wrapper.certsobj", self.certsobj.cert_dir)
        assert (auth.sign_as_creator(
            package, 'creator',
            certsobj=self.certsobj)), "Signing as {} failed.".format('creator')

    def direct_sign_agentpackage_admin(self, package):
        assert (RESTRICTED), "Auth not available"
        assert (auth.sign_as_admin(
            package, 'admin',
            certsobj=self.certsobj)), "Signing as {} failed.".format('admin')

    def direct_sign_agentpackage_initiator(self, package, config_file,
                                           contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file": config_file, "contract": contract}
        assert (auth.sign_as_initiator(package,
                                       'initiator',
                                       files=files,
                                       certsobj=self.certsobj)
                ), "Signing as {} failed.".format('initiator')

    def _aip(self):
        opts = type('Options', (), self.opts)
        aip = AIPplatform(opts)
        aip.setup()
        return aip

    def _install_agent(self, wheel_file, start, vip_identity):
        self.logit('Creating channel for sending the agent.')
        gevent.sleep(0.3)
        self.logit('calling control install agent.')
        self.logit("VOLTTRON_HOME SETTING: {}".format(
            self.env['VOLTTRON_HOME']))
        env = self.env.copy()
        cmd = ['volttron-ctl', '-vv', 'install', wheel_file]
        if vip_identity:
            cmd.extend(['--vip-identity', vip_identity])
        self.logit("cmd: {}".format(cmd))
        res = subprocess.check_output(cmd, env=env)
        assert res, "failed to install wheel:{}".format(wheel_file)
        agent_uuid = res.split(' ')[-2]
        self.logit(agent_uuid)

        if start:
            self.start_agent(agent_uuid)

        return agent_uuid

    def install_multiple_agents(self, agent_configs):
        """
        Installs mutltiple agents on the platform.

        :param agent_configs:list
            A list of 3-tuple that allows the configuration of a platform
            in a single go.  The tuple order is
            1. path to the agent directory.
            2. configuration data (either file or json data)
            3. Whether the agent should be started or not.

        :return:list:
            A list of uuid's associated with the agents that were installed.


        :Note:
            In order for this method to be called the platform must be
            currently running.
        """
        if not self.is_running():
            raise PlatformWrapperError("Instance isn't running!")
        results = []

        for path, config, start in agent_configs:
            results = self.install_agent(agent_dir=path,
                                         config_file=config,
                                         start=start)

        return results

    def install_agent(self,
                      agent_wheel=None,
                      agent_dir=None,
                      config_file=None,
                      start=True,
                      vip_identity=None):
        """
        Install and optionally start an agent on the instance.

        This function allows installation from an agent wheel or an
        agent directory (NOT BOTH).  If an agent_wheel is specified then
        it is assumed to be ready for installation (has a config file).
        If an agent_dir is specified then a config_file file must be
        specified or if it is not specified then it is assumed that the
        file agent_dir/config is to be used as the configuration file.  If
        none of these exist then an assertion error will be thrown.

        This function will return with a uuid of the installed agent.

        :param agent_wheel:
        :param agent_dir:
        :param config_file:
        :param start:
        :param vip_identity:
        :return:
        """

        assert self.is_running(), "Instance must be running to install agent."
        assert agent_wheel or agent_dir, "Invalid agent_wheel or agent_dir."

        if agent_wheel:
            assert not agent_dir
            assert not config_file
            assert os.path.exists(agent_wheel)
            wheel_file = agent_wheel
            agent_uuid = self._install_agent(wheel_file, start, vip_identity)

        # Now if the agent_dir is specified.
        if agent_dir:
            assert not agent_wheel
            if isinstance(config_file, dict):
                from os.path import join, basename
                temp_config = join(self.volttron_home,
                                   basename(agent_dir) + "_config_file")
                with open(temp_config, "w") as fp:
                    fp.write(json.dumps(config_file))
                config_file = temp_config
            elif not config_file:
                if os.path.exists(os.path.join(agent_dir, "config")):
                    config_file = os.path.join(agent_dir, "config")
                else:
                    from os.path import join, basename
                    temp_config = join(self.volttron_home,
                                       basename(agent_dir) + "_config_file")
                    with open(temp_config, "w") as fp:
                        fp.write(json.dumps({}))
                    config_file = temp_config
            elif os.path.exists(config_file):
                pass  # config_file already set!
            else:
                raise ValueError("Can't determine correct config file.")

            script = os.path.join(self.volttron_root,
                                  "scripts/install-agent.py")
            cmd = [
                self.python, script, "--volttron-home", self.volttron_home,
                "--volttron-root", self.volttron_root, "--agent-source",
                agent_dir, "--config", config_file, "--json"
            ]

            if vip_identity:
                cmd.extend(["--vip-identity", vip_identity])
            if start:
                cmd.extend(["--start"])

            results = subprocess.check_output(cmd)

            # Because we are no longer silencing output from the install, the
            # the results object is now much more verbose.  Our assumption is
            # the line before the output we care about has WHEEL at the end
            # of it.
            new_results = ""
            found_wheel = False
            for line in results.split("\n"):
                if line.endswith("WHEEL"):
                    found_wheel = True
                elif found_wheel:
                    new_results += line
            results = new_results

            #
            # Response from results is expected as follows depending on
            # parameters, note this is a json string so parse to get dictionary.
            # {
            #     "started": true,
            #     "agent_pid": 26241,
            #     "starting": true,
            #     "agent_uuid": "ec1fd94e-922a-491f-9878-c392b24dbe50"
            # }
            assert results

            resultobj = jsonapi.loads(str(results))

            if start:
                assert resultobj['started']
            agent_uuid = resultobj['agent_uuid']

        assert agent_uuid is not None

        if start:
            assert self.is_agent_running(agent_uuid)

        return agent_uuid

    def start_agent(self, agent_uuid):
        self.logit('Starting agent {}'.format(agent_uuid))
        self.logit("VOLTTRON_HOME SETTING: {}".format(
            self.env['VOLTTRON_HOME']))
        cmd = ['volttron-ctl']
        cmd.extend(['start', agent_uuid])
        p = Popen(cmd, env=self.env, stdout=sys.stdout, stderr=sys.stderr)
        p.wait()

        # Confirm agent running
        cmd = ['volttron-ctl']
        cmd.extend(['status', agent_uuid])
        res = subprocess.check_output(cmd, env=self.env)
        # 776 TODO: Timing issue where check fails
        time.sleep(.1)
        self.logit("Subprocess res is {}".format(res))
        assert 'running' in res
        pidpos = res.index('[') + 1
        pidend = res.index(']')
        pid = int(res[pidpos:pidend])

        assert psutil.pid_exists(pid), \
            "The pid associated with agent {} does not exist".format(pid)

        self.started_agent_pids.append(pid)
        return pid

    def stop_agent(self, agent_uuid):
        # Confirm agent running
        _log.debug("STOPPING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl']
            cmd.extend(['stop', agent_uuid])
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_pid(agent_uuid)

    def list_agents(self):
        agent = self.build_agent()
        print('PEER LIST: {}'.format(agent.vip.peerlist().get(timeout=10)))
        agent_list = agent.vip.rpc('control', 'list_agents').get(timeout=10)
        agent.core.stop(timeout=3)
        return agent_list

    def remove_agent(self, agent_uuid):
        """Remove the agent specified by agent_uuid"""
        _log.debug("REMOVING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl']
            cmd.extend(['remove', agent_uuid])
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_pid(agent_uuid)

    def is_agent_running(self, agent_uuid):
        return self.agent_pid(agent_uuid) is not None

    def agent_pid(self, agent_uuid):
        """
        Returns the pid of a running agent or None

        :param agent_uuid:
        :return:
        """
        # Confirm agent running
        cmd = ['volttron-ctl']
        cmd.extend(['status', agent_uuid])
        pid = None
        try:
            res = subprocess.check_output(cmd, env=self.env)
            try:
                pidpos = res.index('[') + 1
                pidend = res.index(']')
                pid = int(res[pidpos:pidend])
            except:
                pid = None
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))

        # Handle the following exception that seems to happen when getting a
        # pid of an agent during the platform shutdown phase.
        #
        # Logged from file platformwrapper.py, line 797
        #   AGENT             IDENTITY          TAG STATUS
        # Traceback (most recent call last):
        #   File "/usr/lib/python2.7/logging/__init__.py", line 882, in emit
        #     stream.write(fs % msg)
        #   File "/home/volttron/git/volttron/env/local/lib/python2.7/site-packages/_pytest/capture.py", line 244, in write
        #     self.buffer.write(obj)
        # ValueError: I/O operation on closed file
        except ValueError:
            pass
        # _log.debug("AGENT_PID: {}".format(pid))
        return pid

    def build_agentpackage(self, agent_dir, config_file={}):
        if isinstance(config_file, dict):
            cfg_path = os.path.join(agent_dir, "config_temp")
            with open(cfg_path, "w") as tmp_cfg:
                tmp_cfg.write(jsonapi.dumps(config_file))
            config_file = cfg_path

        # Handle relative paths from the volttron git directory.
        if not os.path.isabs(agent_dir):
            agent_dir = os.path.join(self.volttron_root, agent_dir)

        assert os.path.exists(config_file)
        assert os.path.exists(agent_dir)

        wheel_path = packaging.create_package(agent_dir, self.packaged_dir)
        packaging.add_files_to_package(
            wheel_path, {'config_file': os.path.join('./', config_file)})

        return wheel_path

    def confirm_agent_running(self,
                              agent_name,
                              max_retries=5,
                              timeout_seconds=2):
        running = False
        retries = 0
        while not running and retries < max_retries:
            status = self.test_aip.status_agents()
            print("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(
                    status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running

    # def direct_stop_agent(self, agent_uuid):
    #     result = self.conn.call.stop_agent(agent_uuid)
    #     print result

    def shutdown_platform(self):
        """
        Stop platform here.  First grab a list of all of the agents that are
        running on the platform, then shutdown, then if any of the listed agent
        pids are still running then kill them.
        """

        # Handle cascading calls from multiple levels of fixtures.
        if self._instance_shutdown:
            return

        running_pids = []

        for agnt in self.list_agents():
            pid = self.agent_pid(agnt['uuid'])
            if pid is not None and int(pid) > 0:
                running_pids.append(int(pid))

        # First try and nicely shutdown the platform, which should clean all
        # of the agents up automatically.
        cmd = ['volttron-ctl']
        cmd.extend(['shutdown', '--platform'])
        try:
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError:
            if self.p_process is not None:
                try:
                    gevent.sleep(0.2)
                    self.p_process.terminate()
                    gevent.sleep(0.2)
                except OSError:
                    self.logit('Platform process was terminated.')
            else:
                self.logit("platform process was null")

        for pid in running_pids:
            if psutil.pid_exists(pid):
                self.logit("TERMINATING: {}".format(pid))
                proc = psutil.Process(pid)
                proc.terminate()

        if self.use_twistd and self.t_process is not None:
            self.t_process.kill()
            self.t_process.wait()
        elif self.use_twistd:
            self.logit("twistd process was null")

        if os.environ.get('PRINT_LOG'):
            logpath = os.path.join(self.volttron_home, 'volttron.log')
            if os.path.exists(logpath):
                print("************************* Begin {}".format(logpath))
                with open(logpath) as f:
                    for l in f.readlines():
                        print(l)
                print("************************* End {}".format(logpath))
            else:
                print("######################### No Log Exists: {}".format(
                    logpath))
        if not self.skip_cleanup:
            self.logit('Removing {}'.format(self.volttron_home))
            shutil.rmtree(self.volttron_home, ignore_errors=True)
        self._instance_shutdown = True

    def __repr__(self):
        return str(self)

    def __str__(self):
        data = []
        data.append('volttron_home: {}'.format(self.volttron_home))
        return '\n'.join(data)
Exemplo n.º 3
0
class PlatformWrapper:
    def __init__(self):
        """ Initializes a new VOLTTRON instance

        Creates a temporary VOLTTRON_HOME directory with a packaged directory
        for agents that are built.
        """

        # This is hopefully going to keep us from attempting to shutdown
        # multiple times.  For example if a fixture calls shutdown and a
        # lower level fixture calls shutdown, this won't hang.
        self._instance_shutdown = False

        self.volttron_home = tempfile.mkdtemp()
        self.packaged_dir = os.path.join(self.volttron_home, "packaged")
        os.makedirs(self.packaged_dir)

        # in the context of this platform it is very important not to
        # use the main os.environ for anything.
        self.env = {
            'VOLTTRON_HOME': self.volttron_home,
            'PACKAGED_DIR': self.packaged_dir,
            'DEBUG_MODE': os.environ.get('DEBUG_MODE', ''),
            'DEBUG': os.environ.get('DEBUG', ''),
            'PATH': VOLTTRON_ROOT + ':' + os.environ['PATH']
        }
        self.volttron_root = VOLTTRON_ROOT

        volttron_exe = subprocess.check_output(['which', 'volttron']).strip()

        assert os.path.exists(volttron_exe)
        self.python = os.path.join(os.path.dirname(volttron_exe), 'python')
        assert os.path.exists(self.python)

        # By default no web server should be started.
        self.bind_web_address = None
        self.discovery_address = None
        self.jsonrpc_endpoint = None
        self.volttron_central_address = None
        self.instance_name = None
        self.serverkey = None

        self.p_process = None
        self.t_process = None

        self.started_agent_pids = []
        self.local_vip_address = None
        self.vip_address = None
        self.logit('Creating platform wrapper')

        # This was used when we are testing the SMAP historian.
        self.use_twistd = False

        # Added restricted code properties
        self.certsobj = None

        # Control whether the instance directory is cleaned up when shutdown.
        # if the environment variable DEBUG is set to a True value then the
        # instance is not cleaned up.
        self.skip_cleanup = False

        # This is used as command line entry replacement.  Especially working
        # with older 2.0 agents.
        self.opts = None

        keystorefile = os.path.join(self.volttron_home, 'keystore')
        self.keystore = KeyStore(keystorefile)
        self.keystore.generate()

    def logit(self, message):
        print('{}: {}'.format(self.volttron_home, message))

    def allow_all_connections(self):
        """ Add a /.*/ entry to the auth.json file.
        """
        entry = AuthEntry(credentials="/.*/")
        authfile = AuthFile(self.volttron_home + "/auth.json")
        try:
            authfile.add(entry)
        except AuthFileEntryAlreadyExists:
            pass

    def build_connection(self, peer=None, address=None, identity=None,
                         publickey=None, secretkey=None, serverkey=None,
                         capabilities=[], **kwargs):
        self.logit('Building connection to {}'.format(peer))
        self.allow_all_connections()

        if address is None:
            self.logit(
                'Default address was None so setting to current instances')
            address = self.vip_address
            serverkey = self.serverkey
        if serverkey is None:
            self.logit("serverkey wasn't set but the address was.")
            raise Exception("Invalid state.")

        if publickey is None or secretkey is None:
            self.logit('generating new public secret key pair')
            keyfile = tempfile.mktemp(".keys", "agent", self.volttron_home)
            keys = KeyStore(keyfile)
            keys.generate()
            publickey = keys.public
            secretkey = keys.secret
            entry = AuthEntry(capabilities=capabilities,
                              comments="Added by test",
                              credentials=keys.public)
            file = AuthFile(self.volttron_home + "/auth.json")
            file.add(entry)

        conn = Connection(address=address, peer=peer, publickey=publickey,
                          secretkey=secretkey, serverkey=serverkey,
                          volttron_home=self.volttron_home)

        return conn

    def build_agent(self, address=None, should_spawn=True, identity=None,
                    publickey=None, secretkey=None, serverkey=None,
                    agent_class=Agent, **kwargs):
        """ Build an agent connnected to the passed bus.

        By default the current instance that this class wraps will be the
        vip address of the agent.

        :param address:
        :param should_spawn:
        :param identity:
        :param publickey:
        :param secretkey:
        :param serverkey:
        :param agent_class: Agent class to build
        :return:
        """
        self.logit("Building generic agent.")

        use_ipc = kwargs.pop('use_ipc', False)

        if serverkey is None:
            serverkey = self.serverkey
        if publickey is None:
            self.logit('generating new public secret key pair')
            keyfile = tempfile.mktemp(".keys", "agent", self.volttron_home)
            keys = KeyStore(keyfile)
            keys.generate()
            publickey = keys.public
            secretkey = keys.secret

        if address is None:
            self.logit('Using vip-address ' + self.vip_address)
            address = self.vip_address

        if publickey and not serverkey:
            self.logit('using instance serverkey: {}'.format(self.publickey))
            serverkey = self.publickey

        agent = agent_class(address=address, identity=identity,
                            publickey=publickey, secretkey=secretkey,
                            serverkey=serverkey,
                            volttron_home=self.volttron_home,
                            **kwargs)
        self.logit('platformwrapper.build_agent.address: {}'.format(address))

        # Automatically add agent's credentials to auth.json file
        if publickey:
            self.logit('Adding publickey to auth.json')
            gevent.spawn(self._append_allow_curve_key, publickey)
            gevent.sleep(0.1)

        if should_spawn:
            self.logit('platformwrapper.build_agent spawning')
            event = gevent.event.Event()
            gevent.spawn(agent.core.run, event)  # .join(0)
            event.wait(timeout=2)

            hello = agent.vip.hello().get(timeout=.3)
            self.logit('Got hello response {}'.format(hello))
        agent.publickey = publickey
        return agent

    def _read_auth_file(self):
        auth_path = os.path.join(self.volttron_home, 'auth.json')
        try:
            with open(auth_path, 'r') as fd:
                data = strip_comments(FileObject(fd, close=False).read())
                if data:
                    auth = jsonapi.loads(data)
                else:
                    auth = {}
        except IOError:
            auth = {}
        if 'allow' not in auth:
            auth['allow'] = []
        return auth, auth_path

    def _append_allow_curve_key(self, publickey):
        entry = AuthEntry(credentials=publickey)
        authfile = AuthFile(self.volttron_home + "/auth.json")
        try:
            authfile.add(entry)
        except AuthFileEntryAlreadyExists:
            pass

    def add_vc(self):
        return add_vc_to_instance(self)

    def add_capabilities(self, publickey, capabilities):
        if isinstance(capabilities, basestring):
            capabilities = [capabilities]
        auth, auth_path = self._read_auth_file()
        cred = publickey
        allow = auth['allow']
        entry = next((item for item in allow if item['credentials'] == cred),
                     {})
        caps = entry.get('capabilities', [])
        entry['capabilities'] = list(set(caps + capabilities))

        with open(auth_path, 'w+') as fd:
            json.dump(auth, fd)

    def set_auth_dict(self, auth_dict):
        if auth_dict:
            with open(os.path.join(self.volttron_home, 'auth.json'), 'w') as fd:
                fd.write(json.dumps(auth_dict))

    def startup_platform(self, vip_address, auth_dict=None, use_twistd=False,
                         mode=UNRESTRICTED, bind_web_address=None,
                         volttron_central_address=None,
                         volttron_central_serverkey=None):

        # if not isinstance(vip_address, list):
        #     self.vip_address = [vip_address]
        # else:
        #     self.vip_address = vip_address

        self.vip_address = vip_address
        self.mode = mode
        self.bind_web_address = bind_web_address
        if self.bind_web_address:
            self.discovery_address = "{}/discovery/".format(
                self.bind_web_address)

            # Only available if vc is installed!
            self.jsonrpc_endpoint = "{}/jsonrpc".format(
                self.bind_web_address)

        enable_logging = self.env.get('ENABLE_LOGGING', False)
        debug_mode = self.env.get('DEBUG_MODE', False)
        if not debug_mode:
            debug_mode = self.env.get('DEBUG', False)
        self.skip_cleanup = self.env.get('SKIP_CLEANUP', False)
        if debug_mode:
            self.skip_cleanup = True
            enable_logging = True
        self.logit(
            "In start up platform enable_logging is {} ".format(enable_logging))
        assert self.mode in MODES, 'Invalid platform mode set: ' + str(mode)
        opts = None

        # see main.py for how we handle pub sub addresses.
        ipc = 'ipc://{}{}/run/'.format(
            '@' if sys.platform.startswith('linux') else '',
            self.volttron_home)
        self.local_vip_address = ipc + 'vip.socket'
        self.set_auth_dict(auth_dict)

        self.opts = {'verify_agents': False,
                     'volttron_home': self.volttron_home,
                     'vip_address': vip_address,
                     'vip_local_address': ipc + 'vip.socket',
                     'publish_address': ipc + 'publish',
                     'subscribe_address': ipc + 'subscribe',
                     'bind_web_address': bind_web_address,
                     'volttron_central_address': volttron_central_address,
                     'volttron_central_serverkey': volttron_central_serverkey,
                     'platform_name': None,
                     'log': os.path.join(self.volttron_home, 'volttron.log'),
                     'log_config': None,
                     'monitor': True,
                     'autostart': True,
                     'log_level': logging.DEBUG,
                     'verboseness': logging.DEBUG}

        pconfig = os.path.join(self.volttron_home, 'config')
        config = {}

        # Add platform's public key to known hosts file
        publickey = self.keystore.public
        known_hosts_file = os.path.join(self.volttron_home, 'known_hosts')
        known_hosts = KnownHostsStore(known_hosts_file)
        known_hosts.add(self.opts['vip_local_address'], publickey)
        known_hosts.add(self.opts['vip_address'], publickey)

        # Set up the configuration file based upon the passed parameters.
        parser = configparser.ConfigParser()
        parser.add_section('volttron')
        parser.set('volttron', 'vip-address', vip_address)
        if bind_web_address:
            parser.set('volttron', 'bind-web-address', bind_web_address)
        if volttron_central_address:
            parser.set('volttron', 'volttron-central-address',
                       volttron_central_address)
        if volttron_central_serverkey:
            parser.set('volttron', 'volttron-central-serverkey',
                       volttron_central_serverkey)
        if self.mode == UNRESTRICTED:
            # TODO Restricted code should set with volttron as contianer
            # if RESTRICTED_AVAILABLE:
            #     config['mobility'] = False
            #     config['resource-monitor'] = False
            #     config['verify'] = False
            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
                parser.write(cfg)

        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(self.volttron_home, 'certificates')

            print ("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)

            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
        else:
            raise PlatformWrapperError(
                "Invalid platform mode specified: {}".format(mode))

        log = os.path.join(self.volttron_home, 'volttron.log')
        if enable_logging:
            cmd = ['volttron', '-vv', '-l{}'.format(log)]
        else:
            cmd = ['volttron', '-l{}'.format(log)]

        print('process environment: {}'.format(self.env))
        print('popen params: {}'.format(cmd))
        self.p_process = Popen(cmd, env=self.env, stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)

        assert self.p_process is not None
        # A None value means that the process is still running.
        # A negative means that the process exited with an error.
        assert self.p_process.poll() is None

        self.serverkey = self.keystore.public
        assert self.serverkey
        agent = self.build_agent()

        has_control = False
        times = 0
        while not has_control and times < 10:
            times += 1
            try:
                has_control = agent.vip.peerlist().get(timeout=.2)
            except gevent.Timeout:
                pass

        if not has_control:
            self.shutdown_platform()
            raise "Couldn't connect to core platform!"

        if bind_web_address:
            times = 0
            has_discovery = False
            while times < 10:
                times += 1
                try:
                    resp = requests.get(self.discovery_address)
                    if resp.ok:
                        has_discovery = True
                        break
                except Exception as e:
                    gevent.sleep(0.1)
                    self.logit("Connection error found {}".format(e))
            if not has_discovery:
                raise "Couldn't connect to discovery platform."

        self.use_twistd = use_twistd

        # TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.volttron_home, TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = [TWISTED_START, "-n", "smap", tconfig]
            self.t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)

    def is_running(self):
        self.logit("PROCESS IS RUNNING: {}".format(self.p_process))
        return self.p_process is not None and self.p_process.poll() is None

    def twistd_is_running(self):
        return self.t_process is not None

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print ("wrapper.certsobj", self.certsobj.cert_dir)
        assert (
            auth.sign_as_creator(package, 'creator',
                                 certsobj=self.certsobj)), "Signing as {} failed.".format(
            'creator')

    def direct_sign_agentpackage_admin(self, package):
        assert (RESTRICTED), "Auth not available"
        assert (auth.sign_as_admin(package, 'admin',
                                   certsobj=self.certsobj)), "Signing as {} failed.".format(
            'admin')

    def direct_sign_agentpackage_initiator(self, package, config_file,
                                           contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file": config_file, "contract": contract}
        assert (auth.sign_as_initiator(package, 'initiator', files=files,
                                       certsobj=self.certsobj)), "Signing as {} failed.".format(
            'initiator')

    def _aip(self):
        opts = type('Options', (), self.opts)
        aip = AIPplatform(opts)
        aip.setup()
        return aip

    def _install_agent(self, wheel_file, start, vip_identity):
        self.logit('Creating channel for sending the agent.')
        gevent.sleep(0.3)
        self.logit('calling control install agent.')
        self.logit("VOLTTRON_HOME SETTING: {}".format(
            self.env['VOLTTRON_HOME']))
        env = self.env.copy()
        cmd = ['volttron-ctl', '-vv', 'install', wheel_file]
        if vip_identity:
            cmd.extend(['--vip-identity', vip_identity])
        self.logit("cmd: {}".format(cmd))
        res = subprocess.check_output(cmd, env=env)
        assert res, "failed to install wheel:{}".format(wheel_file)
        agent_uuid = res.split(' ')[-2]
        self.logit(agent_uuid)

        if start:
            self.start_agent(agent_uuid)

        return agent_uuid

    def install_multiple_agents(self, agent_configs):
        """
        Installs mutltiple agents on the platform.

        :param agent_configs:list
            A list of 3-tuple that allows the configuration of a platform
            in a single go.  The tuple order is
            1. path to the agent directory.
            2. configuration data (either file or json data)
            3. Whether the agent should be started or not.

        :return:list:
            A list of uuid's associated with the agents that were installed.


        :Note:
            In order for this method to be called the platform must be
            currently running.
        """
        if not self.is_running():
            raise PlatformWrapperError("Instance isn't running!")
        results = []

        for path, config, start in agent_configs:
            results = self.install_agent(agent_dir=path, config_file=config,
                                         start=start)

        return results

    def install_agent(self, agent_wheel=None, agent_dir=None, config_file=None,
                      start=True, vip_identity=None):
        """
        Install and optionally start an agent on the instance.

        This function allows installation from an agent wheel or an
        agent directory (NOT BOTH).  If an agent_wheel is specified then
        it is assumed to be ready for installation (has a config file).
        If an agent_dir is specified then a config_file file must be
        specified or if it is not specified then it is assumed that the
        file agent_dir/config is to be used as the configuration file.  If
        none of these exist then an assertion error will be thrown.

        This function will return with a uuid of the installed agent.

        :param agent_wheel:
        :param agent_dir:
        :param config_file:
        :param start:
        :param vip_identity:
        :return:
        """

        assert self.is_running(), "Instance must be running to install agent."
        assert agent_wheel or agent_dir, "Invalid agent_wheel or agent_dir."

        if agent_wheel:
            assert not agent_dir
            assert not config_file
            assert os.path.exists(agent_wheel)
            wheel_file = agent_wheel
            agent_uuid = self._install_agent(wheel_file, start, vip_identity)

        # Now if the agent_dir is specified.
        if agent_dir:
            assert not agent_wheel
            if isinstance(config_file, dict):
                from os.path import join, basename
                temp_config = join(self.volttron_home,
                                   basename(agent_dir) + "_config_file")
                with open(temp_config, "w") as fp:
                    fp.write(json.dumps(config_file))
                config_file = temp_config
            elif not config_file:
                if os.path.exists(os.path.join(agent_dir, "config")):
                    config_file = os.path.join(agent_dir, "config")
                else:
                    from os.path import join, basename
                    temp_config = join(self.volttron_home,
                                       basename(agent_dir) + "_config_file")
                    with open(temp_config, "w") as fp:
                        fp.write(json.dumps({}))
                    config_file = temp_config
            elif os.path.exists(config_file):
                pass  # config_file already set!
            else:
                raise ValueError("Can't determine correct config file.")

            script = os.path.join(self.volttron_root,
                                  "scripts/install-agent.py")
            cmd = [self.python, script,
                   "--volttron-home", self.volttron_home,
                   "--volttron-root", self.volttron_root,
                   "--agent-source", agent_dir,
                   "--config", config_file,
                   "--json"]

            if vip_identity:
                cmd.extend(["--vip-identity", vip_identity])
            if start:
                cmd.extend(["--start"])

            results = subprocess.check_output(cmd)

            # Because we are no longer silencing output from the install, the
            # the results object is now much more verbose.  Our assumption is
            # the line before the output we care about has WHEEL at the end
            # of it.
            new_results = ""
            found_wheel = False
            for line in results.split("\n"):
                if line.endswith("WHEEL"):
                    found_wheel = True
                elif found_wheel:
                    new_results += line
            results = new_results

            #
            # Response from results is expected as follows depending on
            # parameters, note this is a json string so parse to get dictionary.
            # {
            #     "started": true,
            #     "agent_pid": 26241,
            #     "starting": true,
            #     "agent_uuid": "ec1fd94e-922a-491f-9878-c392b24dbe50"
            # }
            assert results

            resultobj = jsonapi.loads(str(results))

            if start:
                assert resultobj['started']
            agent_uuid = resultobj['agent_uuid']

        assert agent_uuid is not None

        if start:
            assert self.is_agent_running(agent_uuid)

        return agent_uuid

    def start_agent(self, agent_uuid):
        self.logit('Starting agent {}'.format(agent_uuid))
        self.logit("VOLTTRON_HOME SETTING: {}".format(
            self.env['VOLTTRON_HOME']))
        cmd = ['volttron-ctl']
        cmd.extend(['start', agent_uuid])
        p = Popen(cmd, env=self.env,
                  stdout=sys.stdout, stderr=sys.stderr)
        p.wait()

        # Confirm agent running
        cmd = ['volttron-ctl']
        cmd.extend(['status', agent_uuid])
        res = subprocess.check_output(cmd, env=self.env)
        # 776 TODO: Timing issue where check fails
        time.sleep(.1)
        self.logit("Subprocess res is {}".format(res))
        assert 'running' in res
        pidpos = res.index('[') + 1
        pidend = res.index(']')
        pid = int(res[pidpos: pidend])

        assert psutil.pid_exists(pid), \
            "The pid associated with agent {} does not exist".format(pid)

        self.started_agent_pids.append(pid)
        return pid

    def stop_agent(self, agent_uuid):
        # Confirm agent running
        _log.debug("STOPPING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl']
            cmd.extend(['stop', agent_uuid])
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_pid(agent_uuid)

    def list_agents(self):
        agent = self.build_agent()
        print('PEER LIST: {}'.format(agent.vip.peerlist().get(timeout=10)))
        agent_list = agent.vip.rpc('control', 'list_agents').get(timeout=10)
        agent.core.stop(timeout=3)
        return agent_list

    def remove_agent(self, agent_uuid):
        """Remove the agent specified by agent_uuid"""
        _log.debug("REMOVING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl']
            cmd.extend(['remove', agent_uuid])
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_pid(agent_uuid)

    def is_agent_running(self, agent_uuid):
        return self.agent_pid(agent_uuid) is not None

    def agent_pid(self, agent_uuid):
        """
        Returns the pid of a running agent or None

        :param agent_uuid:
        :return:
        """
        # Confirm agent running
        cmd = ['volttron-ctl']
        cmd.extend(['status', agent_uuid])
        pid = None
        try:
            res = subprocess.check_output(cmd, env=self.env)
            try:
                pidpos = res.index('[') + 1
                pidend = res.index(']')
                pid = int(res[pidpos: pidend])
            except:
                pid = None
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))

        # Handle the following exception that seems to happen when getting a
        # pid of an agent during the platform shutdown phase.
        #
        # Logged from file platformwrapper.py, line 797
        #   AGENT             IDENTITY          TAG STATUS
        # Traceback (most recent call last):
        #   File "/usr/lib/python2.7/logging/__init__.py", line 882, in emit
        #     stream.write(fs % msg)
        #   File "/home/volttron/git/volttron/env/local/lib/python2.7/site-packages/_pytest/capture.py", line 244, in write
        #     self.buffer.write(obj)
        # ValueError: I/O operation on closed file
        except ValueError:
            pass
        # _log.debug("AGENT_PID: {}".format(pid))
        return pid

    def build_agentpackage(self, agent_dir, config_file={}):
        if isinstance(config_file, dict):
            cfg_path = os.path.join(agent_dir, "config_temp")
            with open(cfg_path, "w") as tmp_cfg:
                tmp_cfg.write(jsonapi.dumps(config_file))
            config_file = cfg_path

        # Handle relative paths from the volttron git directory.
        if not os.path.isabs(agent_dir):
            agent_dir = os.path.join(self.volttron_root, agent_dir)

        assert os.path.exists(config_file)
        assert os.path.exists(agent_dir)

        wheel_path = packaging.create_package(agent_dir,
                                              self.packaged_dir)
        packaging.add_files_to_package(wheel_path, {
            'config_file': os.path.join('./', config_file)
        })

        return wheel_path

    def confirm_agent_running(self, agent_name, max_retries=5,
                              timeout_seconds=2):
        running = False
        retries = 0
        while not running and retries < max_retries:
            status = self.test_aip.status_agents()
            print ("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running

    # def direct_stop_agent(self, agent_uuid):
    #     result = self.conn.call.stop_agent(agent_uuid)
    #     print result

    def shutdown_platform(self):
        """
        Stop platform here.  First grab a list of all of the agents that are
        running on the platform, then shutdown, then if any of the listed agent
        pids are still running then kill them.
        """

        # Handle cascading calls from multiple levels of fixtures.
        if self._instance_shutdown:
            return

        running_pids = []

        for agnt in self.list_agents():
            pid = self.agent_pid(agnt['uuid'])
            if pid is not None and int(pid) > 0:
                running_pids.append(int(pid))

        # First try and nicely shutdown the platform, which should clean all
        # of the agents up automatically.
        cmd = ['volttron-ctl']
        cmd.extend(['shutdown', '--platform'])
        try:
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError:
            if self.p_process is not None:
                try:
                    gevent.sleep(0.2)
                    self.p_process.terminate()
                    gevent.sleep(0.2)
                except OSError:
                    self.logit('Platform process was terminated.')
            else:
                self.logit("platform process was null")

        for pid in running_pids:
            if psutil.pid_exists(pid):
                self.logit("TERMINATING: {}".format(pid))
                proc = psutil.Process(pid)
                proc.terminate()

        if self.use_twistd and self.t_process is not None:
            self.t_process.kill()
            self.t_process.wait()
        elif self.use_twistd:
            self.logit("twistd process was null")

        if os.environ.get('PRINT_LOG'):
            logpath = os.path.join(self.volttron_home, 'volttron.log')
            if os.path.exists(logpath):
                print("************************* Begin {}".format(logpath))
                with open(logpath) as f:
                    for l in f.readlines():
                        print(l)
                print("************************* End {}".format(logpath))
            else:
                print("######################### No Log Exists: {}".format(
                    logpath
                ))
        if not self.skip_cleanup:
            self.logit('Removing {}'.format(self.volttron_home))
            shutil.rmtree(self.volttron_home, ignore_errors=True)
        self._instance_shutdown = True

    def __repr__(self):
        return str(self)

    def __str__(self):
        data = []
        data.append('volttron_home: {}'.format(self.volttron_home))
        return '\n'.join(data)
Exemplo n.º 4
0
class PlatformWrapper:
    def __init__(self):
        """ Initializes a new VOLTTRON instance

        Creates a temporary VOLTTRON_HOME directory with a packaged directory
        for agents that are built.
        """

        self.volttron_home = tempfile.mkdtemp()
        self.packaged_dir = os.path.join(self.volttron_home, "packaged")
        os.makedirs(self.packaged_dir)

        # in the context of this platform it is very important not to
        # use the main os.environ for anything.
        self.env = {
            'VOLTTRON_HOME': self.volttron_home,
            'PACKAGED_DIR': self.packaged_dir,
            'DEBUG_MODE': os.environ.get('DEBUG_MODE', ''),
            'DEBUG': os.environ.get('DEBUG', ''),
            'PATH': VOLTTRON_ROOT + ':' + os.environ['PATH']
        }

        # By default no web server should be started.
        self.bind_web_address = None
        self.discovery_address = None
        self.jsonrpc_endpoint = None
        self.volttron_central_address = None
        self.instance_name = None
        self.serverkey = None

        self.p_process = None
        self.t_process = None

        self.started_agent_pids = []
        self.local_vip_address = None
        self.vip_address = None
        self.encrypt = False
        self.logit('Creating platform wrapper')

        # This was used when we are testing the SMAP historian.
        self.use_twistd = False

        # Added restricted code properties
        self.certsobj = None

        # Control whether the instance directory is cleaned up when shutdown.
        # if the environment variable DEBUG is set to a True value then the
        # instance is not cleaned up.
        self.skip_cleanup = False

        # This is used as command line entry replacement.  Especially working
        # with older 2.0 agents.
        self.opts = None

        keystorefile = os.path.join(self.volttron_home, 'keystore')
        self.keystore = KeyStore(keystorefile)
        self.keystore.generate()

    def logit(self, message):
        print('{}: {}'.format(self.volttron_home, message))

    def allow_all_connections(self):
        """ Add a /.*/ entry to the auth.json file.
        """
        entry = AuthEntry(credentials="/.*/")
        authfile = AuthFile(self.volttron_home + "/auth.json")
        authfile.add(entry)

    def build_connection(self,
                         peer=None,
                         address=None,
                         identity=None,
                         publickey=None,
                         secretkey=None,
                         serverkey=None,
                         **kwargs):

        if self.encrypt:
            self.allow_all_connections()

        if address is None:
            address = self.vip_address
            serverkey = self.serverkey

        if publickey is None or secretkey is None:
            self.logit('generating new public secret key pair')
            keyfile = tempfile.mktemp(".keys", "agent", self.volttron_home)
            keys = KeyStore(keyfile)
            keys.generate()
            publickey = keys.public()
            secretkey = keys.secret()
        if self.encrypt:
            conn = Connection(address=address,
                              peer=peer,
                              publickey=publickey,
                              secretkey=secretkey,
                              serverkey=serverkey,
                              volttron_home=self.volttron_home)
        else:
            conn = Connection(address=self.local_vip_address,
                              peer=peer,
                              volttron_home=self.volttron_home)
        return conn

    def build_agent(self,
                    address=None,
                    should_spawn=True,
                    identity=None,
                    publickey=None,
                    secretkey=None,
                    serverkey=None,
                    agent_class=Agent,
                    **kwargs):
        """ Build an agent connnected to the passed bus.

        By default the current instance that this class wraps will be the
        vip address of the agent.

        :param address:
        :param should_spawn:
        :param identity:
        :param publickey:
        :param secretkey:
        :param serverkey:
        :param agent_class: Agent class to build
        :return:
        """
        self.logit("Building generic agent.")

        use_ipc = kwargs.pop('use_ipc', False)
        if self.encrypt:
            if serverkey is None:
                serverkey = self.serverkey
            if publickey is None:
                self.logit('generating new public secret key pair')
                keyfile = tempfile.mktemp(".keys", "agent", self.volttron_home)
                keys = KeyStore(keyfile)
                keys.generate()
                publickey = keys.public()
                secretkey = keys.secret()

        if address is None:
            if not self.encrypt:
                self.logit('Using IPC vip-address')
                address = "ipc://@" + self.volttron_home + "/run/vip.socket"
            else:
                self.logit('Using vip-address ' + self.vip_address)
                address = self.vip_address

        if publickey and not serverkey:
            self.logit('using instance serverkey: {}'.format(self.publickey))
            serverkey = self.publickey

        agent = agent_class(address=address,
                            identity=identity,
                            publickey=publickey,
                            secretkey=secretkey,
                            serverkey=serverkey,
                            volttron_home=self.volttron_home,
                            **kwargs)
        self.logit('platformwrapper.build_agent.address: {}'.format(address))

        # Automatically add agent's credentials to auth.json file
        if publickey:
            self.logit('Adding publickey to auth.json')
            gevent.spawn(self._append_allow_curve_key, publickey)
            gevent.sleep(0.1)

        if should_spawn:
            self.logit('platformwrapper.build_agent spawning')
            event = gevent.event.Event()
            gevent.spawn(agent.core.run, event)  #.join(0)
            event.wait(timeout=2)

            hello = agent.vip.hello().get(timeout=.3)
            self.logit('Got hello response {}'.format(hello))
        agent.publickey = publickey
        return agent

    def _read_auth_file(self):
        auth_path = os.path.join(self.volttron_home, 'auth.json')
        try:
            with open(auth_path, 'r') as fd:
                data = strip_comments(FileObject(fd, close=False).read())
                if data:
                    auth = jsonapi.loads(data)
                else:
                    auth = {}
        except IOError:
            auth = {}
        if not 'allow' in auth:
            auth['allow'] = []
        return auth, auth_path

    def _append_allow_curve_key(self, publickey):
        entry = AuthEntry(credentials=publickey)
        authfile = AuthFile(self.volttron_home + "/auth.json")
        authfile.add(entry)

    def add_vc(self):
        return add_vc_to_instance(self)

    def add_capabilities(self, publickey, capabilities):
        if isinstance(capabilities, basestring):
            capabilities = [capabilities]
        auth, auth_path = self._read_auth_file()
        cred = publickey
        allow = auth['allow']
        entry = next((item for item in allow if item['credentials'] == cred),
                     {})
        caps = entry.get('capabilities', [])
        entry['capabilities'] = list(set(caps + capabilities))

        with open(auth_path, 'w+') as fd:
            json.dump(auth, fd)

    def set_auth_dict(self, auth_dict):
        if auth_dict:
            with open(os.path.join(self.volttron_home, 'auth.json'),
                      'w') as fd:
                fd.write(json.dumps(auth_dict))

    def startup_platform(self,
                         vip_address,
                         auth_dict=None,
                         use_twistd=False,
                         mode=UNRESTRICTED,
                         encrypt=False,
                         bind_web_address=None,
                         volttron_central_address=None,
                         volttron_central_serverkey=None):

        # if not isinstance(vip_address, list):
        #     self.vip_address = [vip_address]
        # else:
        #     self.vip_address = vip_address

        self.vip_address = vip_address
        self.encrypt = encrypt
        self.mode = mode
        self.bind_web_address = bind_web_address
        if self.bind_web_address:
            self.discovery_address = "{}/discovery/".format(
                self.bind_web_address)

            # Only available if vc is installed!
            self.jsonrpc_endpoint = "{}/jsonrpc".format(self.bind_web_address)

        enable_logging = self.env.get('ENABLE_LOGGING', False)
        debug_mode = self.env.get('DEBUG_MODE', False)
        if not debug_mode:
            debug_mode = self.env.get('DEBUG', False)
        self.skip_cleanup = self.env.get('SKIP_CLEANUP', False)
        if debug_mode:
            self.skip_cleanup = True
            enable_logging = True
        self.logit("In start up platform enable_logging is {} ".format(
            enable_logging))
        assert self.mode in MODES, 'Invalid platform mode set: ' + str(mode)
        opts = None

        # see main.py for how we handle pub sub addresses.
        ipc = 'ipc://{}{}/run/'.format(
            '@' if sys.platform.startswith('linux') else '',
            self.volttron_home)
        self.local_vip_address = ipc + 'vip.socket'
        self.set_auth_dict(auth_dict)

        self.opts = {
            'verify_agents': False,
            'volttron_home': self.volttron_home,
            'vip_address': vip_address,
            'vip_local_address': ipc + 'vip.socket',
            'publish_address': ipc + 'publish',
            'subscribe_address': ipc + 'subscribe',
            'bind_web_address': bind_web_address,
            'volttron_central_address': volttron_central_address,
            'volttron_central_serverkey': volttron_central_serverkey,
            'platform_name': None,
            'developer_mode': not encrypt,
            'log': os.path.join(self.volttron_home, 'volttron.log'),
            'log_config': None,
            'monitor': True,
            'autostart': True,
            'log_level': logging.DEBUG,
            'verboseness': logging.DEBUG
        }

        pconfig = os.path.join(self.volttron_home, 'config')
        config = {}

        # Set up the configuration file based upon the passed parameters.
        parser = configparser.ConfigParser()
        parser.add_section('volttron')
        parser.set('volttron', 'vip-address', vip_address)
        if bind_web_address:
            parser.set('volttron', 'bind-web-address', bind_web_address)
        if volttron_central_address:
            parser.set('volttron', 'volttron-central-address',
                       volttron_central_address)

        if self.mode == UNRESTRICTED:
            # TODO Restricted code should set with volttron as contianer
            # if RESTRICTED_AVAILABLE:
            #     config['mobility'] = False
            #     config['resource-monitor'] = False
            #     config['verify'] = False
            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
                parser.write(cfg)

        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(self.volttron_home, 'certificates')

            print("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)

            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
            # opts = type('Options', (), {'resource-monitor':False,
            #                             'verify_agents': True,
            #                             'volttron_home': self.volttron_home})()
        else:
            raise PlatformWrapperError(
                "Invalid platform mode specified: {}".format(mode))

        log = os.path.join(self.volttron_home, 'volttron.log')
        if enable_logging:
            cmd = ['volttron', '-vv', '-l{}'.format(log)]
        else:
            cmd = ['volttron', '-l{}'.format(log)]

        if not encrypt:
            cmd.append('--developer-mode')
        print('process environment: {}'.format(self.env))
        print('popen params: {}'.format(cmd))
        self.p_process = Popen(cmd,
                               env=self.env,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)

        assert self.p_process is not None
        # A None value means that the process is still running.
        # A negative means that the process exited with an error.
        assert self.p_process.poll() is None

        self.serverkey = self.keystore.public()
        assert self.serverkey
        agent = self.build_agent()

        has_control = False
        times = 0
        while not has_control and times < 10:
            times += 1
            try:
                has_control = agent.vip.peerlist().get(timeout=.2)
            except gevent.Timeout:
                pass

        if not has_control:
            self.shutdown_platform()
            raise "Couldn't connect to core platform!"

        if bind_web_address:
            times = 0
            has_discovery = False
            while times < 10:
                times += 1
                try:
                    resp = requests.get(self.discovery_address)
                    if resp.ok:
                        has_discovery = True
                        break
                except Exception as e:
                    gevent.sleep(0.1)
                    self.logit("Connection error found {}".format(e))
            if not has_discovery:
                raise "Couldn't connect to discovery platform."

        self.use_twistd = use_twistd

        # TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.volttron_home,
                                   TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = [TWISTED_START, "-n", "smap", tconfig]
            self.t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)
            #self.t_process = subprocess.Popen(["twistd", "-n", "smap", "test-smap.ini"])

    def is_running(self):
        self.logit("PROCESS IS RUNNING: {}".format(self.p_process))
        return self.p_process is not None and self.p_process.poll() is None

    def twistd_is_running(self):
        return self.t_process is not None

    # def publish(self, topic, data):
    #     '''Publish data to a zmq context.
    #
    #     The publisher is goint to use the platform that is contained within
    #     this wrapper to write data to.
    #     '''
    #     if not self.zmq_context:
    #         self.zmq_context = zmq.Context()
    #     self.logit("binding publisher to: ", self.env['AGENT_PUB_ADDR'])
    #     pub = zmq.Socket(self.zmq_context, zmq.PUB)
    #     pub.bind(self.env['AGENT_PUB_ADDR'])
    #     pub.send_multipart([topic, data])

    # def fillout_file(self, filename, template, config_file):
    #
    #     try:
    #         config = json.loads(open(config_file, 'r').read())
    #     except Exception as e:
    #         sys.stderr.write (str(e))
    #         raise PlatformWrapperError("Could not load configuration file for tests")
    #
    #     config['tmpdir'] = self.tmpdir
    #
    #     outfile = os.path.join(self.tmpdir, filename)
    #     with closing(open(outfile, 'w')) as cfg:
    #         cfg.write(template.format(**config))
    #
    #     return outfile

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print("wrapper.certsobj", self.certsobj.cert_dir)
        assert (auth.sign_as_creator(
            package, 'creator',
            certsobj=self.certsobj)), "Signing as {} failed.".format('creator')

    def direct_sign_agentpackage_admin(self, package):
        assert (RESTRICTED), "Auth not available"
        assert (auth.sign_as_admin(
            package, 'admin',
            certsobj=self.certsobj)), "Signing as {} failed.".format('admin')

    def direct_sign_agentpackage_initiator(self, package, config_file,
                                           contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file": config_file, "contract": contract}
        assert (auth.sign_as_initiator(package,
                                       'initiator',
                                       files=files,
                                       certsobj=self.certsobj)
                ), "Signing as {} failed.".format('initiator')

    def _aip(self):
        opts = type('Options', (), self.opts)
        aip = AIPplatform(opts)
        aip.setup()
        return aip

    # TODO Remove when verified that the other method works properly.
    # def _install_agent(self, wheel_file, start, vip_identity):
    #     aip = self._aip()
    #     auuid = aip.install_agent(wheel_file, vip_identity=vip_identity)
    #     assert auuid is not None
    #     if start:
    #         self.logit('STARTING: {}'.format(wheel_file))
    #         status = self.start_agent(auuid)
    #         # aip.start_agent(auuid)
    #         # status = aip.agent_status(auuid)
    #         self.logit('STATUS NOW: {}'.format(status))
    #         assert status > 0
    #
    #     return auuid

    def _install_agent(self, wheel_file, start, vip_identity):

        agent = self.build_agent()
        self.logit('Creating channel for sending the agent.')
        channel_name = str(uuid.uuid4())
        channel = agent.vip.channel('control', channel_name)
        gevent.sleep(0.3)
        self.logit('calling control install agent.')
        result = agent.vip.rpc.call('control', 'install_agent', wheel_file,
                                    channel_name, vip_identity)
        self.logit('waiting for ready')
        response = channel.recv()
        if response != b'ready':
            raise ValueError(
                'Invalid channel protocol returned {}'.format(response))

        with open(wheel_file, 'rb') as fin:
            _log.debug('sending wheel to control.')
            while True:
                data = fin.read(8125)
                if not data:
                    _log.debug('Finished sending data')
                    break
                channel.send(data)

        _log.debug('sending done message.')
        channel.send('done')
        try:
            # must do this before channel closes or process will hang.
            auuid = result.get(timeout=10)
            _log.debug('closing channel')
        except gevent.Timeout:
            _log.error('Timeout in channel')
        finally:
            channel.close(linger=0)
            del channel

        if start:
            self.start_agent(auuid)

        return auuid

    def install_multiple_agents(self, agent_configs):
        """
        Installs mutltiple agents on the platform.

        :param agent_configs:list
            A list of 3-tuple that allows the configuration of a platform
            in a single go.  The tuple order is
            1. path to the agent directory.
            2. configuration data (either file or json data)
            3. Whether the agent should be started or not.

        :return:list:
            A list of uuid's associated with the agents that were installed.


        :Note:
            In order for this method to be called the platform must be
            currently running.
        """
        if not self.is_running():
            raise PlatformWrapperError("Instance isn't running!")
        results = []

        for path, config, start in agent_configs:
            results = self.install_agent(agent_dir=path,
                                         config_file=config,
                                         start=start)

        return results

    def install_agent(self,
                      agent_wheel=None,
                      agent_dir=None,
                      config_file=None,
                      start=True,
                      vip_identity=None):
        """ Install and optionally start an agent on the instance.

            This function allows installation from an agent wheel or an
            agent directory (NOT BOTH).  If an agent_wheel is specified then
            it is assumed to be ready for installation (has a config file).
            If an agent_dir is specified then a config_file file must be
            specified or if it is not specified then it is assumed that the
            file agent_dir/config is to be used as the configuration file.  If
            none of these exist then an assertion error will be thrown.

            This function will return with a uuid of the installed agent.
        :param agent_wheel:
        :param agent_dir:
        :param config_file:
        :param start:
        :param vip_identity:
        :return:
        """

        assert self.is_running(), "Instance must be running to install agent."
        assert agent_wheel or agent_dir, "Invalid agent_wheel or agent_dir."

        if agent_wheel:
            assert not agent_dir
            assert not config_file
            assert os.path.exists(agent_wheel)
            wheel_file = agent_wheel

        if agent_dir:
            assert not agent_wheel
            if isinstance(config_file, dict):
                from os.path import join, basename
                temp_config = join(self.volttron_home,
                                   basename(agent_dir) + "_config_file")
                with open(temp_config, "w") as fp:
                    fp.write(json.dumps(config_file))
                config_file = temp_config
            elif not config_file:
                assert os.path.exists(os.path.join(agent_dir, "config"))
                config_file = os.path.join(agent_dir, "config")
            elif os.path.exists(config_file):
                pass  # config_file already set!
            else:
                raise ValueError("Can't determine correct config file.")

            self.logit('Building agent package')
            wheel_file = self.build_agentpackage(agent_dir, config_file)
            assert wheel_file

        agent_uuid = self._install_agent(wheel_file, start, vip_identity)

        assert agent_uuid is not None

        if start:
            assert self.is_agent_running(agent_uuid)

        return agent_uuid

    def start_agent(self, agent_uuid):
        self.logit('Starting agent {}'.format(agent_uuid))
        self.logit("VOLTTRON_HOME SETTING: {}".format(
            self.env['VOLTTRON_HOME']))
        cmd = ['volttron-ctl', 'start', agent_uuid]
        p = Popen(cmd, env=self.env, stdout=sys.stdout, stderr=sys.stderr)
        p.wait()

        # Confirm agent running
        cmd = ['volttron-ctl', 'status', agent_uuid]
        res = subprocess.check_output(cmd, env=self.env)
        self.logit("Subprocess res is {}".format(res))
        assert 'running' in res
        pidpos = res.index('[') + 1
        pidend = res.index(']')
        pid = int(res[pidpos:pidend])

        self.started_agent_pids.append(pid)
        return int(pid)

    def stop_agent(self, agent_uuid):
        # Confirm agent running
        _log.debug("STOPPING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl', 'stop', agent_uuid]
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_status(agent_uuid)

    def list_agents(self):
        agent = self.build_agent()
        print('PEER LIST: {}'.format(agent.vip.peerlist().get(timeout=10)))
        agent_list = agent.vip.rpc('control', 'list_agents').get(timeout=10)
        agent.core.stop(timeout=3)
        return agent_list

    def remove_agent(self, agent_uuid):
        """Remove the agent specified by agent_uuid"""
        _log.debug("REMOVING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl', 'remove', agent_uuid]
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_status(agent_uuid)

    def is_agent_running(self, agent_uuid):
        return self.agent_status(agent_uuid) is not None

    def agent_status(self, agent_uuid):
        _log.debug("AGENT_STATUS: {}".format(agent_uuid))
        # Confirm agent running
        cmd = ['volttron-ctl', 'status', agent_uuid]
        pid = None
        try:
            res = subprocess.check_output(cmd, env=self.env)

            try:
                pidpos = res.index('[') + 1
                pidend = res.index(']')
                pid = int(res[pidpos:pidend])
            except:
                pid = None
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))

        return pid

    def build_agentpackage(self, agent_dir, config_file):
        assert os.path.exists(agent_dir)
        assert os.path.exists(config_file)
        wheel_path = packaging.create_package(agent_dir, self.packaged_dir)
        packaging.add_files_to_package(
            wheel_path, {'config_file': os.path.join('./', config_file)})

        return wheel_path

    # def direct_build_agentpackage(self, agent_dir):
    #     self.logit("Building agent_directory ", agent_dir)
    #     wheel_path = packaging.create_package(os.path.join('./', agent_dir),
    #                                           self.packaged_dir)
    #
    #     return wheel_path
    #
    # def direct_send_agent(self, package, target):
    #     pparams = [VCTRL, SEND_AGENT, target, package]
    #     print (pparams, "CWD", os.getcwd())
    #     send_process = subprocess.call(pparams, env=self.env)
    #     print ("Done sending to", target)
    #
    # def direct_configure_agentpackage(self, agent_wheel, config_file):
    #     packaging.add_files_to_package(agent_wheel, {
    #                             'config_file':os.path.join('./', config_file)
    #                         })
    #
    #


#     def direct_build_install_agent(self, agent_dir, config_file):
#         agent_wheel = self.build_agentpackage(agent_dir=agent_dir,
#             config_file=config_file)
#         self.direct_configure_agentpackage(agent_wheel, config_file)
#         assert(agent_wheel is not None,"Agent wheel was not built")
#
#         uuid = self.test_aip.install_agent(agent_wheel)
#         #aip volttron_home, verify_agents
#         return uuid
# #         conn.call.start_agent()

# def direct_build_install_run_agent(self, agent_dir, config_file):
#     agent_uuid = self.direct_build_install_agent(agent_dir, config_file)
#     self.direct_start_agent(agent_uuid)
#     return agent_uuid
#
# def direct_build_send_agent(self, agent_dir, config_file, target):
#     agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
#     self.direct_start_agent(agent_uuid)
#     return agent_uuid

    def confirm_agent_running(self,
                              agent_name,
                              max_retries=5,
                              timeout_seconds=2):
        running = False
        retries = 0
        while not running and retries < max_retries:
            status = self.test_aip.status_agents()
            print("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(
                    status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running

    # def direct_stop_agent(self, agent_uuid):
    #     result = self.conn.call.stop_agent(agent_uuid)
    #     print result

    def shutdown_platform(self):
        '''Stop platform here

           This function will shutdown the platform and attempt to kill any
           process that the platformwrapper has started.
        '''
        import signal
        self.logit('shutting down platform: PIDS: {}'.format(
            self.started_agent_pids))
        while self.started_agent_pids:
            pid = self.started_agent_pids.pop()
            self.logit('ending pid: {}'.format(pid))
            try:
                os.kill(pid, signal.SIGTERM)
            except:
                self.logit('could not kill: {} '.format(pid))

        if self.p_process is not None:
            try:
                gevent.sleep(0.2)
                self.p_process.terminate()
                gevent.sleep(0.2)
            except OSError:
                self.logit('Platform process was terminated.')
        else:
            self.logit("platform process was null")

        if self.use_twistd and self.t_process != None:
            self.t_process.kill()
            self.t_process.wait()
        elif self.use_twistd:
            self.logit("twistd process was null")
        if not self.skip_cleanup:
            self.logit('Removing {}'.format(self.volttron_home))
            shutil.rmtree(self.volttron_home, ignore_errors=True)
Exemplo n.º 5
0
class PlatformWrapper:
    def __init__(self):
        '''Initializes a new volttron environment

        Creates a temporary VOLTTRON_HOME directory with a packaged directory for
        agents that are built.
        '''
        self.__volttron_home = tempfile.mkdtemp()
        self.__packaged_dir = os.path.join(self.volttron_home, "packaged")
        os.makedirs(self.__packaged_dir)
        self.env = os.environ.copy()
        self.env['VOLTTRON_HOME'] = self.volttron_home

        # TODO: does changing os.environ affect the environment external to
        # this script?
        os.environ['VOLTTRON_HOME'] = self.volttron_home

        # By default no web server should be started.
        self.__bind_web_address = None

        self._p_process = None
        self._t_process = None
        self.__publickey = self.generate_key()
        self._started_pids = []
        self.__local_vip_address = None
        self.__vip_address = None
        self.logit('Creating platform wrapper')

    def logit(self, message):
        print('{}: {}'.format(self.volttron_home, message))

    @property
    def bind_web_address(self):
        return self.__bind_web_address

    @property
    def local_vip_address(self):
        return self.__local_vip_address

    @property
    def packaged_dir(self):
        return self.__packaged_dir

    @property
    def publickey(self):
        return self.__publickey

    @property
    def vip_address(self):
        return self.__vip_address

    @property
    def volttron_home(self):
        return self.__volttron_home

    def allow_all_connections(self):
        """ Add a CURVE:.* entry to the auth.json file.
        """
        entry = AuthEntry(credentials="/CURVE:.*/")
        authfile = AuthFile(self.volttron_home + "/auth.json")
        authfile.add(entry)

    def build_agent(self,
                    address=None,
                    should_spawn=True,
                    identity=None,
                    publickey=None,
                    secretkey=None,
                    serverkey=None,
                    generatekeys=False,
                    **kwargs):
        """ Build an agent connnected to the passed bus.

        By default the current instance that this class wraps will be the
        vip address of the agent.

        :param address:
        :param should_spawn:
        :param identity:
        :param publickey:
        :param secretkey:
        :param serverkey:
        :return:
        """
        self.logit("Building generic agent.")

        use_ipc = kwargs.pop('use_ipc', False)
        if address is None:
            if use_ipc:
                self.logit('Using IPC vip-address')
                address = "ipc://@" + self.volttron_home + "/run/vip.socket"
            else:
                self.logit('Using vip-address ' + self.vip_address)
                address = self.vip_address

        if generatekeys:
            self.logit('generating new public secret key pair')
            tf = tempfile.NamedTemporaryFile()
            ks = KeyStore(tf.name)
            ks.generate()
            publickey = ks.public()
            secretkey = ks.secret()

        if publickey and not serverkey:
            self.logit('using instance serverkey: {}'.format(self.publickey))
            serverkey = self.publickey

        agent = Agent(address=address,
                      identity=identity,
                      publickey=publickey,
                      secretkey=secretkey,
                      serverkey=serverkey,
                      **kwargs)
        self.logit('platformwrapper.build_agent.address: {}'.format(address))

        # Automatically add agent's credentials to auth.json file
        if publickey:
            self.logit('Adding publickey to auth.json')
            gevent.spawn(self._append_allow_curve_key, publickey)
            gevent.sleep(0.1)

        if should_spawn:
            self.logit('platformwrapper.build_agent spawning')
            event = gevent.event.Event()
            gevent.spawn(agent.core.run, event)  #.join(0)
            event.wait(timeout=2)

            hello = agent.vip.hello().get(timeout=.3)
            self.logit('Got hello response {}'.format(hello))

        return agent

    def generate_key(self):
        key = ''.join(zmq.curve_keypair())
        with open(os.path.join(self.volttron_home, 'curve.key'), 'w') as fd:
            fd.write(key)
        return encode_key(key[:40])  # public key

    def _read_auth_file(self):
        auth_path = os.path.join(self.volttron_home, 'auth.json')
        try:
            with open(auth_path, 'r') as fd:
                data = strip_comments(FileObject(fd, close=False).read())
                if data:
                    auth = jsonapi.loads(data)
                else:
                    auth = {}
        except IOError:
            auth = {}
        if not 'allow' in auth:
            auth['allow'] = []
        return auth, auth_path

    def _append_allow_curve_key(self, publickey):
        entry = AuthEntry(credentials="CURVE:{}".format(publickey))
        authfile = AuthFile(self.volttron_home + "/auth.json")
        authfile.add(entry)

    def add_capabilities(self, publickey, capabilities):
        if isinstance(capabilities, basestring):
            capabilities = [capabilities]
        auth, auth_path = self._read_auth_file()
        cred = 'CURVE:{}'.format(publickey)
        allow = auth['allow']
        entry = next((item for item in allow if item['credentials'] == cred),
                     {})
        caps = entry.get('capabilities', [])
        entry['capabilities'] = list(set(caps + capabilities))

        with open(auth_path, 'w+') as fd:
            json.dump(auth, fd)

    def set_auth_dict(self, auth_dict):
        if auth_dict:
            with open(os.path.join(self.volttron_home, 'auth.json'),
                      'w') as fd:
                fd.write(json.dumps(auth_dict))

    def startup_platform(self,
                         vip_address,
                         auth_dict=None,
                         use_twistd=False,
                         mode=UNRESTRICTED,
                         encrypt=False,
                         bind_web_address=None):
        # if not isinstance(vip_address, list):
        #     self.vip_address = [vip_address]
        # else:
        #     self.vip_address = vip_address

        self.vip_address = vip_address
        self.mode = mode
        self.bind_web_address = bind_web_address

        enable_logging = os.environ.get('ENABLE_LOGGING', False)
        debug_mode = os.environ.get('DEBUG_MODE', False)
        self.skip_cleanup = os.environ.get('SKIP_CLEANUP', False)
        if debug_mode:
            self.skip_cleanup = True
            enable_logging = True
        self.logit("In start up platform enable_logging is {} ".format(
            enable_logging))
        assert self.mode in MODES, 'Invalid platform mode set: ' + str(mode)
        opts = None

        # see main.py for how we handle pub sub addresses.
        ipc = 'ipc://{}{}/run/'.format(
            '@' if sys.platform.startswith('linux') else '',
            self.volttron_home)
        self.local_vip_address = ipc + 'vip.socket'
        if not encrypt:
            # Remove connection encryption
            with open(os.path.join(self.volttron_home, 'curve.key'), 'w'):
                pass

        self.set_auth_dict(auth_dict)

        self.opts = {
            'verify_agents': False,
            'volttron_home': self.volttron_home,
            'vip_address': vip_address,
            'vip_local_address': ipc + 'vip.socket',
            'publish_address': ipc + 'publish',
            'subscribe_address': ipc + 'subscribe',
            'bind_web_address': bind_web_address,
            'developer_mode': not encrypt,
            'log': os.path.join(self.volttron_home, 'volttron.log'),
            'log_config': None,
            'monitor': True,
            'autostart': True,
            'log_level': logging.DEBUG,
            'verboseness': logging.DEBUG
        }

        pconfig = os.path.join(self.volttron_home, 'config')
        config = {}

        parser = configparser.ConfigParser()
        parser.add_section('volttron')
        parser.set('volttron', 'vip-address', vip_address)
        if bind_web_address:
            parser.set('volttron', 'bind-web-address', bind_web_address)
        if self.mode == UNRESTRICTED:
            if RESTRICTED_AVAILABLE:
                config['mobility'] = False
                config['resource-monitor'] = False
                config['verify'] = False
            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
                parser.write(cfg)

        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(
                os.path.expanduser(self.env['VOLTTRON_HOME']), 'certificates')

            print("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)

            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
            opts = type(
                'Options', (), {
                    'resource-monitor': False,
                    'verify_agents': True,
                    'volttron_home': self.volttron_home
                })()
        else:
            raise PlatformWrapperError(
                "Invalid platform mode specified: {}".format(mode))

        log = os.path.join(self.env['VOLTTRON_HOME'], 'volttron.log')
        if enable_logging:
            cmd = ['volttron', '-vv', '-l{}'.format(log)]
        else:
            cmd = ['volttron', '-l{}'.format(log)]

        if self.opts['developer_mode']:
            cmd.append('--developer-mode')

        self._p_process = Popen(cmd,
                                env=self.env,
                                stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)

        assert self._p_process is not None
        # A None value means that the process is still running.
        # A negative means that the process exited with an error.
        assert self._p_process.poll() is None

        # # make sure we don't return too quickly.
        gevent.sleep(0.2)

        #os.environ['VOLTTRON_HOME'] = self.opts['volttron_home']
        #self._p_process = Process(target=start_volttron_process, args=(self.opts,))
        #self._p_process.daemon = True
        #self._p_process.start()

        gevent.sleep(0.2)
        self.use_twistd = use_twistd

        #TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.volttron_home,
                                   TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = [TWISTED_START, "-n", "smap", tconfig]
            self._t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)
            #self._t_process = subprocess.Popen(["twistd", "-n", "smap", "test-smap.ini"])

    def is_running(self):
        self.logit("PROCESS IS RUNNING: {}".format(self._p_process))
        return self._p_process is not None and self._p_process.poll() is None

    def twistd_is_running(self):
        return self._t_process is not None

    # def publish(self, topic, data):
    #     '''Publish data to a zmq context.
    #
    #     The publisher is goint to use the platform that is contained within
    #     this wrapper to write data to.
    #     '''
    #     if not self.zmq_context:
    #         self.zmq_context = zmq.Context()
    #     self.logit("binding publisher to: ", self.env['AGENT_PUB_ADDR'])
    #     pub = zmq.Socket(self.zmq_context, zmq.PUB)
    #     pub.bind(self.env['AGENT_PUB_ADDR'])
    #     pub.send_multipart([topic, data])

    # def fillout_file(self, filename, template, config_file):
    #
    #     try:
    #         config = json.loads(open(config_file, 'r').read())
    #     except Exception as e:
    #         sys.stderr.write (str(e))
    #         raise PlatformWrapperError("Could not load configuration file for tests")
    #
    #     config['tmpdir'] = self.tmpdir
    #
    #     outfile = os.path.join(self.tmpdir, filename)
    #     with closing(open(outfile, 'w')) as cfg:
    #         cfg.write(template.format(**config))
    #
    #     return outfile

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print("wrapper.certsobj", self.certsobj.cert_dir)
        assert (auth.sign_as_creator(
            package, 'creator',
            certsobj=self.certsobj)), "Signing as {} failed.".format('creator')

    def direct_sign_agentpackage_admin(self, package):
        assert (RESTRICTED), "Auth not available"
        assert (auth.sign_as_admin(
            package, 'admin',
            certsobj=self.certsobj)), "Signing as {} failed.".format('admin')

    def direct_sign_agentpackage_initiator(self, package, config_file,
                                           contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file": config_file, "contract": contract}
        assert (auth.sign_as_initiator(package,
                                       'initiator',
                                       files=files,
                                       certsobj=self.certsobj)
                ), "Signing as {} failed.".format('initiator')

    def _aip(self):
        opts = type('Options', (), self.opts)
        aip = AIPplatform(opts)
        aip.setup()
        return aip

    def _install_agent(self, wheel_file, start):
        aip = self._aip()
        auuid = aip.install_agent(wheel_file)
        assert auuid is not None
        if start:
            self.logit('STARTING: {}'.format(wheel_file))
            status = self.start_agent(auuid)
            #aip.start_agent(auuid)
            #status = aip.agent_status(auuid)
            self.logit('STATUS NOW: {}'.format(status))
            assert status > 0

        return auuid

    def install_multiple_agents(self, agent_configs):
        """
        Installs mutltiple agents on the platform.

        :param agent_configs:list
            A list of 3-tuple that allows the configuration of a platform
            in a single go.  The tuple order is
            1. path to the agent directory.
            2. configuration data (either file or json data)
            3. Whether the agent should be started or not.

        :return:list:
            A list of uuid's associated with the agents that were installed.


        :Note:
            In order for this method to be called the platform must be
            currently running.
        """
        if not self.is_running():
            raise PlatformWrapperError("Instance isn't running!")
        results = []

        for path, config, start in agent_configs:
            results = self.install_agent(agent_dir=path,
                                         config_file=config,
                                         start=start)

        return results

    def install_agent(self,
                      agent_wheel=None,
                      agent_dir=None,
                      config_file=None,
                      start=True):
        '''Install and optionally start an agent on the platform.

            This function allows installation from an agent wheel or an
            agent directory (NOT BOTH).  If an agent_wheel is specified then
            it is assumed to be ready for installation (has a config file).
            If an agent_dir is specified then a config_file file must be
            specified or if it is not specified then it is assumed that the
            file agent_dir/config is to be used as the configuration file.  If
            none of these exist then an assertion error will be thrown.

            This function will return with a uuid of the installed agent.
        '''

        assert self.is_running()
        assert agent_wheel or agent_dir

        if agent_wheel:
            assert not agent_dir
            assert not config_file
            assert os.path.exists(agent_wheel)
            wheel_file = agent_wheel

        if agent_dir:
            assert not agent_wheel
            if not config_file:
                assert os.path.exists(os.path.join(agent_dir, "config"))
                config_file = os.path.join(agent_dir, "config")
            else:
                if isinstance(config_file, dict):
                    from os.path import join, basename
                    temp_config = join(self.volttron_home,
                                       basename(agent_dir) + "_config_file")
                    with open(temp_config, "w") as fp:
                        fp.write(json.dumps(config_file))
                    config_file = temp_config
            self.logit('Building agent package')
            wheel_file = self.build_agentpackage(agent_dir, config_file)
            assert wheel_file

        agent_uuid = self._install_agent(wheel_file, start)

        assert agent_uuid is not None

        if start:
            assert self.is_agent_running(agent_uuid)

        return agent_uuid

    def start_agent(self, agent_uuid):
        self.logit('Starting agent {}'.format(agent_uuid))
        self.logit("VOLTTRONO_HOME SETTING: {}".format(
            os.environ['VOLTTRON_HOME']))
        cmd = ['volttron-ctl', 'start', agent_uuid]
        p = Popen(cmd, env=self.env, stdout=sys.stdout, stderr=sys.stderr)
        p.wait()

        # Confirm agent running
        cmd = ['volttron-ctl', 'status', agent_uuid]
        res = subprocess.check_output(cmd, env=self.env)
        assert 'running' in res
        pidpos = res.index('[') + 1
        pidend = res.index(']')
        pid = int(res[pidpos:pidend])

        self._started_pids.append(pid)
        return int(pid)

    def stop_agent(self, agent_uuid):
        # Confirm agent running
        _log.debug("STOPPING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl', 'stop', agent_uuid]
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_status(agent_uuid)

    def list_agents(self):
        aip = self._aip()
        return aip.list_agents()

    def remove_agent(self, agent_uuid):
        """Remove the agent specified by agent_uuid"""
        _log.debug("REMOVING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl', 'remove', agent_uuid]
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_status(agent_uuid)

    def is_agent_running(self, agent_uuid):
        return self.agent_status(agent_uuid) is not None

    def agent_status(self, agent_uuid):
        _log.debug("AGENT_STATUS: {}".format(agent_uuid))
        # Confirm agent running
        cmd = ['volttron-ctl', 'status', agent_uuid]
        pid = None
        try:
            res = subprocess.check_output(cmd, env=self.env)

            try:
                pidpos = res.index('[') + 1
                pidend = res.index(']')
                pid = int(res[pidpos:pidend])
            except:
                pid = None
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))

        return pid

    def build_agentpackage(self, agent_dir, config_file):
        assert os.path.exists(agent_dir)
        assert os.path.exists(config_file)
        wheel_path = packaging.create_package(agent_dir, self.packaged_dir)
        packaging.add_files_to_package(
            wheel_path, {'config_file': os.path.join('./', config_file)})

        return wheel_path

    # def direct_build_agentpackage(self, agent_dir):
    #     self.logit("Building agent_directory ", agent_dir)
    #     wheel_path = packaging.create_package(os.path.join('./', agent_dir),
    #                                           self.packaged_dir)
    #
    #     return wheel_path
    #
    # def direct_send_agent(self, package, target):
    #     pparams = [VCTRL, SEND_AGENT, target, package]
    #     print (pparams, "CWD", os.getcwd())
    #     send_process = subprocess.call(pparams, env=self.env)
    #     print ("Done sending to", target)
    #
    # def direct_configure_agentpackage(self, agent_wheel, config_file):
    #     packaging.add_files_to_package(agent_wheel, {
    #                             'config_file':os.path.join('./', config_file)
    #                         })
    #
    #


#     def direct_build_install_agent(self, agent_dir, config_file):
#         agent_wheel = self.build_agentpackage(agent_dir=agent_dir,
#             config_file=config_file)
#         self.direct_configure_agentpackage(agent_wheel, config_file)
#         assert(agent_wheel is not None,"Agent wheel was not built")
#
#         uuid = self.test_aip.install_agent(agent_wheel)
#         #aip volttron_home, verify_agents
#         return uuid
# #         conn.call.start_agent()

# def direct_build_install_run_agent(self, agent_dir, config_file):
#     agent_uuid = self.direct_build_install_agent(agent_dir, config_file)
#     self.direct_start_agent(agent_uuid)
#     return agent_uuid
#
# def direct_build_send_agent(self, agent_dir, config_file, target):
#     agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
#     self.direct_start_agent(agent_uuid)
#     return agent_uuid

    def confirm_agent_running(self,
                              agent_name,
                              max_retries=5,
                              timeout_seconds=2):
        running = False
        retries = 0
        while (not running and retries < max_retries):
            status = self.test_aip.status_agents()
            print("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(
                    status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running

    # def direct_stop_agent(self, agent_uuid):
    #     result = self.conn.call.stop_agent(agent_uuid)
    #     print result

    def shutdown_platform(self):
        '''Stop platform here

           This function will shutdown the platform and attempt to kill any
           process that the platformwrapper has started.
        '''
        import signal
        self.logit('shutting down platform: PIDS: {}'.format(
            self._started_pids))
        while self._started_pids:
            pid = self._started_pids.pop()
            self.logit('ending pid: {}'.format(pid))
            try:
                os.kill(pid, signal.SIGTERM)
            except:
                self.logit('could not kill: {} '.format(pid))
        if self._p_process != None:
            try:
                gevent.sleep(0.2)
                self._p_process.terminate()
                gevent.sleep(0.2)
            except OSError:
                self.logit('Platform process was terminated.')
        else:
            self.logit("platform process was null")

        if self.use_twistd and self._t_process != None:
            self._t_process.kill()
            self._t_process.wait()
        elif self.use_twistd:
            self.logit("twistd process was null")
Exemplo n.º 6
0
def run(command, use_sudo=False, user='', group='', freturn=False, err_to_out=False, input=None, force=False, **kwargs):
    """Execute command on host via ssh or subprocess.

    TODO: check on windows - maybe it will not work on it

    Factory uses pipes for communication with subprocess.
    So, there is no way to use popen and automatically write passwords for ssh and sudo on localhost,
    because "smart" programs like ssh and sudo uses tty directly.
    Also active tty required (needed check it) and for sudo uses "sudo -S".
    Alternatives:
      1) Use paramico like fabric = no ssh sockets.
      2) Use pty.fork, waitpid, execv as pexcpect and sh = only unix, no separated stderr, hard to communicate.
      3) Use ssh-copy-id like sh module recommended = ask passwords only one first time.
      4) Use sshpass like ansible = external dependencies.
      5) Use local ssh server and run all commands through it instead of popen.

    Args:
      command (str): command for executing
      use_sudo (bool): running with sudo prefix if True and current user not root, default is False
      user (str): username for sudo -u prefix
      group (str): group for sudo -g prefix
      freturn (bool): return tuple if True, else return str, default is False
      err_to_out (bool): redirect stderr to stdout if True, default is False
      input (str or tuple of str): str will be flushed to stdin after executed command, default is None
      force (bool): executing full operations even if envs.common.dry_run is True
      **kwargs (dict): add only for supporting dry-run replacing

    Return:
      str if freturn is False: string that contained all stdout messages
      tuple if freturn is True:
        string that contained all stdout messages
        string that contained all stderr
        int that mean return code of command

    """
    # hack for dry-run
    if envs.common.dry_run and not force:
        from dry_operations import run
        return run(command, use_sudo, user, group, freturn, err_to_out, input, **kwargs)

    logger = envs.connect.logger
    interactive = envs.common.interactive
    parallel = envs.common.parallel
    host_string = ''.join((envs.connect.user,
                           '@',
                           envs.connect.host))
    logger.debug('executing run function')
    logger.debug('arguments for executing and another locals: %s', locals())

    command = command_patching_for_sudo(command, use_sudo, user, group)

    # logging
    write_message_to_log(command, 'in: ')

    stderr = PIPE
    if err_to_out:
        stderr = STDOUT
    logger.debug('stderr: %s', stderr)
    # open new connect
    if envs.connect.host in envs.common.localhost:
        logger.debug('executing command %s with shell=True', command)
        p = Popen(command, stdout=PIPE, stderr=stderr, stdin=PIPE, shell=True)
    else:
        scommand = [
            envs.common.ssh_binary,
            envs.common.ssh_port_option,
            str(envs.connect.port),
            host_string,
        ]
        scommand += envs.common.ssh_args.split()
        scommand += envs.connect.con_args.split()
        scommand += [command]
        logger.debug('executing command %s', scommand)
        p = Popen(scommand, stdout=PIPE, stderr=stderr, stdin=PIPE)
    # flush input
    if input:
        if type(input) is str:
            input = [input]
        for s in input:
            s = str(s)
            if s[-1] not in ('\n', '\r'):
                s += '\n'
            logger.debug('flushing input %s', s)
            p.stdin.write(s)
            p.stdin.flush()
    # run another command
    if parallel:
        gevent.sleep(0)
        logger.debug('run another command with gevent.sleep(0)')
    # processing std loop
    threads = []
    if interactive:
        args = (p, copy(envs.common), copy(envs.connect))
        gin = gevent.spawn(in_loop, *args)
        logger.debug('executing in_loop with args %s', args)
        threads.append(gin)

    args = (p, copy(envs.common), copy(envs.connect))
    gout = gevent.spawn(out_loop, *args)
    logger.debug('executing out_loop with args %s', args)
    threads.append(gout)

    if not err_to_out:
        args = (p, copy(envs.common), copy(envs.connect), True)
        gerr = gevent.spawn(out_loop, *args)
        logger.debug('executing err_loop with args %s', args)
        threads.append(gerr)

    gevent.joinall(threads)
    logger.debug('child process has terminated with status %s', p.returncode)
    #TODO: check returncode if returncode==None
    sumout = gout.value
    sumerr = gerr.value if not err_to_out else ''
    status = p.returncode
    if p.poll() is None:
        p.terminate()
        p.kill()
    if freturn:
        logger.debug('return sumout %s, sumerr %s, status %s', sumout, sumerr, status)
        return (sumout, sumerr, status)
    logger.debug('return sumout %s', sumout)
    return sumout
Exemplo n.º 7
0
class PlatformWrapper:
    def __init__(self):
        '''Initializes a new volttron environment

        Creates a temporary VOLTTRON_HOME directory with a packaged directory for
        agents that are built.
        '''
        self.__volttron_home = tempfile.mkdtemp()
        self.__packaged_dir = os.path.join(self.volttron_home, "packaged")
        os.makedirs(self.__packaged_dir)
        self.env = os.environ.copy()
        self.env['VOLTTRON_HOME'] = self.volttron_home

        # TODO: does changing os.environ affect the environment external to
        # this script?
        os.environ['VOLTTRON_HOME'] = self.volttron_home

        # By default no web server should be started.
        self.__bind_web_address = None

        self._p_process = None
        self._t_process = None
        self.__publickey = self.generate_key()
        self._started_pids = []
        self.__local_vip_address = None
        self.__vip_address = None
        self.logit('Creating platform wrapper')

    def logit(self, message):
        print('{}: {}'.format(self.volttron_home, message))

    @property
    def bind_web_address(self):
        return self.__bind_web_address

    @property
    def local_vip_address(self):
        return self.__local_vip_address

    @property
    def packaged_dir(self):
        return self.__packaged_dir

    @property
    def publickey(self):
        return self.__publickey

    @property
    def vip_address(self):
        return self.__vip_address

    @property
    def volttron_home(self):
        return self.__volttron_home

    def allow_all_connections(self):
        """ Add a CURVE:.* entry to the auth.json file.
        """
        entry = AuthEntry(credentials="/CURVE:.*/")
        authfile = AuthFile(self.volttron_home+"/auth.json")
        authfile.add(entry)

    def build_agent(self, address=None, should_spawn=True, identity=None,
                    publickey=None, secretkey=None, serverkey=None,
                    generatekeys=False, **kwargs):
        """ Build an agent connnected to the passed bus.

        By default the current instance that this class wraps will be the
        vip address of the agent.

        :param address:
        :param should_spawn:
        :param identity:
        :param publickey:
        :param secretkey:
        :param serverkey:
        :return:
        """
        self.logit("Building generic agent.")

        use_ipc = kwargs.pop('use_ipc', False)
        if address is None:
            if use_ipc:
                self.logit('Using IPC vip-address')
                address = "ipc://@"+self.volttron_home+"/run/vip.socket"
            else:
                self.logit('Using vip-address '+self.vip_address)
                address = self.vip_address

        if generatekeys:
            self.logit('generating new public secret key pair')
            tf = tempfile.NamedTemporaryFile()
            ks = KeyStore(tf.name)
            ks.generate()
            publickey = ks.public()
            secretkey = ks.secret()

        if publickey and not serverkey:
            self.logit('using instance serverkey: {}'.format(self.publickey))
            serverkey = self.publickey

        agent = Agent(address=address, identity=identity, publickey=publickey,
                      secretkey=secretkey, serverkey=serverkey, **kwargs)
        self.logit('platformwrapper.build_agent.address: {}'.format(address))

        # Automatically add agent's credentials to auth.json file
        if publickey:
            self.logit('Adding publickey to auth.json')
            gevent.spawn(self._append_allow_curve_key, publickey)
            gevent.sleep(0.1)


        if should_spawn:
            self.logit('platformwrapper.build_agent spawning')
            event = gevent.event.Event()
            gevent.spawn(agent.core.run, event)#.join(0)
            event.wait(timeout=2)

            hello = agent.vip.hello().get(timeout=.3)
            self.logit('Got hello response {}'.format(hello))

        return agent

    def generate_key(self):
        key = ''.join(zmq.curve_keypair())
        with open(os.path.join(self.volttron_home, 'curve.key'), 'w') as fd:
            fd.write(key)
        return encode_key(key[:40]) # public key

    def _read_auth_file(self):
        auth_path = os.path.join(self.volttron_home, 'auth.json')
        try:
            with open(auth_path, 'r') as fd:
                data = strip_comments(FileObject(fd, close=False).read())
                if data:
                    auth = jsonapi.loads(data)
                else:
                    auth = {}
        except IOError:
            auth = {}
        if not 'allow' in auth:
            auth['allow'] = []
        return auth, auth_path

    def _append_allow_curve_key(self, publickey):
        entry = AuthEntry(credentials="CURVE:{}".format(publickey))
        authfile = AuthFile(self.volttron_home+"/auth.json")
        authfile.add(entry)

    def add_capabilities(self, publickey, capabilities):
        if isinstance(capabilities, basestring):
            capabilities = [capabilities]
        auth, auth_path = self._read_auth_file()
        cred = 'CURVE:{}'.format(publickey)
        allow = auth['allow']
        entry = next((item for item in allow if item['credentials'] == cred), {})
        caps = entry.get('capabilities', [])
        entry['capabilities'] = list(set(caps + capabilities))

        with open(auth_path, 'w+') as fd:
            json.dump(auth, fd)

    def set_auth_dict(self, auth_dict):
        if auth_dict:
            with open(os.path.join(self.volttron_home, 'auth.json'), 'w') as fd:
                fd.write(json.dumps(auth_dict))

    def startup_platform(self, vip_address, auth_dict=None, use_twistd=False,
        mode=UNRESTRICTED, encrypt=False, bind_web_address=None):
        # if not isinstance(vip_address, list):
        #     self.vip_address = [vip_address]
        # else:
        #     self.vip_address = vip_address

        self.vip_address = vip_address
        self.mode = mode
        self.bind_web_address = bind_web_address

        enable_logging = os.environ.get('ENABLE_LOGGING', False)
        debug_mode = os.environ.get('DEBUG_MODE', False)
        self.skip_cleanup = os.environ.get('SKIP_CLEANUP', False)
        if debug_mode:
            self.skip_cleanup = True
            enable_logging = True
        self.logit("In start up platform enable_logging is {} ".format(enable_logging))
        assert self.mode in MODES, 'Invalid platform mode set: '+str(mode)
        opts = None

        # see main.py for how we handle pub sub addresses.
        ipc = 'ipc://{}{}/run/'.format(
            '@' if sys.platform.startswith('linux') else '',
            self.volttron_home)
        self.local_vip_address = ipc + 'vip.socket'
        if not encrypt:
            # Remove connection encryption
            with open(os.path.join(self.volttron_home, 'curve.key'), 'w'):
                pass

        self.set_auth_dict(auth_dict)

        self.opts = {'verify_agents': False,
                'volttron_home': self.volttron_home,
                'vip_address': vip_address,
                'vip_local_address': ipc + 'vip.socket',
                'publish_address': ipc + 'publish',
                'subscribe_address': ipc + 'subscribe',
                'bind_web_address': bind_web_address,
                'developer_mode': not encrypt,
                'log': os.path.join(self.volttron_home,'volttron.log'),
                'log_config': None,
                'monitor': True,
                'autostart': True,
                'log_level': logging.DEBUG,
                'verboseness': logging.DEBUG}

        pconfig = os.path.join(self.volttron_home, 'config')
        config = {}

        parser =  configparser.ConfigParser()
        parser.add_section('volttron')
        parser.set('volttron', 'vip-address', vip_address)
        if bind_web_address:
            parser.set('volttron', 'bind-web-address', bind_web_address)
        if self.mode == UNRESTRICTED:
            if RESTRICTED_AVAILABLE:
                config['mobility'] = False
                config['resource-monitor'] = False
                config['verify'] = False
            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_UNRESTRICTED.format(**config))
                parser.write(cfg)


        elif self.mode == RESTRICTED:
            if not RESTRICTED_AVAILABLE:
                raise ValueError("restricted is not available.")

            certsdir = os.path.join(os.path.expanduser(self.env['VOLTTRON_HOME']),
                                    'certificates')

            print ("certsdir", certsdir)
            self.certsobj = certs.Certs(certsdir)


            with closing(open(pconfig, 'wb')) as cfg:
                cfg.write(PLATFORM_CONFIG_RESTRICTED.format(**config))
            opts = type('Options', (), {'resource-monitor':False,
                                        'verify_agents': True,
                                        'volttron_home': self.volttron_home})()
        else:
            raise PlatformWrapperError("Invalid platform mode specified: {}".format(mode))

        log = os.path.join(self.env['VOLTTRON_HOME'], 'volttron.log')
        if enable_logging:
            cmd = ['volttron', '-vv', '-l{}'.format(log)]
        else:
            cmd = ['volttron', '-l{}'.format(log)]

        if self.opts['developer_mode']:
            cmd.append('--developer-mode')

        self._p_process = Popen(cmd, env=self.env, stdout=subprocess.PIPE,
                                stderr=subprocess.PIPE)

        assert self._p_process is not None
        # A None value means that the process is still running.
        # A negative means that the process exited with an error.
        assert self._p_process.poll() is None

        # # make sure we don't return too quickly.
        gevent.sleep(0.2)

        #os.environ['VOLTTRON_HOME'] = self.opts['volttron_home']
        #self._p_process = Process(target=start_volttron_process, args=(self.opts,))
        #self._p_process.daemon = True
        #self._p_process.start()

        gevent.sleep(0.2)
        self.use_twistd = use_twistd

        #TODO: Revise this to start twistd with platform.
        if self.use_twistd:
            tconfig = os.path.join(self.volttron_home, TMP_SMAP_CONFIG_FILENAME)

            with closing(open(tconfig, 'w')) as cfg:
                cfg.write(TWISTED_CONFIG.format(**config))

            tparams = [TWISTED_START, "-n", "smap", tconfig]
            self._t_process = subprocess.Popen(tparams, env=self.env)
            time.sleep(5)
            #self._t_process = subprocess.Popen(["twistd", "-n", "smap", "test-smap.ini"])

    def is_running(self):
        self.logit("PROCESS IS RUNNING: {}".format(self._p_process))
        return self._p_process is not None and self._p_process.poll() is None

    def twistd_is_running(self):
        return self._t_process is not None

    # def publish(self, topic, data):
    #     '''Publish data to a zmq context.
    #
    #     The publisher is goint to use the platform that is contained within
    #     this wrapper to write data to.
    #     '''
    #     if not self.zmq_context:
    #         self.zmq_context = zmq.Context()
    #     self.logit("binding publisher to: ", self.env['AGENT_PUB_ADDR'])
    #     pub = zmq.Socket(self.zmq_context, zmq.PUB)
    #     pub.bind(self.env['AGENT_PUB_ADDR'])
    #     pub.send_multipart([topic, data])

    # def fillout_file(self, filename, template, config_file):
    #
    #     try:
    #         config = json.loads(open(config_file, 'r').read())
    #     except Exception as e:
    #         sys.stderr.write (str(e))
    #         raise PlatformWrapperError("Could not load configuration file for tests")
    #
    #     config['tmpdir'] = self.tmpdir
    #
    #     outfile = os.path.join(self.tmpdir, filename)
    #     with closing(open(outfile, 'w')) as cfg:
    #         cfg.write(template.format(**config))
    #
    #     return outfile

    def direct_sign_agentpackage_creator(self, package):
        assert (RESTRICTED), "Auth not available"
        print ("wrapper.certsobj", self.certsobj.cert_dir)
        assert(auth.sign_as_creator(package, 'creator', certsobj=self.certsobj)), "Signing as {} failed.".format('creator')


    def direct_sign_agentpackage_admin(self, package):
        assert (RESTRICTED), "Auth not available"
        assert(auth.sign_as_admin(package, 'admin', certsobj=self.certsobj)), "Signing as {} failed.".format('admin')


    def direct_sign_agentpackage_initiator(self, package, config_file, contract):
        assert (RESTRICTED), "Auth not available"
        files = {"config_file":config_file,"contract":contract}
        assert(auth.sign_as_initiator(package, 'initiator', files=files,
                                      certsobj=self.certsobj)), "Signing as {} failed.".format('initiator')

    def _aip(self):
        opts = type('Options', (), self.opts)
        aip = AIPplatform(opts)
        aip.setup()
        return aip

    def _install_agent(self, wheel_file, start):
        aip = self._aip()
        auuid = aip.install_agent(wheel_file)
        assert auuid is not None
        if start:
            self.logit('STARTING: {}'.format(wheel_file))
            status = self.start_agent(auuid)
            #aip.start_agent(auuid)
            #status = aip.agent_status(auuid)
            self.logit('STATUS NOW: {}'.format(status))
            assert status > 0

        return auuid

    def install_multiple_agents(self, agent_configs):
        """
        Installs mutltiple agents on the platform.

        :param agent_configs:list
            A list of 3-tuple that allows the configuration of a platform
            in a single go.  The tuple order is
            1. path to the agent directory.
            2. configuration data (either file or json data)
            3. Whether the agent should be started or not.

        :return:list:
            A list of uuid's associated with the agents that were installed.


        :Note:
            In order for this method to be called the platform must be
            currently running.
        """
        if not self.is_running():
            raise PlatformWrapperError("Instance isn't running!")
        results = []

        for path, config, start  in agent_configs:
            results = self.install_agent(agent_dir=path, config_file=config,
                                         start=start)

        return results

    def install_agent(self, agent_wheel=None, agent_dir=None, config_file=None,
        start=True):
        '''Install and optionally start an agent on the platform.

            This function allows installation from an agent wheel or an
            agent directory (NOT BOTH).  If an agent_wheel is specified then
            it is assumed to be ready for installation (has a config file).
            If an agent_dir is specified then a config_file file must be
            specified or if it is not specified then it is assumed that the
            file agent_dir/config is to be used as the configuration file.  If
            none of these exist then an assertion error will be thrown.

            This function will return with a uuid of the installed agent.
        '''

        assert self.is_running()
        assert agent_wheel or agent_dir

        if agent_wheel:
            assert not agent_dir
            assert not config_file
            assert os.path.exists(agent_wheel)
            wheel_file = agent_wheel

        if agent_dir:
            assert not agent_wheel
            if not config_file:
                assert os.path.exists(os.path.join(agent_dir, "config"))
                config_file = os.path.join(agent_dir, "config")
            else:
                if isinstance(config_file, dict):
                    from os.path import join, basename
                    temp_config=join(self.volttron_home, basename(agent_dir) + "_config_file")
                    with open(temp_config,"w") as fp:
                        fp.write(json.dumps(config_file))
                    config_file = temp_config
            self.logit('Building agent package')
            wheel_file = self.build_agentpackage(agent_dir, config_file)
            assert wheel_file

        agent_uuid = self._install_agent(wheel_file, start)

        assert agent_uuid is not None

        if start:
            assert self.is_agent_running(agent_uuid)

        return agent_uuid

    def start_agent(self, agent_uuid):
        self.logit('Starting agent {}'.format(agent_uuid))
        self.logit("VOLTTRONO_HOME SETTING: {}".format(os.environ['VOLTTRON_HOME']))
        cmd = ['volttron-ctl', 'start', agent_uuid]
        p = Popen(cmd, env=self.env,
                  stdout=sys.stdout, stderr=sys.stderr)
        p.wait()

        # Confirm agent running
        cmd = ['volttron-ctl', 'status', agent_uuid]
        res = subprocess.check_output(cmd, env=self.env)
        assert 'running' in res
        pidpos = res.index('[') + 1
        pidend = res.index(']')
        pid = int(res[pidpos: pidend])

        self._started_pids.append(pid)
        return int(pid)

    def stop_agent(self, agent_uuid):
        # Confirm agent running
        _log.debug("STOPPING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl', 'stop', agent_uuid]
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_status(agent_uuid)

    def list_agents(self):
        aip = self._aip()
        return aip.list_agents()

    def remove_agent(self, agent_uuid):
        """Remove the agent specified by agent_uuid"""
        _log.debug("REMOVING AGENT: {}".format(agent_uuid))
        try:
            cmd = ['volttron-ctl', 'remove', agent_uuid]
            res = subprocess.check_output(cmd, env=self.env)
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))
        return self.agent_status(agent_uuid)

    def is_agent_running(self, agent_uuid):
        return self.agent_status(agent_uuid) is not None

    def agent_status(self, agent_uuid):
        _log.debug("AGENT_STATUS: {}".format(agent_uuid))
        # Confirm agent running
        cmd = ['volttron-ctl', 'status', agent_uuid]
        pid = None
        try:
            res = subprocess.check_output(cmd, env=self.env)

            try:
                pidpos = res.index('[') + 1
                pidend = res.index(']')
                pid = int(res[pidpos: pidend])
            except:
                pid = None
        except CalledProcessError as ex:
            _log.error("Exception: {}".format(ex))

        return pid

    def build_agentpackage(self, agent_dir, config_file):
        assert os.path.exists(agent_dir)
        assert os.path.exists(config_file)
        wheel_path = packaging.create_package(agent_dir,
                                              self.packaged_dir)
        packaging.add_files_to_package(wheel_path, {
                'config_file': os.path.join('./', config_file)
            })

        return wheel_path

    # def direct_build_agentpackage(self, agent_dir):
    #     self.logit("Building agent_directory ", agent_dir)
    #     wheel_path = packaging.create_package(os.path.join('./', agent_dir),
    #                                           self.packaged_dir)
    #
    #     return wheel_path
    #
    # def direct_send_agent(self, package, target):
    #     pparams = [VCTRL, SEND_AGENT, target, package]
    #     print (pparams, "CWD", os.getcwd())
    #     send_process = subprocess.call(pparams, env=self.env)
    #     print ("Done sending to", target)
    #
    # def direct_configure_agentpackage(self, agent_wheel, config_file):
    #     packaging.add_files_to_package(agent_wheel, {
    #                             'config_file':os.path.join('./', config_file)
    #                         })
    #
    #

#     def direct_build_install_agent(self, agent_dir, config_file):
#         agent_wheel = self.build_agentpackage(agent_dir=agent_dir,
#             config_file=config_file)
#         self.direct_configure_agentpackage(agent_wheel, config_file)
#         assert(agent_wheel is not None,"Agent wheel was not built")
#
#         uuid = self.test_aip.install_agent(agent_wheel)
#         #aip volttron_home, verify_agents
#         return uuid
# #         conn.call.start_agent()



    # def direct_build_install_run_agent(self, agent_dir, config_file):
    #     agent_uuid = self.direct_build_install_agent(agent_dir, config_file)
    #     self.direct_start_agent(agent_uuid)
    #     return agent_uuid
    #
    # def direct_build_send_agent(self, agent_dir, config_file, target):
    #     agent_uuid = self.direct_buid_install_agent(agent_dir, config_file)
    #     self.direct_start_agent(agent_uuid)
    #     return agent_uuid


    def confirm_agent_running(self, agent_name, max_retries=5, timeout_seconds=2):
        running = False
        retries = 0
        while (not running and retries < max_retries):
            status = self.test_aip.status_agents()
            print ("Status", status)
            if len(status) > 0:
                status_name = status[0][1]
                assert status_name == agent_name

                assert len(status[0][2]) == 2, 'Unexpected agent status message'
                status_agent_status = status[0][2][1]
                running = not isinstance(status_agent_status, int)
            retries += 1
            time.sleep(timeout_seconds)
        return running


    # def direct_stop_agent(self, agent_uuid):
    #     result = self.conn.call.stop_agent(agent_uuid)
    #     print result


    def shutdown_platform(self):
        '''Stop platform here

           This function will shutdown the platform and attempt to kill any
           process that the platformwrapper has started.
        '''
        import signal
        self.logit('shutting down platform: PIDS: {}'.format(self._started_pids))
        while self._started_pids:
            pid = self._started_pids.pop()
            self.logit('ending pid: {}'.format(pid))
            try:
                os.kill(pid, signal.SIGTERM)
            except:
                self.logit('could not kill: {} '.format(pid))
        if self._p_process != None:
            try:
                gevent.sleep(0.2)
                self._p_process.terminate()
                gevent.sleep(0.2)
            except OSError:
                self.logit('Platform process was terminated.')
        else:
            self.logit("platform process was null")



        if self.use_twistd and self._t_process != None:
            self._t_process.kill()
            self._t_process.wait()
        elif self.use_twistd:
            self.logit("twistd process was null")
Exemplo n.º 8
0
class TaskExecutor(object):
    def __init__(self, balancer, index):
        self.balancer = balancer
        self.index = index
        self.task = None
        self.proc = None
        self.pid = None
        self.conn = None
        self.state = None
        self.key = str(uuid.uuid4())
        self.checked_in = Event()
        self.result = AsyncResult()
        self.exiting = False
        self.thread = gevent.spawn(self.executor)

    def checkin(self, conn):
        self.balancer.logger.debug('Check-in of worker #{0} (key {1})'.format(self.index, self.key))
        self.conn = conn
        self.state = WorkerState.IDLE
        self.checked_in.set()

    def get_status(self):
        if not self.conn:
            return None

        try:
            st = TaskStatus(0)
            if issubclass(self.task.clazz, MasterProgressTask):
                progress_subtask_info = self.conn.call_client_sync(
                    'taskproxy.get_master_progress_info'
                )
                if progress_subtask_info['increment_progress'] != 0:
                    progress_subtask_info['progress'] += progress_subtask_info['increment_progress']
                    progress_subtask_info['increment_progress'] = 0
                    self.conn.call_client_sync(
                        'taskproxy.set_master_progress_detail',
                        {
                            'progress': progress_subtask_info['progress'],
                            'increment_progress': progress_subtask_info['increment_progress']
                        }
                    )
                if progress_subtask_info['active_tids']:
                    progress_to_increment = 0
                    concurent_weight = progress_subtask_info['concurent_subtask_detail']['average_weight']
                    for tid in progress_subtask_info['concurent_subtask_detail']['tids']:
                        subtask_status = self.balancer.get_task(tid).executor.get_status()
                        progress_to_increment += subtask_status.percentage * concurent_weight * \
                            progress_subtask_info['subtask_weights'][str(tid)]
                    for tid in set(progress_subtask_info['active_tids']).symmetric_difference(
                        set(progress_subtask_info['concurent_subtask_detail']['tids'])
                    ):
                        subtask_status = self.balancer.get_task(tid).executor.get_status()
                        progress_to_increment += subtask_status.percentage * \
                            progress_subtask_info['subtask_weights'][str(tid)]
                    progress_subtask_info['progress'] += int(progress_to_increment)
                    if progress_subtask_info['pass_subtask_details']:
                        progress_subtask_info['message'] = subtask_status.message
                st = TaskStatus(
                    progress_subtask_info['progress'], progress_subtask_info['message']
                )
            else:
                st.__setstate__(self.conn.call_client_sync('taskproxy.get_status'))
            return st

        except RpcException as err:
            self.balancer.logger.error(
                "Cannot obtain status from task #{0}: {1}".format(self.task.id, str(err))
            )
            self.proc.terminate()

    def put_status(self, status):
        # Try to collect rusage at this point, when process is still alive
        try:
            kinfo = bsd.kinfo_getproc(self.pid)
            self.task.rusage = kinfo.rusage
        except LookupError:
            pass

        if status['status'] == 'ROLLBACK':
            self.task.set_state(TaskState.ROLLBACK)

        if status['status'] == 'FINISHED':
            self.result.set(status['result'])

        if status['status'] == 'FAILED':
            error = status['error']
            cls = TaskException

            if error['type'] == 'task.TaskAbortException':
                cls = TaskAbortException

            if error['type'] == 'ValidationException':
                cls = ValidationException

            self.result.set_exception(cls(
                code=error['code'],
                message=error['message'],
                stacktrace=error['stacktrace'],
                extra=error.get('extra')
            ))

    def put_warning(self, warning):
        self.task.add_warning(warning)

    def run(self, task):
        self.result = AsyncResult()
        self.task = task
        self.task.set_state(TaskState.EXECUTING)

        filename = None
        module_name = inspect.getmodule(task.clazz).__name__
        for dir in self.balancer.dispatcher.plugin_dirs:
            found = False
            try:
                for root, _, files in os.walk(dir):
                    file = first_or_default(lambda f: module_name in f, files)
                    if file:
                        filename = os.path.join(root, file)
                        found = True
                        break

                if found:
                    break
            except FileNotFoundError:
                continue

        self.conn.call_client_sync('taskproxy.run', {
            'id': task.id,
            'class': task.clazz.__name__,
            'filename': filename,
            'args': task.args,
            'debugger': task.debugger,
            'environment': task.environment
        })

        try:
            self.result.get()
        except BaseException as e:
            if not isinstance(e, TaskException):
                self.balancer.dispatcher.report_error(
                    'Task {0} raised exception other than TaskException'.format(self.task.name),
                    e
                )

            if isinstance(e, TaskAbortException):
                self.task.set_state(TaskState.ABORTED, TaskStatus(0, 'aborted'))
            else:
                self.task.error = serialize_error(e)
                self.task.set_state(TaskState.FAILED, TaskStatus(0, str(e), extra={
                    "stacktrace": traceback.format_exc()
                }))

            self.task.ended.set()
            self.balancer.task_exited(self.task)
            self.state = WorkerState.IDLE
            return

        self.task.result = self.result.value
        self.task.set_state(TaskState.FINISHED, TaskStatus(100, ''))
        self.task.ended.set()
        self.balancer.task_exited(self.task)
        self.state = WorkerState.IDLE

    def abort(self):
        self.balancer.logger.info("Trying to abort task #{0}".format(self.task.id))
        # Try to abort via RPC. If this fails, kill process
        try:
            self.conn.call_client_sync('taskproxy.abort')
        except RpcException as err:
            self.balancer.logger.warning("Failed to abort task #{0} gracefully: {1}".format(self.task.id, str(err)))
            self.balancer.logger.warning("Killing process {0}".format(self.pid))
            self.proc.terminate()

    def executor(self):
        while not self.exiting:
            try:
                self.proc = Popen(
                    [TASKWORKER_PATH, self.key],
                    close_fds=True,
                    preexec_fn=os.setpgrp,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT)

                self.pid = self.proc.pid
                self.balancer.logger.debug('Started executor #{0} as PID {1}'.format(self.index, self.pid))
            except OSError:
                self.result.set_exception(TaskException(errno.EFAULT, 'Cannot spawn task executor'))
                self.balancer.logger.error('Cannot spawn task executor #{0}'.format(self.index))
                return

            for line in self.proc.stdout:
                line = line.decode('utf8')
                self.balancer.logger.debug('Executor #{0}: {1}'.format(self.index, line.strip()))
                if self.task:
                    self.task.output += line

            self.proc.wait()

            if self.proc.returncode == -signal.SIGTERM:
                self.balancer.logger.info(
                    'Executor process with PID {0} was terminated gracefully'.format(
                        self.proc.pid
                    )
                )
            else:
                self.balancer.logger.error('Executor process with PID {0} died abruptly with exit code {1}'.format(
                    self.proc.pid,
                    self.proc.returncode)
                )

            self.result.set_exception(TaskException(errno.EFAULT, 'Task executor died'))
            gevent.sleep(1)

    def die(self):
        self.exiting = True
        if self.proc:
            try:
                self.proc.terminate()
            except ProcessLookupError:
                self.balancer.logger.warning('Executor process with PID {0} already dead'.format(self.proc.pid))
Exemplo n.º 9
0
class Process(object):
    # TODO: handle bot stdout and stderr
    # TODO: refactor into TTY, Process and TTYProcess?
    def __init__(self, args, env=None, executable=None, shell=False):
        master, slave = pty.openpty()
        fcntl.fcntl(master, fcntl.F_SETFL, os.O_NONBLOCK)

        self._finished = Event()
        self._master = master
        self._read_event = get_hub().loop.io(master, 1)
        self._write_event = get_hub().loop.io(master, 2)
        self._args = args
        self._proc = Popen(args,
                           env=env,
                           executable=executable,
                           shell=shell,
                           stdin=slave,
                           stdout=slave,
                           stderr=slave,
                           bufsize=0,
                           universal_newlines=False,
                           close_fds=True)

    def __repr__(self):
        return "Process:%x %r" % (id(self), self._args)

    @property
    def finished(self):
        return self._finished.ready()

    def _waitclosed(self):
        self._proc.wait()
        self.stop()

    def _writer(self, inch):
        """
        This greenlet will block until messages are ready to be written to pty
        """
        try:
            sock = self._master
            for msg in inch.watch():
                if 'resize' in msg:
                    set_winsize(sock, msg['resize']['width'],
                                msg['resize']['height'])
                if 'data' in msg:
                    buf = msg['data']
                    while not self.finished and len(buf):
                        try:
                            wait(self._write_event)
                        except Exception:
                            break
                        nwritten = os.write(sock, msg['data'])
                        buf = buf[nwritten:]
        except Exception:
            LOG.exception("In Process._writer")

    def run(self, task):
        writer_task = gevent.spawn(self._writer, task.input)
        gevent.spawn(self._waitclosed)
        proc = self._proc
        try:
            sock = self._master
            while not self.finished:
                try:
                    wait(self._read_event)
                except Exception:
                    break
                data = os.read(sock, 1024)
                if len(data) == 0 or data is StopIteration:
                    break
                if sock == proc.stderr:
                    task.output.send(dict(error=data))
                else:
                    task.output.send(dict(data=data))
        except Exception:
            LOG.exception("While reading from process")
        finally:
            writer_task.kill()
            self.stop()

    def stop(self):
        if not self.finished:
            cancel_wait(self._read_event)
            cancel_wait(self._write_event)
            try:
                os.close(self._master)
            except Exception:
                pass
            if not self._proc.poll():
                self._proc.terminate()
                self._proc.wait()
            self._finished.set()
Exemplo n.º 10
0
class TaskExecutor(object):
    def __init__(self, balancer, index):
        self.balancer = balancer
        self.index = index
        self.task = None
        self.proc = None
        self.pid = None
        self.conn = None
        self.state = None
        self.key = str(uuid.uuid4())
        self.checked_in = Event()
        self.result = AsyncResult()
        gevent.spawn(self.executor)

    def checkin(self, conn):
        self.balancer.logger.debug('Check-in of worker #{0} (key {1})'.format(self.index, self.key))
        self.conn = conn
        self.state = WorkerState.IDLE
        self.checked_in.set()

    def get_status(self):
        if not self.conn:
            return None

        try:
            st = TaskStatus(0)
            st.__setstate__(self.conn.call_client_sync('taskproxy.get_status'))
            return st
        except RpcException as err:
            self.balancer.logger.error("Cannot obtain status from task #{0}: {1}".format(self.task.id, str(err)))
            self.proc.terminate()

    def put_status(self, status):
        # Try to collect rusage at this point, when process is still alive
        try:
            kinfo = bsd.kinfo_getproc(self.pid)
            self.task.rusage = kinfo.rusage
        except LookupError:
            pass

        if status['status'] == 'FINISHED':
            self.result.set(status['result'])

        if status['status'] == 'FAILED':
            error = status['error']
            self.result.set_exception(TaskException(
                code=error['code'],
                message=error['message'],
                stacktrace=error['stacktrace'],
                extra=error.get('extra')
            ))

    def run(self, task):
        self.result = AsyncResult()
        self.task = task
        self.task.set_state(TaskState.EXECUTING)

        self.conn.call_client_sync('taskproxy.run', {
            'id': task.id,
            'class': task.clazz.__name__,
            'filename': inspect.getsourcefile(task.clazz),
            'args': task.args,
            'debugger': task.debugger
        })

        try:
            self.result.get()
        except BaseException as e:
            if not isinstance(e, TaskException):
                self.balancer.dispatcher.report_error(
                    'Task {0} raised exception other than TaskException'.format(self.task.name),
                    e
                )

            self.task.error = serialize_error(e)
            self.task.set_state(TaskState.FAILED, TaskStatus(0, str(e), extra={
                "stacktrace": traceback.format_exc()
            }))

            self.task.ended.set()
            self.balancer.task_exited(self.task)
            self.state = WorkerState.IDLE
            return

        self.task.result = self.result.value
        self.task.set_state(TaskState.FINISHED, TaskStatus(100, ''))
        self.task.ended.set()
        self.balancer.task_exited(self.task)
        self.state = WorkerState.IDLE

    def abort(self):
        self.balancer.logger.info("Trying to abort task #{0}".format(self.task.id))
        # Try to abort via RPC. If this fails, kill process
        try:
            self.conn.call_client_sync('taskproxy.abort')
        except RpcException as err:
            self.balancer.logger.warning("Failed to abort task #{0} gracefully: {1}".format(self.task.id, str(err)))
            self.balancer.logger.warning("Killing process {0}".format(self.pid))
            self.proc.terminate()

    def executor(self):
        while True:
            try:
                self.proc = Popen(
                    [TASKWORKER_PATH, self.key],
                    close_fds=True,
                    stdout=subprocess.PIPE,
                    stderr=subprocess.STDOUT)

                self.pid = self.proc.pid
                self.balancer.logger.debug('Started executor #{0} as PID {1}'.format(self.index, self.pid))
            except OSError:
                self.result.set_exception(TaskException(errno.EFAULT, 'Cannot spawn task executor'))
                return

            for line in self.proc.stdout:
                line = line.decode('utf8')
                self.balancer.logger.debug('Executor #{0}: {1}'.format(self.index, line.strip()))
                if self.task:
                    self.task.output += line

            self.proc.wait()
            self.balancer.logger.error('Executor process with PID {0} died abruptly with exit code {1}'.format(
                self.proc.pid,
                self.proc.returncode)
            )

            self.result.set_exception(TaskException(errno.EFAULT, 'Task executor died'))
            gevent.sleep(1)

    def die(self):
        if self.proc:
            self.proc.terminate()