<commit_before># (c) 2012-2014, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

from ansible.errors import AnsibleError

from ansible.parsing.mod_args import ModuleArgsParser
from ansible.parsing.splitter import parse_kv
from ansible.parsing.yaml.objects import AnsibleBaseYAMLObject, AnsibleMapping

from ansible.plugins import module_loader, lookup_loader

from ansible.playbook.attribute import Attribute, FieldAttribute
from ansible.playbook.base import Base
from ansible.playbook.become import Become
from ansible.playbook.block import Block
from ansible.playbook.conditional import Conditional
from ansible.playbook.role import Role
from ansible.playbook.taggable import Taggable

__all__ = ['Task']

class Task(Base, Conditional, Taggable, Become):

    """
    A task is a language feature that represents a call to a module, with given arguments and other parameters.
    A handler is a subclass of a task.

    Usage:

       Task.load(datastructure) -> Task
       Task.something(...)
    """

    # =================================================================================
    # ATTRIBUTES
    # load_<attribute_name> and
    # validate_<attribute_name>
    # will be used if defined
    # might be possible to define others

    _args                 = FieldAttribute(isa='dict', default=dict())
    _action               = FieldAttribute(isa='string')

    _always_run           = FieldAttribute(isa='bool')
    _any_errors_fatal     = FieldAttribute(isa='bool')
    _async                = FieldAttribute(isa='int', default=0)
    _changed_when         = FieldAttribute(isa='string')
    _delay                = FieldAttribute(isa='int', default=5)
    _delegate_to          = FieldAttribute(isa='string')
    _failed_when          = FieldAttribute(isa='string')
    _first_available_file = FieldAttribute(isa='list')
    _ignore_errors        = FieldAttribute(isa='bool')

    _loop                 = FieldAttribute(isa='string', private=True)
    _loop_args            = FieldAttribute(isa='list', private=True)
    _local_action         = FieldAttribute(isa='string')

    # FIXME: this should not be a Task
    _meta                 = FieldAttribute(isa='string')

    _name                 = FieldAttribute(isa='string', default='')

    _notify               = FieldAttribute(isa='list')
    _poll                 = FieldAttribute(isa='int')
    _register             = FieldAttribute(isa='string')
    _retries              = FieldAttribute(isa='int', default=1)
    _run_once             = FieldAttribute(isa='bool')
    _until                = FieldAttribute(isa='list') # ?

    def __init__(self, block=None, role=None, task_include=None):
        ''' constructors a task, without the Task.load classmethod, it will be pretty blank '''

        self._block        = block
        self._role         = role
        self._task_include = task_include

        super(Task, self).__init__()

    def get_name(self):
       ''' return the name of the task '''

       if self._role and self.name:
           return "%s : %s" % (self._role.get_name(), self.name)
       elif self.name:
           return self.name
       else:
           flattened_args = self._merge_kv(self.args)
           if self._role:
               return "%s : %s %s" % (self._role.get_name(), self.action, flattened_args)
           else:
               return "%s %s" % (self.action, flattened_args)

    def _merge_kv(self, ds):
        if ds is None:
            return ""
        elif isinstance(ds, basestring):
            return ds
        elif isinstance(ds, dict):
            buf = ""
            for (k,v) in ds.iteritems():
                if k.startswith('_'):
                    continue
                buf = buf + "%s=%s " % (k,v)
            buf = buf.strip()
            return buf

    @staticmethod
    def load(data, block=None, role=None, task_include=None, variable_manager=None, loader=None):
        t = Task(block=block, role=role, task_include=task_include)
        return t.load_data(data, variable_manager=variable_manager, loader=loader)

    def __repr__(self):
        ''' returns a human readable representation of the task '''
        return "TASK: %s" % self.get_name()

    def _preprocess_loop(self, ds, new_ds, k, v):
        ''' take a lookup plugin name and store it correctly '''

        loop_name = k.replace("with_", "")
        if new_ds.get('loop') is not None:
            raise AnsibleError("duplicate loop in task: %s" % loop_name, obj=ds)
        if v is None:
            raise AnsibleError("you must specify a value when using %s" % k, obj=ds)
        new_ds['loop'] = loop_name
        new_ds['loop_args'] = v

    def preprocess_data(self, ds):
        '''
        tasks are especially complex arguments so need pre-processing.
        keep it short.
        '''

        assert isinstance(ds, dict)

        # the new, cleaned datastructure, which will have legacy
        # items reduced to a standard structure suitable for the
        # attributes of the task class
        new_ds = AnsibleMapping()
        if isinstance(ds, AnsibleBaseYAMLObject):
            new_ds.ansible_pos = ds.ansible_pos

        # use the args parsing class to determine the action, args,
        # and the delegate_to value from the various possible forms
        # supported as legacy
        args_parser = ModuleArgsParser(task_ds=ds)
        (action, args, delegate_to) = args_parser.parse()

        new_ds['action']      = action
        new_ds['args']        = args
        new_ds['delegate_to'] = delegate_to

        for (k,v) in ds.iteritems():
            if k in ('action', 'local_action', 'args', 'delegate_to') or k == action or k == 'shell':
                # we don't want to re-assign these values, which were
                # determined by the ModuleArgsParser() above
                continue
            elif k.replace("with_", "") in lookup_loader:
                self._preprocess_loop(ds, new_ds, k, v)
            else:
                new_ds[k] = v

        return super(Task, self).preprocess_data(new_ds)

    def post_validate(self, templar):
        '''
        Override of base class post_validate, to also do final validation on
        the block and task include (if any) to which this task belongs.
        '''

        if self._block:
            self._block.post_validate(templar)
        if self._task_include:
            self._task_include.post_validate(templar)

        super(Task, self).post_validate(templar)

    def get_vars(self):
        all_vars = self.vars.copy()
        if self._block:
            all_vars.update(self._block.get_vars())
        if self._task_include:
            all_vars.update(self._task_include.get_vars())

        #if isinstance(self.args, dict):
        #    all_vars.update(self.args)

        if 'tags' in all_vars:
            del all_vars['tags']
        if 'when' in all_vars:
            del all_vars['when']
        return all_vars

    def copy(self, exclude_block=False):
        new_me = super(Task, self).copy()

        new_me._block = None
        if self._block and not exclude_block:
            new_me._block = self._block.copy()

        new_me._role = None
        if self._role:
            new_me._role = self._role

        new_me._task_include = None
        if self._task_include:
            new_me._task_include = self._task_include.copy()

        return new_me

    def serialize(self):
        data = super(Task, self).serialize()

        if self._block:
            data['block'] = self._block.serialize()

        if self._role:
            data['role'] = self._role.serialize()

        if self._task_include:
            data['task_include'] = self._task_include.serialize()

        return data

    def deserialize(self, data):

        # import is here to avoid import loops
        #from ansible.playbook.task_include import TaskInclude

        block_data = data.get('block')

        if block_data:
            b = Block()
            b.deserialize(block_data)
            self._block = b
            del data['block']

        role_data = data.get('role')
        if role_data:
            r = Role()
            r.deserialize(role_data)
            self._role = r
            del data['role']

        ti_data = data.get('task_include')
        if ti_data:
            #ti = TaskInclude()
            ti = Task()
            ti.deserialize(ti_data)
            self._task_include = ti
            del data['task_include']

        super(Task, self).deserialize(data)

    def evaluate_conditional(self, templar, all_vars):
        if self._block is not None:
            if not self._block.evaluate_conditional(templar, all_vars):
                return False
        if self._task_include is not None:
            if not self._task_include.evaluate_conditional(templar, all_vars):
                return False
        return super(Task, self).evaluate_conditional(templar, all_vars)

    def set_loader(self, loader):
        '''
        Sets the loader on this object and recursively on parent, child objects.
        This is used primarily after the Task has been serialized/deserialized, which
        does not preserve the loader.
        '''

        self._loader = loader

        if self._block:
            self._block.set_loader(loader)
        if self._task_include:
            self._task_include.set_loader(loader)

    def _get_parent_attribute(self, attr, extend=False):
        '''
        Generic logic to get the attribute or parent attribute for a task value.
        '''
        value = self._attributes[attr]
        if self._block and (value is None or extend):
            parent_value = getattr(self._block, attr)
            if extend:
                value = self._extend_value(value, parent_value)
            else:
                value = parent_value
        if self._task_include and (value is None or extend):
            parent_value = getattr(self._task_include, attr)
            if extend:
                value = self._extend_value(value, parent_value)
            else:
                value = parent_value
        return value

    def _get_attr_environment(self):
        '''
        Override for the 'tags' getattr fetcher, used from Base.
        '''
        environment = self._attributes['tags']
        if environment is None:
            environment = dict()

        environment = self._get_parent_attribute('environment', extend=True)

        return environment

<commit_msg>Make sure role name doesn't end up in task name more than once<commit_after># (c) 2012-2014, Michael DeHaan <michael.dehaan@gmail.com>
#
# This file is part of Ansible
#
# Ansible is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Ansible is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Ansible.  If not, see <http://www.gnu.org/licenses/>.

# Make coding more python3-ish
from __future__ import (absolute_import, division, print_function)
__metaclass__ = type

from ansible.errors import AnsibleError

from ansible.parsing.mod_args import ModuleArgsParser
from ansible.parsing.splitter import parse_kv
from ansible.parsing.yaml.objects import AnsibleBaseYAMLObject, AnsibleMapping

from ansible.plugins import module_loader, lookup_loader

from ansible.playbook.attribute import Attribute, FieldAttribute
from ansible.playbook.base import Base
from ansible.playbook.become import Become
from ansible.playbook.block import Block
from ansible.playbook.conditional import Conditional
from ansible.playbook.role import Role
from ansible.playbook.taggable import Taggable

__all__ = ['Task']

class Task(Base, Conditional, Taggable, Become):

    """
    A task is a language feature that represents a call to a module, with given arguments and other parameters.
    A handler is a subclass of a task.

    Usage:

       Task.load(datastructure) -> Task
       Task.something(...)
    """

    # =================================================================================
    # ATTRIBUTES
    # load_<attribute_name> and
    # validate_<attribute_name>
    # will be used if defined
    # might be possible to define others

    _args                 = FieldAttribute(isa='dict', default=dict())
    _action               = FieldAttribute(isa='string')

    _always_run           = FieldAttribute(isa='bool')
    _any_errors_fatal     = FieldAttribute(isa='bool')
    _async                = FieldAttribute(isa='int', default=0)
    _changed_when         = FieldAttribute(isa='string')
    _delay                = FieldAttribute(isa='int', default=5)
    _delegate_to          = FieldAttribute(isa='string')
    _failed_when          = FieldAttribute(isa='string')
    _first_available_file = FieldAttribute(isa='list')
    _ignore_errors        = FieldAttribute(isa='bool')

    _loop                 = FieldAttribute(isa='string', private=True)
    _loop_args            = FieldAttribute(isa='list', private=True)
    _local_action         = FieldAttribute(isa='string')

    # FIXME: this should not be a Task
    _meta                 = FieldAttribute(isa='string')

    _name                 = FieldAttribute(isa='string', default='')

    _notify               = FieldAttribute(isa='list')
    _poll                 = FieldAttribute(isa='int')
    _register             = FieldAttribute(isa='string')
    _retries              = FieldAttribute(isa='int', default=1)
    _run_once             = FieldAttribute(isa='bool')
    _until                = FieldAttribute(isa='list') # ?

    def __init__(self, block=None, role=None, task_include=None):
        ''' constructors a task, without the Task.load classmethod, it will be pretty blank '''

        self._block        = block
        self._role         = role
        self._task_include = task_include

        super(Task, self).__init__()

    def get_name(self):
       ''' return the name of the task '''

       if self._role and self.name and not self.name.startswith("%s :" % self._role.get_name()):
           return "%s : %s" % (self._role.get_name(), self.name)
       elif self.name:
           return self.name
       else:
           flattened_args = self._merge_kv(self.args)
           if self._role:
               return "%s : %s %s" % (self._role.get_name(), self.action, flattened_args)
           else:
               return "%s %s" % (self.action, flattened_args)

    def _merge_kv(self, ds):
        if ds is None:
            return ""
        elif isinstance(ds, basestring):
            return ds
        elif isinstance(ds, dict):
            buf = ""
            for (k,v) in ds.iteritems():
                if k.startswith('_'):
                    continue
                buf = buf + "%s=%s " % (k,v)
            buf = buf.strip()
            return buf

    @staticmethod
    def load(data, block=None, role=None, task_include=None, variable_manager=None, loader=None):
        t = Task(block=block, role=role, task_include=task_include)
        return t.load_data(data, variable_manager=variable_manager, loader=loader)

    def __repr__(self):
        ''' returns a human readable representation of the task '''
        return "TASK: %s" % self.get_name()

    def _preprocess_loop(self, ds, new_ds, k, v):
        ''' take a lookup plugin name and store it correctly '''

        loop_name = k.replace("with_", "")
        if new_ds.get('loop') is not None:
            raise AnsibleError("duplicate loop in task: %s" % loop_name, obj=ds)
        if v is None:
            raise AnsibleError("you must specify a value when using %s" % k, obj=ds)
        new_ds['loop'] = loop_name
        new_ds['loop_args'] = v

    def preprocess_data(self, ds):
        '''
        tasks are especially complex arguments so need pre-processing.
        keep it short.
        '''

        assert isinstance(ds, dict)

        # the new, cleaned datastructure, which will have legacy
        # items reduced to a standard structure suitable for the
        # attributes of the task class
        new_ds = AnsibleMapping()
        if isinstance(ds, AnsibleBaseYAMLObject):
            new_ds.ansible_pos = ds.ansible_pos

        # use the args parsing class to determine the action, args,
        # and the delegate_to value from the various possible forms
        # supported as legacy
        args_parser = ModuleArgsParser(task_ds=ds)
        (action, args, delegate_to) = args_parser.parse()

        new_ds['action']      = action
        new_ds['args']        = args
        new_ds['delegate_to'] = delegate_to

        for (k,v) in ds.iteritems():
            if k in ('action', 'local_action', 'args', 'delegate_to') or k == action or k == 'shell':
                # we don't want to re-assign these values, which were
                # determined by the ModuleArgsParser() above
                continue
            elif k.replace("with_", "") in lookup_loader:
                self._preprocess_loop(ds, new_ds, k, v)
            else:
                new_ds[k] = v

        return super(Task, self).preprocess_data(new_ds)

    def post_validate(self, templar):
        '''
        Override of base class post_validate, to also do final validation on
        the block and task include (if any) to which this task belongs.
        '''

        if self._block:
            self._block.post_validate(templar)
        if self._task_include:
            self._task_include.post_validate(templar)

        super(Task, self).post_validate(templar)

    def get_vars(self):
        all_vars = self.vars.copy()
        if self._block:
            all_vars.update(self._block.get_vars())
        if self._task_include:
            all_vars.update(self._task_include.get_vars())

        #if isinstance(self.args, dict):
        #    all_vars.update(self.args)

        if 'tags' in all_vars:
            del all_vars['tags']
        if 'when' in all_vars:
            del all_vars['when']
        return all_vars

    def copy(self, exclude_block=False):
        new_me = super(Task, self).copy()

        new_me._block = None
        if self._block and not exclude_block:
            new_me._block = self._block.copy()

        new_me._role = None
        if self._role:
            new_me._role = self._role

        new_me._task_include = None
        if self._task_include:
            new_me._task_include = self._task_include.copy()

        return new_me

    def serialize(self):
        data = super(Task, self).serialize()

        if self._block:
            data['block'] = self._block.serialize()

        if self._role:
            data['role'] = self._role.serialize()

        if self._task_include:
            data['task_include'] = self._task_include.serialize()

        return data

    def deserialize(self, data):

        # import is here to avoid import loops
        #from ansible.playbook.task_include import TaskInclude

        block_data = data.get('block')

        if block_data:
            b = Block()
            b.deserialize(block_data)
            self._block = b
            del data['block']

        role_data = data.get('role')
        if role_data:
            r = Role()
            r.deserialize(role_data)
            self._role = r
            del data['role']

        ti_data = data.get('task_include')
        if ti_data:
            #ti = TaskInclude()
            ti = Task()
            ti.deserialize(ti_data)
            self._task_include = ti
            del data['task_include']

        super(Task, self).deserialize(data)

    def evaluate_conditional(self, templar, all_vars):
        if self._block is not None:
            if not self._block.evaluate_conditional(templar, all_vars):
                return False
        if self._task_include is not None:
            if not self._task_include.evaluate_conditional(templar, all_vars):
                return False
        return super(Task, self).evaluate_conditional(templar, all_vars)

    def set_loader(self, loader):
        '''
        Sets the loader on this object and recursively on parent, child objects.
        This is used primarily after the Task has been serialized/deserialized, which
        does not preserve the loader.
        '''

        self._loader = loader

        if self._block:
            self._block.set_loader(loader)
        if self._task_include:
            self._task_include.set_loader(loader)

    def _get_parent_attribute(self, attr, extend=False):
        '''
        Generic logic to get the attribute or parent attribute for a task value.
        '''
        value = self._attributes[attr]
        if self._block and (value is None or extend):
            parent_value = getattr(self._block, attr)
            if extend:
                value = self._extend_value(value, parent_value)
            else:
                value = parent_value
        if self._task_include and (value is None or extend):
            parent_value = getattr(self._task_include, attr)
            if extend:
                value = self._extend_value(value, parent_value)
            else:
                value = parent_value
        return value

    def _get_attr_environment(self):
        '''
        Override for the 'tags' getattr fetcher, used from Base.
        '''
        environment = self._attributes['tags']
        if environment is None:
            environment = dict()

        environment = self._get_parent_attribute('environment', extend=True)

        return environment

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:

import os
import os.path as op
import time
import argparse
import yaml

from nipype import logging
logger = logging.getLogger('workflow')

class QAProtocolCLI:
    """
    This class and the associated _run_workflow function implement what
    the former scripts (qap_anatomical_spatial.py, etc.) contained
    """

    def __init__(self):
        parser = argparse.ArgumentParser()

        group = parser.add_argument_group(
            "Regular Use Inputs (non-cloud runs)")
        cloudgroup = parser.add_argument_group(
            "AWS Cloud Inputs (only required for AWS Cloud runs)")
        req = parser.add_argument_group("Required Inputs")

        cloudgroup.add_argument('--subj_idx', type=int,
                                help='Subject index to run')
        cloudgroup.add_argument(
            '--s3_dict_yml', type=str,
            help='Path to YAML file containing S3 input filepaths dictionary')

        # Subject list (YAML file)
        group.add_argument(
            "--sublist", type=str, help="filepath to subject list YAML")
        req.add_argument(
            "config", type=str, help="filepath to pipeline configuration YAML")

        # Write PDF reports
        group.add_argument(
            "--with-reports", action='store_true', default=False,
            help="Write a summary report in PDF format.")

        args = parser.parse_args()

        # checks
        if args.subj_idx and not args.s3_dict_yml and not args.sublist:
            raise RuntimeError(
                "\n[!] You provided --subj_idx, but not --s3_dict_yml. "
                "When executing cloud-based runs, please provide both "
                "inputs.\n")

        elif args.s3_dict_yml and not args.subj_idx and not args.sublist:
            raise RuntimeError(
                "\n[!] You provided --s3_dict_yml, but not --subj_idx. "
                "When executing cloud-based runs, please provide both "
                "inputs.\n")

        elif not args.sublist and not args.subj_idx and not args.s3_dict_yml:
            raise RuntimeError(
                "\n[!] Either --sublist is required for regular runs, or "
                "both --subj_idx and --s3_dict_yml for cloud-based runs.\n")

        elif args.sublist and args.subj_idx and args.s3_dict_yml:
            raise RuntimeError(
                "\n[!] Either --sublist is required for regular runs, or "
                "both --subj_idx and --s3_dict_yml for cloud-based runs, "
                "but not all three. (I'm not sure which you are trying to "
                "do!)\n")

        elif args.sublist and (args.subj_idx or args.s3_dict_yml):
            raise RuntimeError(
                "\n[!] Either --sublist is required for regular runs, or "
                "both --subj_idx and --s3_dict_yml for cloud-based runs. "
                "(I'm not sure which you are trying to do!)\n")

        self._cloudify = False

        if args.subj_idx and args.s3_dict_yml:
            self._cloudify = True

            # ---- Cloud-ify! ----
            # Import packages
            from cloud_utils import dl_subj_from_s3, upl_qap_output
            # Download and build a one-subject dictionary from S3
            self._sub_dict = dl_subj_from_s3(
                args.subj_idx, args.config, args.s3_dict_yml)

            if not self._sub_dict:
                err = "\n[!] Subject dictionary was not successfully " \
                      "downloaded from the S3 bucket!\n"
                raise RuntimeError(err)

        elif args.sublist:
            self._sub_dict = args.sublist

        else:
            raise RuntimeError(
                "\n[!] Arguments were parsed, but no appropriate run found")

        # Load config
        with open(args.config, "r") as f:
            self._config = yaml.load(f)

        self._config['pipeline_config_yaml'] = args.config
        self._config['qap_type'] = parser.prog[4:-3]

        if args.with_reports:
            self._config['write_report'] = True

    def _run_here(self, run_name):
        ns_at_once = self._config.get('num_subjects_at_once', 1)
        with open(self._sub_dict, "r") as f:
            subdict = yaml.load(f)

        flat_sub_dict = {}
        sites_dict = {}

        # Preamble: generate flat_subdict
        for subid in subdict.keys():
            # sessions
            for session in subdict[subid].keys():
                # resource files
                for resource in subdict[subid][session].keys():
                    if type(subdict[subid][session][resource]) is dict:
                        # then this has sub-scans defined
                        for scan in subdict[subid][session][resource].keys():
                            filepath = subdict[subid][session][resource][scan]
                            resource_dict = {}
                            resource_dict[resource] = filepath
                            sub_info_tuple = (subid, session, scan)
                            if sub_info_tuple not in flat_sub_dict.keys():
                                flat_sub_dict[sub_info_tuple] = {}

                            flat_sub_dict[sub_info_tuple].update(resource_dict)

                    elif resource == "site_name":
                        sites_dict[subid] = subdict[subid][session][resource]

                    else:
                        filepath = subdict[subid][session][resource]
                        resource_dict = {}
                        resource_dict[resource] = filepath
                        sub_info_tuple = (subid, session, None)

                        if sub_info_tuple not in flat_sub_dict.keys():
                            flat_sub_dict[sub_info_tuple] = {}

                        flat_sub_dict[sub_info_tuple].update(resource_dict)

            # in case some subjects have site names and others don't
            if len(sites_dict.keys()) > 0:
                for subid in subdict.keys():
                    if subid not in sites_dict.keys():
                        sites_dict[subid] = None

        # Start the magic
        logger.info('There are %d subjects in the pool' %
                    len(flat_sub_dict.keys()))

        # Stack workflow args
        wfargs = [(flat_sub_dict[sub_info], self._config, sub_info,
                  run_name, sites_dict.get(sub_info[0], None))
                  for sub_info in flat_sub_dict.keys()]

        results = []

        # skip parallel machinery if we are running only one subject at once
        if ns_at_once == 1:
            for a in wfargs:
                results.append(_run_workflow(a))
        else:
            from multiprocessing import Pool

            try:
                pool = Pool(processes=ns_at_once, masktasksperchild=50)
            except TypeError:  # Make python <2.7 compatible
                pool = Pool(processes=ns_at_once)

            results = pool.map(_run_workflow, wfargs)
            pool.close()
            pool.terminate()
        return results

    def _run_cloud(self, run_name):
        from cloud_utils import upl_qap_output
        # get the site name!
        for resource_path in subject_list[sub]:
            if ".nii" in resource_path:
                filepath = resource_path
                break

        filesplit = filepath.split(self._config["bucket_prefix"])
        site_name = filesplit[1].split("/")[1]

        rt = _run_workflow(
            subject_list[sub], self._config, sub, run_name, site_name)

        # upload results
        upl_qap_output(self._config)
        return rt

    def run(self):
        # Get configurations and settings
        config = self._config
        subject_list = self._sub_dict
        cloudify = self._cloudify
        ns_at_once = config.get('num_subjects_at_once', 1)
        write_report = config.get('write_report', False)

        # Create output directory
        try:
            os.makedirs(config["output_directory"])
        except:
            if not op.isdir(config["output_directory"]):
                err = "[!] Output directory unable to be created.\n" \
                      "Path: %s\n\n" % config["output_directory"]
                raise Exception(err)
            else:
                pass

        # Create working directory
        try:
            os.makedirs(config["working_directory"])
        except:
            if not op.isdir(config["working_directory"]):
                err = "[!] Output directory unable to be created.\n" \
                      "Path: %s\n\n" % config["working_directory"]
                raise Exception(err)
            else:
                pass

        run_name = config['pipeline_config_yaml'].split("/")[-1].split(".")[0]

        results = None

        try:
            if not cloudify:
                results = self._run_here(run_name)
            else:
                results = self._run_cloud(run_name)
        except Exception as e:
            raise Exception("\n\n%s\n\n%s\n\n" % (e, globals()))

        # PDF reporting
        if write_report:
            from qap.viz.reports import workflow_report
            logger.info('Writing PDF reports')
            qap_type = 'qap_' + config['qap_type']
            in_csv = op.join(config['output_directory'], '%s.csv' % qap_type)

            reports = workflow_report(in_csv, qap_type, run_name, results,
                                      out_dir=config['output_directory'])

            for k, v in reports.iteritems():
                if v['success']:
                    logger.info('Written report (%s) in %s' % (k, v['path']))

def _run_workflow(args):

    # build pipeline for each subject, individually
    # ~ 5 min 20 sec per subject
    # (roughly 320 seconds)

    import os
    import os.path as op
    import sys

    import nipype.interfaces.io as nio
    import nipype.pipeline.engine as pe

    import nipype.interfaces.utility as util
    import nipype.interfaces.fsl.maths as fsl

    import glob

    import time
    from time import strftime
    from nipype import config as nyconfig

    resource_pool, config, subject_info, run_name, site_name = args
    sub_id = str(subject_info[0])

    qap_type = config['qap_type']

    if subject_info[1]:
        session_id = subject_info[1]
    else:
        session_id = "session_0"

    if subject_info[2]:
        scan_id = subject_info[2]
    else:
        scan_id = "scan_0"

    # Read and apply general settings in config
    keep_outputs = config.get('write_all_outputs', False)
    output_dir = op.join(config["output_directory"], run_name,
                         sub_id, session_id, scan_id)

    try:
        os.makedirs(output_dir)
    except:
        if not op.isdir(output_dir):
            err = "[!] Output directory unable to be created.\n" \
                  "Path: %s\n\n" % output_dir
            raise Exception(err)
        else:
            pass

    log_dir = output_dir

    # set up logging
    nyconfig.update_config(
        {'logging': {'log_directory': log_dir, 'log_to_file': True}})
    logging.update_logging(nyconfig)

    # take date+time stamp for run identification purposes
    unique_pipeline_id = strftime("%Y%m%d%H%M%S")
    pipeline_start_stamp = strftime("%Y-%m-%d_%H:%M:%S")

    pipeline_start_time = time.time()

    logger.info("Pipeline start time: %s" % pipeline_start_stamp)
    logger.info("Contents of resource pool:\n" + str(resource_pool))
    logger.info("Configuration settings:\n" + str(config))

    # for QAP spreadsheet generation only
    config.update({"subject_id": sub_id, "session_id": session_id,
                   "scan_id": scan_id, "run_name": run_name})

    if site_name:
        config["site_name"] = site_name

    workflow = pe.Workflow(name=scan_id)
    workflow.base_dir = op.join(config["working_directory"], sub_id,
                                session_id)

    # set up crash directory
    workflow.config['execution'] = \
        {'crashdump_dir': config["output_directory"]}

    # update that resource pool with what's already in the output directory
    for resource in os.listdir(output_dir):
        if (op.isdir(op.join(output_dir, resource)) and
                resource not in resource_pool.keys()):
            resource_pool[resource] = glob.glob(op.join(output_dir,
                                                        resource, "*"))[0]

    # resource pool check
    invalid_paths = []

    for resource in resource_pool.keys():
        if not op.isfile(resource_pool[resource]):
            invalid_paths.append((resource, resource_pool[resource]))

    if len(invalid_paths) > 0:
        err = "\n\n[!] The paths provided in the subject list to the " \
              "following resources are not valid:\n"

        for path_tuple in invalid_paths:
            err = err + path_tuple[0] + ": " + path_tuple[1] + "\n"

        err = err + "\n\n"
        raise Exception(err)

    # start connecting the pipeline
    if 'qap_' + qap_type not in resource_pool.keys():
        from qap import qap_workflows as qw
        wf_builder = getattr(qw, 'qap_' + qap_type + '_workflow')
        workflow, resource_pool = wf_builder(workflow, resource_pool, config)

    # set up the datasinks
    new_outputs = 0

    out_list = ['qap_' + qap_type]
    if keep_outputs:
        out_list = resource_pool.keys()

    # Save reports to out_dir if necessary
    if config.get('write_report', False):
        out_list += ['qap_mosaic']

        # The functional temporal also has an FD plot
        if 'functional_temporal' in qap_type:
            out_list += ['qap_fd']

    for output in out_list:
        # we use a check for len()==2 here to select those items in the
        # resource pool which are tuples of (node, node_output), instead
        # of the items which are straight paths to files

        # resource pool items which are in the tuple format are the
        # outputs that have been created in this workflow because they
        # were not present in the subject list YML (the starting resource
        # pool) and had to be generated
        if len(resource_pool[output]) == 2:
            ds = pe.Node(nio.DataSink(), name='datasink_%s' % output)
            ds.inputs.base_directory = output_dir
            node, out_file = resource_pool[output]
            workflow.connect(node, out_file, ds, output)
            new_outputs += 1

    rt = {'id': sub_id, 'session': session_id, 'scan': scan_id,
          'status': 'started'}
    # run the pipeline (if there is anything to do)
    if new_outputs > 0:
        if config.get('write_graph', False):
            workflow.write_graph(
                dotfilename=op.join(output_dir, run_name + ".dot"),
                simple_form=False)

        nc_per_subject = config.get('num_cores_per_subject', 1)
        runargs = {'plugin': 'Linear', 'plugin_args': {}}
        if nc_per_subject > 1:
            runargs['plugin'] = 'MultiProc',
            runargs['plugin_args'] = {'n_procs': nc_per_subject}

        try:
            workflow.run(**runargs)
            rt['status'] = 'finished'
        except Exception as e:  # TODO We should be more specific here ...
            rt.update({'status': 'failed', 'msg': e})
            # ... however this is run inside a pool.map: do not raise Execption

    else:
        rt['status'] = 'cached'
        logger.info("\nEverything is already done for subject %s." % sub_id)

    # Remove working directory when done
    if not keep_outputs:
        try:
            work_dir = op.join(workflow.base_dir, scan_id)

            if op.exists(work_dir):
                import shutil
                shutil.rmtree(work_dir)
        except:
            logger.warn("Couldn\'t remove the working directory!")
            pass

    pipeline_end_stamp = strftime("%Y-%m-%d_%H:%M:%S")
    pipeline_end_time = time.time()
    logger.info("Elapsed time (minutes) since last start: %s"
                % ((pipeline_end_time - pipeline_start_time) / 60))
    logger.info("Pipeline end time: %s" % pipeline_end_stamp)
    return rt
<commit_msg>Fix attempt.<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:

import os
import os.path as op
import time
import argparse
import yaml

from nipype import logging
logger = logging.getLogger('workflow')

class QAProtocolCLI:
    """
    This class and the associated _run_workflow function implement what
    the former scripts (qap_anatomical_spatial.py, etc.) contained
    """

    def __init__(self):
        parser = argparse.ArgumentParser()

        group = parser.add_argument_group(
            "Regular Use Inputs (non-cloud runs)")
        cloudgroup = parser.add_argument_group(
            "AWS Cloud Inputs (only required for AWS Cloud runs)")
        req = parser.add_argument_group("Required Inputs")

        cloudgroup.add_argument('--subj_idx', type=int,
                                help='Subject index to run')
        cloudgroup.add_argument(
            '--s3_dict_yml', type=str,
            help='Path to YAML file containing S3 input filepaths dictionary')

        # Subject list (YAML file)
        group.add_argument(
            "--sublist", type=str, help="filepath to subject list YAML")
        req.add_argument(
            "config", type=str, help="filepath to pipeline configuration YAML")

        # Write PDF reports
        group.add_argument(
            "--with-reports", action='store_true', default=False,
            help="Write a summary report in PDF format.")

        args = parser.parse_args()

        # checks
        if args.subj_idx and not args.s3_dict_yml and not args.sublist:
            raise RuntimeError(
                "\n[!] You provided --subj_idx, but not --s3_dict_yml. "
                "When executing cloud-based runs, please provide both "
                "inputs.\n")

        elif args.s3_dict_yml and not args.subj_idx and not args.sublist:
            raise RuntimeError(
                "\n[!] You provided --s3_dict_yml, but not --subj_idx. "
                "When executing cloud-based runs, please provide both "
                "inputs.\n")

        elif not args.sublist and not args.subj_idx and not args.s3_dict_yml:
            raise RuntimeError(
                "\n[!] Either --sublist is required for regular runs, or "
                "both --subj_idx and --s3_dict_yml for cloud-based runs.\n")

        elif args.sublist and args.subj_idx and args.s3_dict_yml:
            raise RuntimeError(
                "\n[!] Either --sublist is required for regular runs, or "
                "both --subj_idx and --s3_dict_yml for cloud-based runs, "
                "but not all three. (I'm not sure which you are trying to "
                "do!)\n")

        elif args.sublist and (args.subj_idx or args.s3_dict_yml):
            raise RuntimeError(
                "\n[!] Either --sublist is required for regular runs, or "
                "both --subj_idx and --s3_dict_yml for cloud-based runs. "
                "(I'm not sure which you are trying to do!)\n")

        self._cloudify = False

        if args.subj_idx and args.s3_dict_yml:
            self._cloudify = True

            # ---- Cloud-ify! ----
            # Import packages
            from cloud_utils import dl_subj_from_s3, upl_qap_output
            # Download and build a one-subject dictionary from S3
            self._sub_dict = dl_subj_from_s3(
                args.subj_idx, args.config, args.s3_dict_yml)

            if not self._sub_dict:
                err = "\n[!] Subject dictionary was not successfully " \
                      "downloaded from the S3 bucket!\n"
                raise RuntimeError(err)

        elif args.sublist:
            self._sub_dict = args.sublist

        else:
            raise RuntimeError(
                "\n[!] Arguments were parsed, but no appropriate run found")

        # Load config
        with open(args.config, "r") as f:
            self._config = yaml.load(f)

        self._config['pipeline_config_yaml'] = args.config
        self._config['qap_type'] = parser.prog[4:-3]

        if args.with_reports:
            self._config['write_report'] = True

    def _run_here(self, run_name):
        ns_at_once = self._config.get('num_subjects_at_once', 1)
        with open(self._sub_dict, "r") as f:
            subdict = yaml.load(f)

        flat_sub_dict = {}
        sites_dict = {}

        # Preamble: generate flat_subdict
        for subid in subdict.keys():
            # sessions
            for session in subdict[subid].keys():
                # resource files
                for resource in subdict[subid][session].keys():
                    if type(subdict[subid][session][resource]) is dict:
                        # then this has sub-scans defined
                        for scan in subdict[subid][session][resource].keys():
                            filepath = subdict[subid][session][resource][scan]
                            resource_dict = {}
                            resource_dict[resource] = filepath
                            sub_info_tuple = (subid, session, scan)
                            if sub_info_tuple not in flat_sub_dict.keys():
                                flat_sub_dict[sub_info_tuple] = {}

                            flat_sub_dict[sub_info_tuple].update(resource_dict)

                    elif resource == "site_name":
                        sites_dict[subid] = subdict[subid][session][resource]

                    else:
                        filepath = subdict[subid][session][resource]
                        resource_dict = {}
                        resource_dict[resource] = filepath
                        sub_info_tuple = (subid, session, None)

                        if sub_info_tuple not in flat_sub_dict.keys():
                            flat_sub_dict[sub_info_tuple] = {}

                        flat_sub_dict[sub_info_tuple].update(resource_dict)

            # in case some subjects have site names and others don't
            if len(sites_dict.keys()) > 0:
                for subid in subdict.keys():
                    if subid not in sites_dict.keys():
                        sites_dict[subid] = None

        # Start the magic
        logger.info('There are %d subjects in the pool' %
                    len(flat_sub_dict.keys()))

        # Stack workflow args
        wfargs = [(flat_sub_dict[sub_info], self._config, sub_info,
                  run_name, sites_dict.get(sub_info[0], None))
                  for sub_info in flat_sub_dict.keys()]

        results = []

        # skip parallel machinery if we are running only one subject at once
        if ns_at_once == 1:
            for a in wfargs:
                results.append(_run_workflow(a))
        else:
            from multiprocessing import Pool

            try:
                pool = Pool(processes=ns_at_once, masktasksperchild=50)
            except TypeError:  # Make python <2.7 compatible
                pool = Pool(processes=ns_at_once)

            results = pool.map(_run_workflow, wfargs)
            pool.close()
            pool.terminate()
        return results

    def _run_cloud(self, run_name, subject_list):
        from cloud_utils import upl_qap_output
        # get the site name!

        for sub in subject_list.keys():

            for resource_path in subject_list[sub]:
                if ".nii" in resource_path:
                    filepath = resource_path
                    break

            filesplit = filepath.split(self._config["bucket_prefix"])
            site_name = filesplit[1].split("/")[1]

            rt = _run_workflow(
                subject_list[sub], self._config, sub, run_name, site_name)

        # upload results
        upl_qap_output(self._config)

        return rt

    def run(self):
        # Get configurations and settings
        config = self._config
        subject_list = self._sub_dict
        cloudify = self._cloudify
        ns_at_once = config.get('num_subjects_at_once', 1)
        write_report = config.get('write_report', False)

        # Create output directory
        try:
            os.makedirs(config["output_directory"])
        except:
            if not op.isdir(config["output_directory"]):
                err = "[!] Output directory unable to be created.\n" \
                      "Path: %s\n\n" % config["output_directory"]
                raise Exception(err)
            else:
                pass

        # Create working directory
        try:
            os.makedirs(config["working_directory"])
        except:
            if not op.isdir(config["working_directory"]):
                err = "[!] Output directory unable to be created.\n" \
                      "Path: %s\n\n" % config["working_directory"]
                raise Exception(err)
            else:
                pass

        run_name = config['pipeline_config_yaml'].split("/")[-1].split(".")[0]

        results = None

        try:
            if not cloudify:
                results = self._run_here(run_name)
            else:
                results = self._run_cloud(run_name, subject_list)
        except Exception as e:
            raise Exception("\n\n%s\n\n%s\n\n" % (e, locals()))

        # PDF reporting
        if write_report:
            from qap.viz.reports import workflow_report
            logger.info('Writing PDF reports')
            qap_type = 'qap_' + config['qap_type']
            in_csv = op.join(config['output_directory'], '%s.csv' % qap_type)

            reports = workflow_report(in_csv, qap_type, run_name, results,
                                      out_dir=config['output_directory'])

            for k, v in reports.iteritems():
                if v['success']:
                    logger.info('Written report (%s) in %s' % (k, v['path']))

def _run_workflow(args):

    # build pipeline for each subject, individually
    # ~ 5 min 20 sec per subject
    # (roughly 320 seconds)

    import os
    import os.path as op
    import sys

    import nipype.interfaces.io as nio
    import nipype.pipeline.engine as pe

    import nipype.interfaces.utility as util
    import nipype.interfaces.fsl.maths as fsl

    import glob

    import time
    from time import strftime
    from nipype import config as nyconfig

    resource_pool, config, subject_info, run_name, site_name = args
    sub_id = str(subject_info[0])

    qap_type = config['qap_type']

    if subject_info[1]:
        session_id = subject_info[1]
    else:
        session_id = "session_0"

    if subject_info[2]:
        scan_id = subject_info[2]
    else:
        scan_id = "scan_0"

    # Read and apply general settings in config
    keep_outputs = config.get('write_all_outputs', False)
    output_dir = op.join(config["output_directory"], run_name,
                         sub_id, session_id, scan_id)

    try:
        os.makedirs(output_dir)
    except:
        if not op.isdir(output_dir):
            err = "[!] Output directory unable to be created.\n" \
                  "Path: %s\n\n" % output_dir
            raise Exception(err)
        else:
            pass

    log_dir = output_dir

    # set up logging
    nyconfig.update_config(
        {'logging': {'log_directory': log_dir, 'log_to_file': True}})
    logging.update_logging(nyconfig)

    # take date+time stamp for run identification purposes
    unique_pipeline_id = strftime("%Y%m%d%H%M%S")
    pipeline_start_stamp = strftime("%Y-%m-%d_%H:%M:%S")

    pipeline_start_time = time.time()

    logger.info("Pipeline start time: %s" % pipeline_start_stamp)
    logger.info("Contents of resource pool:\n" + str(resource_pool))
    logger.info("Configuration settings:\n" + str(config))

    # for QAP spreadsheet generation only
    config.update({"subject_id": sub_id, "session_id": session_id,
                   "scan_id": scan_id, "run_name": run_name})

    if site_name:
        config["site_name"] = site_name

    workflow = pe.Workflow(name=scan_id)
    workflow.base_dir = op.join(config["working_directory"], sub_id,
                                session_id)

    # set up crash directory
    workflow.config['execution'] = \
        {'crashdump_dir': config["output_directory"]}

    # update that resource pool with what's already in the output directory
    for resource in os.listdir(output_dir):
        if (op.isdir(op.join(output_dir, resource)) and
                resource not in resource_pool.keys()):
            resource_pool[resource] = glob.glob(op.join(output_dir,
                                                        resource, "*"))[0]

    # resource pool check
    invalid_paths = []

    for resource in resource_pool.keys():
        if not op.isfile(resource_pool[resource]):
            invalid_paths.append((resource, resource_pool[resource]))

    if len(invalid_paths) > 0:
        err = "\n\n[!] The paths provided in the subject list to the " \
              "following resources are not valid:\n"

        for path_tuple in invalid_paths:
            err = err + path_tuple[0] + ": " + path_tuple[1] + "\n"

        err = err + "\n\n"
        raise Exception(err)

    # start connecting the pipeline
    if 'qap_' + qap_type not in resource_pool.keys():
        from qap import qap_workflows as qw
        wf_builder = getattr(qw, 'qap_' + qap_type + '_workflow')
        workflow, resource_pool = wf_builder(workflow, resource_pool, config)

    # set up the datasinks
    new_outputs = 0

    out_list = ['qap_' + qap_type]
    if keep_outputs:
        out_list = resource_pool.keys()

    # Save reports to out_dir if necessary
    if config.get('write_report', False):
        out_list += ['qap_mosaic']

        # The functional temporal also has an FD plot
        if 'functional_temporal' in qap_type:
            out_list += ['qap_fd']

    for output in out_list:
        # we use a check for len()==2 here to select those items in the
        # resource pool which are tuples of (node, node_output), instead
        # of the items which are straight paths to files

        # resource pool items which are in the tuple format are the
        # outputs that have been created in this workflow because they
        # were not present in the subject list YML (the starting resource
        # pool) and had to be generated
        if len(resource_pool[output]) == 2:
            ds = pe.Node(nio.DataSink(), name='datasink_%s' % output)
            ds.inputs.base_directory = output_dir
            node, out_file = resource_pool[output]
            workflow.connect(node, out_file, ds, output)
            new_outputs += 1

    rt = {'id': sub_id, 'session': session_id, 'scan': scan_id,
          'status': 'started'}
    # run the pipeline (if there is anything to do)
    if new_outputs > 0:
        if config.get('write_graph', False):
            workflow.write_graph(
                dotfilename=op.join(output_dir, run_name + ".dot"),
                simple_form=False)

        nc_per_subject = config.get('num_cores_per_subject', 1)
        runargs = {'plugin': 'Linear', 'plugin_args': {}}
        if nc_per_subject > 1:
            runargs['plugin'] = 'MultiProc',
            runargs['plugin_args'] = {'n_procs': nc_per_subject}

        try:
            workflow.run(**runargs)
            rt['status'] = 'finished'
        except Exception as e:  # TODO We should be more specific here ...
            rt.update({'status': 'failed', 'msg': e})
            # ... however this is run inside a pool.map: do not raise Execption

    else:
        rt['status'] = 'cached'
        logger.info("\nEverything is already done for subject %s." % sub_id)

    # Remove working directory when done
    if not keep_outputs:
        try:
            work_dir = op.join(workflow.base_dir, scan_id)

            if op.exists(work_dir):
                import shutil
                shutil.rmtree(work_dir)
        except:
            logger.warn("Couldn\'t remove the working directory!")
            pass

    pipeline_end_stamp = strftime("%Y-%m-%d_%H:%M:%S")
    pipeline_end_time = time.time()
    logger.info("Elapsed time (minutes) since last start: %s"
                % ((pipeline_end_time - pipeline_start_time) / 60))
    logger.info("Pipeline end time: %s" % pipeline_end_stamp)
    return rt
<|endoftext|>


<commit_before># Copyright (C) 2010-2013 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import hashlib
import binascii
import logging
from datetime import datetime

try:
    import magic
except ImportError:
    pass

try:
    import pydeep
    HAVE_SSDEEP = True
except ImportError:
    HAVE_SSDEEP = False

try:
    import yara
    HAVE_YARA = True
except ImportError:
    HAVE_YARA = False

from lib.cuckoo.common.utils import convert_to_printable
from lib.cuckoo.common.constants import CUCKOO_ROOT

log = logging.getLogger(__name__)

FILE_CHUNK_SIZE = 16 * 1024

class Dictionary(dict):
    """Cuckoo custom dict."""

    def __getattr__(self, key):
        return self.get(key, None)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

class URL:
    """URL base object."""

    def __init__(self, url):
        """@param url: URL"""
        self.url = url

class File:
    """Basic file object class with all useful utilities."""

    def __init__(self, file_path):
        """@param file_path: file path."""
        self.file_path = file_path

        # these will be populated when first accessed
        self._file_data = None
        self._crc32     = None
        self._md5       = None
        self._sha1      = None
        self._sha256    = None
        self._sha512    = None

    def get_name(self):
        """Get file name.
        @return: file name.
        """
        file_name = os.path.basename(self.file_path)
        return file_name

    def valid(self):
        return os.path.exists(self.file_path) and \
            os.path.isfile(self.file_path) and \
            os.path.getsize(self.file_path) != 0

    def get_data(self):
        """Read file contents.
        @return: data.
        """
        return self.file_data

    def get_chunks(self):
        """Read file contents in chunks (generator)."""

        fd = open(self.file_path, "rb")
        while True:
            chunk = fd.read(FILE_CHUNK_SIZE)
            if not chunk: break
            yield chunk
        fd.close()

    def calc_hashes(self):
        """Calculate all possible hashes for this file."""
        crc     = 0
        md5     = hashlib.md5()
        sha1    = hashlib.sha1()
        sha256  = hashlib.sha256()
        sha512  = hashlib.sha512()

        for chunk in self.get_chunks():
            crc = binascii.crc32(chunk, crc)
            md5.update(chunk)
            sha1.update(chunk)
            sha256.update(chunk)
            sha512.update(chunk)

        self._crc32     = "".join("%02X" % ((crc>>i)&0xff) for i in [24, 16, 8, 0])
        self._md5       = md5.hexdigest()
        self._sha1      = sha1.hexdigest()
        self._sha256    = sha256.hexdigest()
        self._sha512    = sha512.hexdigest()

    @property
    def file_data(self):
        log.warning("Usage of File.file_data is deprecated, use chunks of files only.")
        if not self._file_data: self._file_data = open(self.file_path, "rb").read()
        return self._file_data

    def get_size(self):
        """Get file size.
        @return: file size.
        """
        return os.path.getsize(self.file_path)

    def get_crc32(self):
        """Get CRC32.
        @return: CRC32.
        """
        if not self._crc32: self.calc_hashes()
        return self._crc32

    def get_md5(self):
        """Get MD5.
        @return: MD5.
        """
        if not self._md5: self.calc_hashes()
        return self._md5

    def get_sha1(self):
        """Get SHA1.
        @return: SHA1.
        """
        if not self._sha1: self.calc_hashes()
        return self._sha1

    def get_sha256(self):
        """Get SHA256.
        @return: SHA256.
        """
        if not self._sha256: self.calc_hashes()
        return self._sha256

    def get_sha512(self):
        """
        Get SHA512.
        @return: SHA512.
        """
        if not self._sha512: self.calc_hashes()
        return self._sha512

    def get_ssdeep(self):
        """Get SSDEEP.
        @return: SSDEEP.
        """
        if not HAVE_SSDEEP:
            return None

        try:
            return pydeep.hash_file(self.file_path)
        except Exception:
            return None

    def get_type(self):
        """Get MIME file type.
        @return: file type.
        """
        try:
            ms = magic.open(magic.MAGIC_NONE)
            ms.load()
            file_type = ms.file(self.file_path)
        except:
            try:
                file_type = magic.from_file(self.file_path)
            except:
                try:
                    import subprocess
                    file_process = subprocess.Popen(["file",
                                                     "-b",
                                                     self.file_path],
                                                    stdout = subprocess.PIPE)
                    file_type = file_process.stdout.read().strip()
                except:
                    return ""
        finally:
            try:
                ms.close()
            except:
                pass

        return file_type

    def get_yara(self, rulepath=os.path.join(CUCKOO_ROOT, "data", "yara", "index_binary.yar")):
        """Get Yara signatures matches.
        @return: matched Yara signatures.
        """
        matches = []

        if HAVE_YARA:
            try:
                rules = yara.compile(rulepath)

                for match in rules.match(self.file_path):
                    strings = []
                    for s in match.strings:
                        # Beware, spaghetti code ahead.
                        try:
                            new = s[2].encode("utf-8")
                        except UnicodeDecodeError:
                            s = s[2].lstrip("uU").encode("hex").upper()
                            s = " ".join(s[i:i+2] for i in range(0, len(s), 2))
                            new = "{ %s }" % s

                        if new not in strings:
                            strings.append(new)

                    matches.append({"name" : match.rule,
                                    "meta" : match.meta,
                                    "strings" : strings})
            except yara.Error as e:
                log.warning("Unable to match Yara signatures: %s", e)

        return matches

    def get_all(self):
        """Get all information available.
        @return: information dict.
        """
        infos = {}
        infos["name"] = self.get_name()
        infos["path"] = self.file_path
        infos["size"] = self.get_size()
        infos["crc32"] = self.get_crc32()
        infos["md5"] = self.get_md5()
        infos["sha1"] = self.get_sha1()
        infos["sha256"] = self.get_sha256()
        infos["sha512"] = self.get_sha512()
        infos["ssdeep"] = self.get_ssdeep()
        infos["type"] = self.get_type()
        infos["yara"] = self.get_yara()

        return infos
<commit_msg>Removed not used imports<commit_after># Copyright (C) 2010-2013 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import hashlib
import binascii
import logging

try:
    import magic
except ImportError:
    pass

try:
    import pydeep
    HAVE_SSDEEP = True
except ImportError:
    HAVE_SSDEEP = False

try:
    import yara
    HAVE_YARA = True
except ImportError:
    HAVE_YARA = False

from lib.cuckoo.common.constants import CUCKOO_ROOT

log = logging.getLogger(__name__)

FILE_CHUNK_SIZE = 16 * 1024

class Dictionary(dict):
    """Cuckoo custom dict."""

    def __getattr__(self, key):
        return self.get(key, None)

    __setattr__ = dict.__setitem__
    __delattr__ = dict.__delitem__

class URL:
    """URL base object."""

    def __init__(self, url):
        """@param url: URL"""
        self.url = url

class File:
    """Basic file object class with all useful utilities."""

    def __init__(self, file_path):
        """@param file_path: file path."""
        self.file_path = file_path

        # these will be populated when first accessed
        self._file_data = None
        self._crc32     = None
        self._md5       = None
        self._sha1      = None
        self._sha256    = None
        self._sha512    = None

    def get_name(self):
        """Get file name.
        @return: file name.
        """
        file_name = os.path.basename(self.file_path)
        return file_name

    def valid(self):
        return os.path.exists(self.file_path) and \
            os.path.isfile(self.file_path) and \
            os.path.getsize(self.file_path) != 0

    def get_data(self):
        """Read file contents.
        @return: data.
        """
        return self.file_data

    def get_chunks(self):
        """Read file contents in chunks (generator)."""

        fd = open(self.file_path, "rb")
        while True:
            chunk = fd.read(FILE_CHUNK_SIZE)
            if not chunk: break
            yield chunk
        fd.close()

    def calc_hashes(self):
        """Calculate all possible hashes for this file."""
        crc     = 0
        md5     = hashlib.md5()
        sha1    = hashlib.sha1()
        sha256  = hashlib.sha256()
        sha512  = hashlib.sha512()

        for chunk in self.get_chunks():
            crc = binascii.crc32(chunk, crc)
            md5.update(chunk)
            sha1.update(chunk)
            sha256.update(chunk)
            sha512.update(chunk)

        self._crc32     = "".join("%02X" % ((crc>>i)&0xff) for i in [24, 16, 8, 0])
        self._md5       = md5.hexdigest()
        self._sha1      = sha1.hexdigest()
        self._sha256    = sha256.hexdigest()
        self._sha512    = sha512.hexdigest()

    @property
    def file_data(self):
        log.warning("Usage of File.file_data is deprecated, use chunks of files only.")
        if not self._file_data: self._file_data = open(self.file_path, "rb").read()
        return self._file_data

    def get_size(self):
        """Get file size.
        @return: file size.
        """
        return os.path.getsize(self.file_path)

    def get_crc32(self):
        """Get CRC32.
        @return: CRC32.
        """
        if not self._crc32: self.calc_hashes()
        return self._crc32

    def get_md5(self):
        """Get MD5.
        @return: MD5.
        """
        if not self._md5: self.calc_hashes()
        return self._md5

    def get_sha1(self):
        """Get SHA1.
        @return: SHA1.
        """
        if not self._sha1: self.calc_hashes()
        return self._sha1

    def get_sha256(self):
        """Get SHA256.
        @return: SHA256.
        """
        if not self._sha256: self.calc_hashes()
        return self._sha256

    def get_sha512(self):
        """
        Get SHA512.
        @return: SHA512.
        """
        if not self._sha512: self.calc_hashes()
        return self._sha512

    def get_ssdeep(self):
        """Get SSDEEP.
        @return: SSDEEP.
        """
        if not HAVE_SSDEEP:
            return None

        try:
            return pydeep.hash_file(self.file_path)
        except Exception:
            return None

    def get_type(self):
        """Get MIME file type.
        @return: file type.
        """
        try:
            ms = magic.open(magic.MAGIC_NONE)
            ms.load()
            file_type = ms.file(self.file_path)
        except:
            try:
                file_type = magic.from_file(self.file_path)
            except:
                try:
                    import subprocess
                    file_process = subprocess.Popen(["file",
                                                     "-b",
                                                     self.file_path],
                                                    stdout = subprocess.PIPE)
                    file_type = file_process.stdout.read().strip()
                except:
                    return ""
        finally:
            try:
                ms.close()
            except:
                pass

        return file_type

    def get_yara(self, rulepath=os.path.join(CUCKOO_ROOT, "data", "yara", "index_binary.yar")):
        """Get Yara signatures matches.
        @return: matched Yara signatures.
        """
        matches = []

        if HAVE_YARA:
            try:
                rules = yara.compile(rulepath)

                for match in rules.match(self.file_path):
                    strings = []
                    for s in match.strings:
                        # Beware, spaghetti code ahead.
                        try:
                            new = s[2].encode("utf-8")
                        except UnicodeDecodeError:
                            s = s[2].lstrip("uU").encode("hex").upper()
                            s = " ".join(s[i:i+2] for i in range(0, len(s), 2))
                            new = "{ %s }" % s

                        if new not in strings:
                            strings.append(new)

                    matches.append({"name" : match.rule,
                                    "meta" : match.meta,
                                    "strings" : strings})
            except yara.Error as e:
                log.warning("Unable to match Yara signatures: %s", e)

        return matches

    def get_all(self):
        """Get all information available.
        @return: information dict.
        """
        infos = {}
        infos["name"] = self.get_name()
        infos["path"] = self.file_path
        infos["size"] = self.get_size()
        infos["crc32"] = self.get_crc32()
        infos["md5"] = self.get_md5()
        infos["sha1"] = self.get_sha1()
        infos["sha256"] = self.get_sha256()
        infos["sha512"] = self.get_sha512()
        infos["ssdeep"] = self.get_ssdeep()
        infos["type"] = self.get_type()
        infos["yara"] = self.get_yara()

        return infos
<|endoftext|>


<commit_before># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import time
import shutil
import logging
from threading import Thread, Lock

from lib.cuckoo.common.exceptions import CuckooAnalysisError, CuckooMachineError, CuckooGuestError
from lib.cuckoo.common.abstracts import Dictionary, MachineManager
from lib.cuckoo.common.utils import File, create_folders
from lib.cuckoo.common.config import Config
from lib.cuckoo.core.database import Database
from lib.cuckoo.core.guest import GuestManager
from lib.cuckoo.core.sniffer import Sniffer
from lib.cuckoo.core.processor import Processor
from lib.cuckoo.core.reporter import Reporter

log = logging.getLogger(__name__)

mmanager = None
machine_lock = Lock()

class AnalysisManager(Thread):
    """Analysis manager thread."""

    def __init__(self, task):
        """@param task: task."""
        Thread.__init__(self)
        Thread.daemon = True
        self.task = task
        self.cfg = Config()
        self.analysis = Dictionary()

    def init_storage(self):
        """Initialize analyses storage folder.
        @raise CuckooAnalysisError: if storage folder already exists."""
        self.analysis.results_folder = os.path.join(os.path.join(os.getcwd(), "storage", "analyses"), str(self.task.id))

        if os.path.exists(self.analysis.results_folder):
            raise CuckooAnalysisError("Analysis results folder already exists at path \"%s\", analysis aborted" % self.analysis.results_folder)

        os.mkdir(self.analysis.results_folder)

    def store_file(self):
        """Store sample file.
        @raise CuckooAnalysisError: if unable to store file."""
        md5 = File(self.task.file_path).get_md5()
        self.analysis.stored_file_path = os.path.join(os.path.join(os.getcwd(), "storage", "binaries"), md5)

        if os.path.exists(self.analysis.stored_file_path):
            log.info("File already exists at \"%s\"" % self.analysis.stored_file_path)
        else:
            try:
                shutil.copy(self.task.file_path, self.analysis.stored_file_path)
            except (IOError, shutil.Error) as e:
                raise CuckooAnalysisError("Unable to store file from \"%s\" to \"%s\", analysis aborted"
                                          % (self.task.file_path, self.analysis.stored_file_path))

        try:
            os.symlink(self.analysis.stored_file_path, os.path.join(self.analysis.results_folder, "binary"))
        except OSError as e:
            raise CuckooAnalysisError("Unable to create symlink from \"%s\" to \"%s\"" % (self.analysis.stored_file_path, self.analysis.results_folder))

    def build_options(self):
        """Get analysis options.
        @return: options dict.
        """
        options = {}

        if self.task.timeout:
            timeout = self.task.timeout
        else:
            timeout = self.cfg.cuckoo.analysis_timeout

        options["file_path"] = self.task.file_path
        options["file_name"] = File(self.task.file_path).get_name()
        options["file_type"] = File(self.task.file_path).get_type()
        options["package"] = self.task.package
        options["options"] = self.task.options
        options["timeout"] = timeout
        options["started"] = time.time()

        return options

    def launch_analysis(self):
        """Start analysis.
        @raise CuckooAnalysisError: if unable to start analysis.
        """
        log.info("Starting analysis of file \"%s\"" % self.task.file_path)

        if not os.path.exists(self.task.file_path):
            raise CuckooAnalysisError("The file to analyze does not exist at path \"%s\", analysis aborted" % self.task.file_path)

        self.init_storage()
        self.store_file()
        options = self.build_options()

        while True:
            machine_lock.acquire()
            vm = mmanager.acquire(label=self.task.machine, platform=self.task.platform)
            machine_lock.release()
            if not vm:
                log.debug("No machine available")
                time.sleep(1)
            else:
                log.info("Acquired machine %s (Label: %s)" % (vm.id, vm.label))
                break

        # Initialize sniffer
        if self.cfg.cuckoo.use_sniffer:
            sniffer = Sniffer(self.cfg.cuckoo.tcpdump)
            sniffer.start(interface=self.cfg.cuckoo.interface, host=vm.ip, file_path=os.path.join(self.analysis.results_folder, "dump.pcap"))
        else:
            sniffer = False

        try:
            # Start machine
            mmanager.start(vm.label)
            # Initialize guest manager
            guest = GuestManager(vm.ip, vm.platform)
            # Launch analysis
            guest.start_analysis(options)
            # Wait for analysis to complete
            success = guest.wait_for_completion()
            # Stop sniffer
            if sniffer:
                sniffer.stop()

            if not success:
                raise CuckooAnalysisError("Analysis failed, review previous errors")
            # Save results
            guest.save_results(self.analysis.results_folder)
        except (CuckooMachineError, CuckooGuestError) as e:
            raise CuckooAnalysisError(e.message)
        finally:
            # Stop machine
            mmanager.stop(vm.label)
            # Release the machine from lock
            mmanager.release(vm.label)

        # Launch reports generation
        Reporter(self.analysis.results_folder).run(Processor(self.analysis.results_folder).run())

        log.info("Reports generation completed (path=%s)" % self.analysis.results_folder)

    def run(self):
        """Run manager thread."""
        success = True

        db = Database()
        db.lock(self.task.id)

        try:
            self.launch_analysis()
        except CuckooMachineError as e:
            log.error("Please check virtual machine status: %s" % e.message)
            success = False
        except CuckooAnalysisError as e:
            log.error(e.message)
            success = False
        finally:
            db.complete(self.task.id, success)

class Scheduler:
    """Task scheduler."""

    def __init__(self):
        self.running = True
        self.cfg = Config()
        self.db = Database()

    def initialize(self):
        """Initialize machine manager."""
        global mmanager

        name = "modules.machinemanagers.%s" % self.cfg.cuckoo.machine_manager
        try:
            __import__(name, globals(), locals(), ["dummy"], -1)
        except ImportError as e:
            raise CuckooMachineError("Unable to import machine manager plugin: %s" % e)

        MachineManager()
        module = MachineManager.__subclasses__()[0]
        mmanager = module()
        mmanager.initialize(self.cfg.cuckoo.machine_manager)

        if len(mmanager.machines) == 0:
            raise CuckooMachineError("No machines available")
        else:
            log.info("Loaded %s machine/s" % len(mmanager.machines))

    def stop(self):
        """Stop scheduler."""
        self.running = False

        # Shutdown vm alive.
        # TODO: in future this code may be moved.
        if len(mmanager.running()) > 0:
            log.info("Shutting down guests")
            for machine in mmanager.running():
                try:
                    mmanager.stop(machine.label)
                except CuckooMachineError as e:
                    log.error("Unble to shudown machine %s, please check manually. Error: %s" % (machine.label, e.message))

    def start(self):
        """Start scheduler."""
        self.initialize()

        while self.running:
            time.sleep(1)

            if mmanager.availables() == 0:
                log.debug("No machines available, try again")
                continue

            task = self.db.fetch()

            if not task:
                log.debug("No pending tasks, try again")
                continue

            analysis = AnalysisManager(task)
            analysis.daemon = True
            analysis.start()
<commit_msg>Found another case where the CUCKOO_ROOT should be considered in scheduler.py, rather than using the cwd<commit_after># Copyright (C) 2010-2012 Cuckoo Sandbox Developers.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import time
import shutil
import logging
from threading import Thread, Lock

from lib.cuckoo.common.exceptions import CuckooAnalysisError, CuckooMachineError, CuckooGuestError
from lib.cuckoo.common.abstracts import Dictionary, MachineManager
from lib.cuckoo.common.utils import File, create_folders
from lib.cuckoo.common.config import Config
from lib.cuckoo.core.database import Database
from lib.cuckoo.core.guest import GuestManager
from lib.cuckoo.core.sniffer import Sniffer
from lib.cuckoo.core.processor import Processor
from lib.cuckoo.core.reporter import Reporter
from lib.cuckoo.common.constants import CUCKOO_ROOT

log = logging.getLogger(__name__)

mmanager = None
machine_lock = Lock()

class AnalysisManager(Thread):
    """Analysis manager thread."""

    def __init__(self, task):
        """@param task: task."""
        Thread.__init__(self)
        Thread.daemon = True
        self.task = task
        self.cfg = Config()
        self.analysis = Dictionary()

    def init_storage(self):
        """Initialize analyses storage folder.
        @raise CuckooAnalysisError: if storage folder already exists."""
        self.analysis.results_folder = os.path.join(os.path.join(CUCKOO_ROOT, "storage", "analyses"), str(self.task.id))

        if os.path.exists(self.analysis.results_folder):
            raise CuckooAnalysisError("Analysis results folder already exists at path \"%s\", analysis aborted" % self.analysis.results_folder)

        os.mkdir(self.analysis.results_folder)

    def store_file(self):
        """Store sample file.
        @raise CuckooAnalysisError: if unable to store file."""
        md5 = File(self.task.file_path).get_md5()
        self.analysis.stored_file_path = os.path.join(os.path.join(os.getcwd(), "storage", "binaries"), md5)

        if os.path.exists(self.analysis.stored_file_path):
            log.info("File already exists at \"%s\"" % self.analysis.stored_file_path)
        else:
            try:
                shutil.copy(self.task.file_path, self.analysis.stored_file_path)
            except (IOError, shutil.Error) as e:
                raise CuckooAnalysisError("Unable to store file from \"%s\" to \"%s\", analysis aborted"
                                          % (self.task.file_path, self.analysis.stored_file_path))

        try:
            os.symlink(self.analysis.stored_file_path, os.path.join(self.analysis.results_folder, "binary"))
        except OSError as e:
            raise CuckooAnalysisError("Unable to create symlink from \"%s\" to \"%s\"" % (self.analysis.stored_file_path, self.analysis.results_folder))

    def build_options(self):
        """Get analysis options.
        @return: options dict.
        """
        options = {}

        if self.task.timeout:
            timeout = self.task.timeout
        else:
            timeout = self.cfg.cuckoo.analysis_timeout

        options["file_path"] = self.task.file_path
        options["file_name"] = File(self.task.file_path).get_name()
        options["file_type"] = File(self.task.file_path).get_type()
        options["package"] = self.task.package
        options["options"] = self.task.options
        options["timeout"] = timeout
        options["started"] = time.time()

        return options

    def launch_analysis(self):
        """Start analysis.
        @raise CuckooAnalysisError: if unable to start analysis.
        """
        log.info("Starting analysis of file \"%s\"" % self.task.file_path)

        if not os.path.exists(self.task.file_path):
            raise CuckooAnalysisError("The file to analyze does not exist at path \"%s\", analysis aborted" % self.task.file_path)

        self.init_storage()
        self.store_file()
        options = self.build_options()

        while True:
            machine_lock.acquire()
            vm = mmanager.acquire(label=self.task.machine, platform=self.task.platform)
            machine_lock.release()
            if not vm:
                log.debug("No machine available")
                time.sleep(1)
            else:
                log.info("Acquired machine %s (Label: %s)" % (vm.id, vm.label))
                break

        # Initialize sniffer
        if self.cfg.cuckoo.use_sniffer:
            sniffer = Sniffer(self.cfg.cuckoo.tcpdump)
            sniffer.start(interface=self.cfg.cuckoo.interface, host=vm.ip, file_path=os.path.join(self.analysis.results_folder, "dump.pcap"))
        else:
            sniffer = False

        try:
            # Start machine
            mmanager.start(vm.label)
            # Initialize guest manager
            guest = GuestManager(vm.ip, vm.platform)
            # Launch analysis
            guest.start_analysis(options)
            # Wait for analysis to complete
            success = guest.wait_for_completion()
            # Stop sniffer
            if sniffer:
                sniffer.stop()

            if not success:
                raise CuckooAnalysisError("Analysis failed, review previous errors")
            # Save results
            guest.save_results(self.analysis.results_folder)
        except (CuckooMachineError, CuckooGuestError) as e:
            raise CuckooAnalysisError(e.message)
        finally:
            # Stop machine
            mmanager.stop(vm.label)
            # Release the machine from lock
            mmanager.release(vm.label)

        # Launch reports generation
        Reporter(self.analysis.results_folder).run(Processor(self.analysis.results_folder).run())

        log.info("Reports generation completed (path=%s)" % self.analysis.results_folder)

    def run(self):
        """Run manager thread."""
        success = True

        db = Database()
        db.lock(self.task.id)

        try:
            self.launch_analysis()
        except CuckooMachineError as e:
            log.error("Please check virtual machine status: %s" % e.message)
            success = False
        except CuckooAnalysisError as e:
            log.error(e.message)
            success = False
        finally:
            db.complete(self.task.id, success)

class Scheduler:
    """Task scheduler."""

    def __init__(self):
        self.running = True
        self.cfg = Config()
        self.db = Database()

    def initialize(self):
        """Initialize machine manager."""
        global mmanager

        name = "modules.machinemanagers.%s" % self.cfg.cuckoo.machine_manager
        try:
            __import__(name, globals(), locals(), ["dummy"], -1)
        except ImportError as e:
            raise CuckooMachineError("Unable to import machine manager plugin: %s" % e)

        MachineManager()
        module = MachineManager.__subclasses__()[0]
        mmanager = module()
        mmanager.initialize(self.cfg.cuckoo.machine_manager)

        if len(mmanager.machines) == 0:
            raise CuckooMachineError("No machines available")
        else:
            log.info("Loaded %s machine/s" % len(mmanager.machines))

    def stop(self):
        """Stop scheduler."""
        self.running = False

        # Shutdown vm alive.
        # TODO: in future this code may be moved.
        if len(mmanager.running()) > 0:
            log.info("Shutting down guests")
            for machine in mmanager.running():
                try:
                    mmanager.stop(machine.label)
                except CuckooMachineError as e:
                    log.error("Unble to shudown machine %s, please check manually. Error: %s" % (machine.label, e.message))

    def start(self):
        """Start scheduler."""
        self.initialize()

        while self.running:
            time.sleep(1)

            if mmanager.availables() == 0:
                log.debug("No machines available, try again")
                continue

            task = self.db.fetch()

            if not task:
                log.debug("No pending tasks, try again")
                continue

            analysis = AnalysisManager(task)
            analysis.daemon = True
            analysis.start()
<|endoftext|>