<commit_before>#!/usr/bin/env python

"""
@package mi.core.instrument.data_particle_generator Base data particle generator
@file mi/core/instrument/data_particle_generator.py
@author Steve Foley
@brief Contains logic to generate data particles to be exchanged between
the driver and agent. This involves a JSON interchange format
"""

import time
import ntplib
import base64
import json

from mi.core.common import BaseEnum
from mi.core.exceptions import SampleException, ReadOnlyException, NotImplementedException, InstrumentParameterException
from mi.core.log import get_logger

log = get_logger()

__author__ = 'Steve Foley'
__license__ = 'Apache 2.0'

class CommonDataParticleType(BaseEnum):
    """
    This enum defines all the common particle types defined in the modules.  Currently there is only one, but by
    using an enum here we have the opportunity to define more common data particles.
    """
    RAW = "raw"

class DataParticleKey(BaseEnum):
    PKT_FORMAT_ID = "pkt_format_id"
    PKT_VERSION = "pkt_version"
    STREAM_NAME = "stream_name"
    INTERNAL_TIMESTAMP = "internal_timestamp"
    PORT_TIMESTAMP = "port_timestamp"
    DRIVER_TIMESTAMP = "driver_timestamp"
    PREFERRED_TIMESTAMP = "preferred_timestamp"
    QUALITY_FLAG = "quality_flag"
    VALUES = "values"
    VALUE_ID = "value_id"
    VALUE = "value"
    BINARY = "binary"
    NEW_SEQUENCE = "new_sequence"

class DataParticleValue(BaseEnum):
    JSON_DATA = "JSON_Data"
    ENG = "eng"
    OK = "ok"
    CHECKSUM_FAILED = "checksum_failed"
    OUT_OF_RANGE = "out_of_range"
    INVALID = "invalid"
    QUESTIONABLE = "questionable"

class DataParticle(object):
    """
    This class is responsible for storing and ultimately generating data
    particles in the designated format from the associated inputs. It
    fills in fields as necessary, and is a valid Data Particle
    that can be sent up to the InstrumentAgent.

    It is the intent that this class is subclassed as needed if an instrument must
    modify fields in the outgoing packet. The hope is to have most of the superclass
    code be called by the child class with just values overridden as needed.
    """

    # data particle type is intended to be defined in each derived data particle class.  This value should be unique
    # for all data particles.  Best practice is to access this variable using the accessor method:
    # data_particle_type()
    _data_particle_type = None

    def __init__(self, raw_data,
                 port_timestamp=None,
                 internal_timestamp=None,
                 preferred_timestamp=None,
                 quality_flag=DataParticleValue.OK,
                 new_sequence=None):
        """ Build a particle seeded with appropriate information

        @param raw_data The raw data used in the particle
        """
        if new_sequence is not None and not isinstance(new_sequence, bool):
            raise TypeError("new_sequence is not a bool")

        self.contents = {
            DataParticleKey.PKT_FORMAT_ID: DataParticleValue.JSON_DATA,
            DataParticleKey.PKT_VERSION: 1,
            DataParticleKey.PORT_TIMESTAMP: port_timestamp,
            DataParticleKey.INTERNAL_TIMESTAMP: internal_timestamp,
            DataParticleKey.DRIVER_TIMESTAMP: ntplib.system_to_ntp_time(time.time()),
            DataParticleKey.PREFERRED_TIMESTAMP: preferred_timestamp,
            DataParticleKey.QUALITY_FLAG: quality_flag,
        }
        self._encoding_errors = []
        if new_sequence is not None:
            self.contents[DataParticleKey.NEW_SEQUENCE] = new_sequence

        self.raw_data = raw_data
        self._values = None

    def __eq__(self, arg):
        """
        Quick equality check for testing purposes. If they have the same raw
        data, timestamp, they are the same enough for this particle
        """
        allowed_diff = .000001
        if self._data_particle_type != arg._data_particle_type:
            log.debug('Data particle type does not match: %s %s', self._data_particle_type, arg._data_particle_type)
            return False

        if self.raw_data != arg.raw_data:
            log.debug('Raw data does not match')
            return False

        t1 = self.contents[DataParticleKey.INTERNAL_TIMESTAMP]
        t2 = arg.contents[DataParticleKey.INTERNAL_TIMESTAMP]

        if (t1 is None) or (t2 is None):
            tdiff = allowed_diff
        else:
            tdiff = abs(t1 - t2)

        if tdiff > allowed_diff:
            log.debug('Timestamp %s does not match %s', t1, t2)
            return False

        generated1 = json.loads(self.generate())
        generated2 = json.loads(arg.generate())
        missing, differing = self._compare(generated1, generated2, ignore_keys=[DataParticleKey.DRIVER_TIMESTAMP,
                                                                                DataParticleKey.PREFERRED_TIMESTAMP])
        if missing:
            log.error('Key mismatch between particle dictionaries: %r', missing)
            return False

        if differing:
            log.error('Value mismatch between particle dictionaries: %r', differing)

        return True

    @staticmethod
    def _compare(d1, d2, ignore_keys=None):
        ignore_keys = ignore_keys if ignore_keys else []
        missing = set(d1).symmetric_difference(d2)
        differing = {}
        for k in d1:
            if k in ignore_keys or k in missing:
                continue
            if d1[k] != d2[k]:
                differing[k] = (d1[k], d2[k])

        return missing, differing

    def set_internal_timestamp(self, timestamp=None, unix_time=None):
        """
        Set the internal timestamp
        @param timestamp: NTP timestamp to set
        @param unit_time: Unix time as returned from time.time()
        @raise InstrumentParameterException if timestamp or unix_time not supplied
        """
        if timestamp is None and unix_time is None:
            raise InstrumentParameterException("timestamp or unix_time required")

        if unix_time is not None:
            timestamp = ntplib.system_to_ntp_time(unix_time)

        # Do we want this to happen here or in down stream processes?
        # if(not self._check_timestamp(timestamp)):
        #    raise InstrumentParameterException("invalid timestamp")

        self.contents[DataParticleKey.INTERNAL_TIMESTAMP] = float(timestamp)

    def set_port_timestamp(self, timestamp=None, unix_time=None):
        """
        Set the port timestamp
        @param timestamp: NTP timestamp to set
        @param unix_time: Unix time as returned from time.time()
        @raise InstrumentParameterException if timestamp or unix_time not supplied
        """
        if timestamp is None and unix_time is None:
            raise InstrumentParameterException("timestamp or unix_time required")

        if unix_time is not None:
            timestamp = ntplib.system_to_ntp_time(unix_time)

        # Do we want this to happen here or in down stream processes?
        if not self._check_timestamp(timestamp):
            raise InstrumentParameterException("invalid timestamp")

        self.contents[DataParticleKey.PORT_TIMESTAMP] = float(timestamp)

    def set_value(self, id, value):
        """
        Set a content value, restricted as necessary

        @param id The ID of the value to set, should be from DataParticleKey
        @param value The value to set
        @raises ReadOnlyException If the parameter cannot be set
        """
        if (id == DataParticleKey.INTERNAL_TIMESTAMP) and (self._check_timestamp(value)):
            self.contents[DataParticleKey.INTERNAL_TIMESTAMP] = value
        else:
            raise ReadOnlyException("Parameter %s not able to be set to %s after object creation!" %
                                    (id, value))

    def get_value(self, id):
        """ Return a stored value

        @param id The ID (from DataParticleKey) for the parameter to return
        @raises NotImplementedException If there is an invalid id
        """
        if DataParticleKey.has(id):
            return self.contents[id]
        else:
            raise NotImplementedException("Value %s not available in particle!", id)

    def data_particle_type(self):
        """
        Return the data particle type (aka stream name)
        @raise: NotImplementedException if _data_particle_type is not set
        """
        if self._data_particle_type is None:
            raise NotImplementedException("_data_particle_type not initialized")

        return self._data_particle_type

    def generate_dict(self):
        """
        Generate a simple dictionary of sensor data and timestamps, without
        going to JSON. This is useful for the times when JSON is not needed to
        go across an interface. There are times when particles are used
        internally to a component/process/module/etc.
        @retval A python dictionary with the proper timestamps and data values
        @throws InstrumentDriverException if there is a problem wtih the inputs
        """
        # verify preferred timestamp exists in the structure...
        if not self._check_preferred_timestamps():
            raise SampleException("Preferred timestamp not in particle!")

        # build response structure
        self._encoding_errors = []
        if self._values is None:
            self._values = self._build_parsed_values()
        result = self._build_base_structure()
        result[DataParticleKey.STREAM_NAME] = self.data_particle_type()
        result[DataParticleKey.VALUES] = self._values

        return result

    def generate(self, sorted=False):
        """
        Generates a JSON_parsed packet from a sample dictionary of sensor data and
        associates a timestamp with it

        @param sorted Returned sorted json dict, useful for testing, but slow,
           so dont do it unless it is important
        @return A JSON_raw string, properly structured with port agent time stamp
           and driver timestamp
        @throws InstrumentDriverException If there is a problem with the inputs
        """
        json_result = json.dumps(self.generate_dict(), sort_keys=sorted)
        return json_result

    def _build_parsed_values(self):
        """
        Build values of a parsed structure. Just the values are built so
        so that a child class can override this class, but call it with
        super() to get the base structure before modification

        @return the values tag for this data structure ready to JSONify
        @raises SampleException when parsed values can not be properly returned
        """
        raise SampleException("Parsed values block not overridden")

    def _build_base_structure(self):
        """
        Build the base/header information for an output structure.
        Follow on methods can then modify it by adding or editing values.

        @return A fresh copy of a core structure to be exported
        """
        result = dict(self.contents)
        # clean out optional fields that were missing
        if not self.contents[DataParticleKey.PORT_TIMESTAMP]:
            del result[DataParticleKey.PORT_TIMESTAMP]
        if not self.contents[DataParticleKey.INTERNAL_TIMESTAMP]:
            del result[DataParticleKey.INTERNAL_TIMESTAMP]
        return result

    def _check_timestamp(self, timestamp):
        """
        Check to make sure the timestamp is reasonable

        @param timestamp An NTP4 formatted timestamp (64bit)
        @return True if timestamp is okay or None, False otherwise
        """
        if timestamp is None:
            return True
        if not isinstance(timestamp, float):
            return False

        # is it sufficiently in the future to be unreasonable?
        if timestamp > ntplib.system_to_ntp_time(time.time() + (86400 * 365)):
            return False
        else:
            return True

    def _check_preferred_timestamps(self):
        """
        Check to make sure the preferred timestamp indicated in the
        particle is actually listed, possibly adjusting to 2nd best
        if not there.

        @throws SampleException When there is a problem with the preferred
            timestamp in the sample.
        """
        if self.contents[DataParticleKey.PREFERRED_TIMESTAMP] is None:
            raise SampleException("Missing preferred timestamp, %s, in particle" %
                                  self.contents[DataParticleKey.PREFERRED_TIMESTAMP])

        # This should be handled downstream.  Don't want to not publish data because
        # the port agent stopped putting out timestamps
        # if self.contents[self.contents[DataParticleKey.PREFERRED_TIMESTAMP]] == None:
        #    raise SampleException("Preferred timestamp, %s, is not defined" %
        #                          self.contents[DataParticleKey.PREFERRED_TIMESTAMP])

        return True

    def _encode_value(self, name, value, encoding_function, value_range=None):
        """
        Encode a value using the encoding function, if it fails store the error in a queue

        :param value_range  tuple containing min/max numerical values or min/max lengths
        """
        encoded_val = None

        # noinspection PyBroadException
        # - custom encoding_function exceptions are not known a priori
        try:
            encoded_val = encoding_function(value)

        except ValueError as e:
            log.error('Unable to convert %s to %s.', encoded_val, encoding_function)
            self._encoding_errors.append({name: value})
        except Exception as e:
            log.error('Data particle error encoding. Name: %s Value: %s, Encoding: %s', name, value, encoding_function)
            self._encoding_errors.append({name: value})

        # optional range checking
        if value_range:
            try:
                vmin, vmax = value_range

            except ValueError as e:  # this only occurs as a programming error and should cause the parser to exit
                log.exception('_encode_value must have exactly two values for tuple argument value_range')
                raise ValueError(e)

            if encoding_function in [int, float]:
                if vmin and encoded_val < vmin:
                    log.error('Particle value (%s) below minimum threshold (%s < %s). Value', name, value, vmin)
                    self._encoding_errors.append({name: value})
                elif vmax and encoded_val > vmax:
                    log.error('Particle value (%s) exceeds maximum threshold (%s > %s). Value', name, value, vmax)
                    self._encoding_errors.append({name: value})
            elif hasattr(encoded_val, '__len__'):
                try:
                    if vmin and len(encoded_val) < vmin:
                        log.error('Particle value (%s) length below minimum threshold (%s < %s). Value',
                                  name, value, vmin)
                        self._encoding_errors.append({name: value})
                    elif vmax and len(encoded_val) > vmax:
                        log.error('Particle value (%s) length exceeds maximum threshold (%s > %s). Value',
                                  name, value, vmax)
                        self._encoding_errors.append({name: value})
                # in the unlikely event that a range was specified and the encoding object created a bogus len()
                # we'll just ignore the range check
                except TypeError:
                    log.warning('_encode_value received an encoding function (%s) that claimed to implement len() but '
                                'does not. Unable to apply range test to %s', encoding_function, name)

        return {DataParticleKey.VALUE_ID: name,
                DataParticleKey.VALUE: encoded_val}

    def get_encoding_errors(self):
        """
        Return the encoding errors list
        """
        return self._encoding_errors

class RawDataParticleKey(BaseEnum):
    PAYLOAD = "raw"
    LENGTH = "length"
    TYPE = "type"
    CHECKSUM = "checksum"

class RawDataParticle(DataParticle):
    """
    This class a common data particle for generating data particles of raw
    data.

    It essentially is a translation of the port agent packet
    """
    _data_particle_type = CommonDataParticleType.RAW

    def _build_parsed_values(self):
        """
        Build a particle out of a port agent packet.
        @returns A list that is ready to be added to the "values" tag before
           the structure is JSONified
        """

        port_agent_packet = self.raw_data
        if not isinstance(port_agent_packet, dict):
            raise SampleException("raw data not a dictionary")

        for param in ["raw", "length", "type", "checksum"]:
            if param not in port_agent_packet:
                raise SampleException("raw data not a complete port agent packet. missing %s" % param)

        payload = None
        length = None
        type = None
        checksum = None

        # Attempt to convert values
        try:
            payload = base64.b64encode(port_agent_packet.get("raw"))
        except TypeError:
            pass

        try:
            length = int(port_agent_packet.get("length"))
        except TypeError:
            pass

        try:
            type = int(port_agent_packet.get("type"))
        except TypeError:
            pass

        try:
            checksum = int(port_agent_packet.get("checksum"))
        except TypeError:
            pass

        result = [{
            DataParticleKey.VALUE_ID: RawDataParticleKey.PAYLOAD,
            DataParticleKey.VALUE: payload,
            DataParticleKey.BINARY: True},
            {
                DataParticleKey.VALUE_ID: RawDataParticleKey.LENGTH,
                DataParticleKey.VALUE: length},
            {
                DataParticleKey.VALUE_ID: RawDataParticleKey.TYPE,
                DataParticleKey.VALUE: type},
            {
                DataParticleKey.VALUE_ID: RawDataParticleKey.CHECKSUM,
                DataParticleKey.VALUE: checksum},
        ]

        return result
<commit_msg>Corrected logging per review.<commit_after>#!/usr/bin/env python

"""
@package mi.core.instrument.data_particle_generator Base data particle generator
@file mi/core/instrument/data_particle_generator.py
@author Steve Foley
@brief Contains logic to generate data particles to be exchanged between
the driver and agent. This involves a JSON interchange format
"""

import time
import ntplib
import base64
import json

from mi.core.common import BaseEnum
from mi.core.exceptions import SampleException, ReadOnlyException, NotImplementedException, InstrumentParameterException
from mi.core.log import get_logger

log = get_logger()

__author__ = 'Steve Foley'
__license__ = 'Apache 2.0'

class CommonDataParticleType(BaseEnum):
    """
    This enum defines all the common particle types defined in the modules.  Currently there is only one, but by
    using an enum here we have the opportunity to define more common data particles.
    """
    RAW = "raw"

class DataParticleKey(BaseEnum):
    PKT_FORMAT_ID = "pkt_format_id"
    PKT_VERSION = "pkt_version"
    STREAM_NAME = "stream_name"
    INTERNAL_TIMESTAMP = "internal_timestamp"
    PORT_TIMESTAMP = "port_timestamp"
    DRIVER_TIMESTAMP = "driver_timestamp"
    PREFERRED_TIMESTAMP = "preferred_timestamp"
    QUALITY_FLAG = "quality_flag"
    VALUES = "values"
    VALUE_ID = "value_id"
    VALUE = "value"
    BINARY = "binary"
    NEW_SEQUENCE = "new_sequence"

class DataParticleValue(BaseEnum):
    JSON_DATA = "JSON_Data"
    ENG = "eng"
    OK = "ok"
    CHECKSUM_FAILED = "checksum_failed"
    OUT_OF_RANGE = "out_of_range"
    INVALID = "invalid"
    QUESTIONABLE = "questionable"

class DataParticle(object):
    """
    This class is responsible for storing and ultimately generating data
    particles in the designated format from the associated inputs. It
    fills in fields as necessary, and is a valid Data Particle
    that can be sent up to the InstrumentAgent.

    It is the intent that this class is subclassed as needed if an instrument must
    modify fields in the outgoing packet. The hope is to have most of the superclass
    code be called by the child class with just values overridden as needed.
    """

    # data particle type is intended to be defined in each derived data particle class.  This value should be unique
    # for all data particles.  Best practice is to access this variable using the accessor method:
    # data_particle_type()
    _data_particle_type = None

    def __init__(self, raw_data,
                 port_timestamp=None,
                 internal_timestamp=None,
                 preferred_timestamp=None,
                 quality_flag=DataParticleValue.OK,
                 new_sequence=None):
        """ Build a particle seeded with appropriate information

        @param raw_data The raw data used in the particle
        """
        if new_sequence is not None and not isinstance(new_sequence, bool):
            raise TypeError("new_sequence is not a bool")

        self.contents = {
            DataParticleKey.PKT_FORMAT_ID: DataParticleValue.JSON_DATA,
            DataParticleKey.PKT_VERSION: 1,
            DataParticleKey.PORT_TIMESTAMP: port_timestamp,
            DataParticleKey.INTERNAL_TIMESTAMP: internal_timestamp,
            DataParticleKey.DRIVER_TIMESTAMP: ntplib.system_to_ntp_time(time.time()),
            DataParticleKey.PREFERRED_TIMESTAMP: preferred_timestamp,
            DataParticleKey.QUALITY_FLAG: quality_flag,
        }
        self._encoding_errors = []
        if new_sequence is not None:
            self.contents[DataParticleKey.NEW_SEQUENCE] = new_sequence

        self.raw_data = raw_data
        self._values = None

    def __eq__(self, arg):
        """
        Quick equality check for testing purposes. If they have the same raw
        data, timestamp, they are the same enough for this particle
        """
        allowed_diff = .000001
        if self._data_particle_type != arg._data_particle_type:
            log.debug('Data particle type does not match: %s %s', self._data_particle_type, arg._data_particle_type)
            return False

        if self.raw_data != arg.raw_data:
            log.debug('Raw data does not match')
            return False

        t1 = self.contents[DataParticleKey.INTERNAL_TIMESTAMP]
        t2 = arg.contents[DataParticleKey.INTERNAL_TIMESTAMP]

        if (t1 is None) or (t2 is None):
            tdiff = allowed_diff
        else:
            tdiff = abs(t1 - t2)

        if tdiff > allowed_diff:
            log.debug('Timestamp %s does not match %s', t1, t2)
            return False

        generated1 = json.loads(self.generate())
        generated2 = json.loads(arg.generate())
        missing, differing = self._compare(generated1, generated2, ignore_keys=[DataParticleKey.DRIVER_TIMESTAMP,
                                                                                DataParticleKey.PREFERRED_TIMESTAMP])
        if missing:
            log.error('Key mismatch between particle dictionaries: %r', missing)
            return False

        if differing:
            log.error('Value mismatch between particle dictionaries: %r', differing)

        return True

    @staticmethod
    def _compare(d1, d2, ignore_keys=None):
        ignore_keys = ignore_keys if ignore_keys else []
        missing = set(d1).symmetric_difference(d2)
        differing = {}
        for k in d1:
            if k in ignore_keys or k in missing:
                continue
            if d1[k] != d2[k]:
                differing[k] = (d1[k], d2[k])

        return missing, differing

    def set_internal_timestamp(self, timestamp=None, unix_time=None):
        """
        Set the internal timestamp
        @param timestamp: NTP timestamp to set
        @param unit_time: Unix time as returned from time.time()
        @raise InstrumentParameterException if timestamp or unix_time not supplied
        """
        if timestamp is None and unix_time is None:
            raise InstrumentParameterException("timestamp or unix_time required")

        if unix_time is not None:
            timestamp = ntplib.system_to_ntp_time(unix_time)

        # Do we want this to happen here or in down stream processes?
        # if(not self._check_timestamp(timestamp)):
        #    raise InstrumentParameterException("invalid timestamp")

        self.contents[DataParticleKey.INTERNAL_TIMESTAMP] = float(timestamp)

    def set_port_timestamp(self, timestamp=None, unix_time=None):
        """
        Set the port timestamp
        @param timestamp: NTP timestamp to set
        @param unix_time: Unix time as returned from time.time()
        @raise InstrumentParameterException if timestamp or unix_time not supplied
        """
        if timestamp is None and unix_time is None:
            raise InstrumentParameterException("timestamp or unix_time required")

        if unix_time is not None:
            timestamp = ntplib.system_to_ntp_time(unix_time)

        # Do we want this to happen here or in down stream processes?
        if not self._check_timestamp(timestamp):
            raise InstrumentParameterException("invalid timestamp")

        self.contents[DataParticleKey.PORT_TIMESTAMP] = float(timestamp)

    def set_value(self, id, value):
        """
        Set a content value, restricted as necessary

        @param id The ID of the value to set, should be from DataParticleKey
        @param value The value to set
        @raises ReadOnlyException If the parameter cannot be set
        """
        if (id == DataParticleKey.INTERNAL_TIMESTAMP) and (self._check_timestamp(value)):
            self.contents[DataParticleKey.INTERNAL_TIMESTAMP] = value
        else:
            raise ReadOnlyException("Parameter %s not able to be set to %s after object creation!" %
                                    (id, value))

    def get_value(self, id):
        """ Return a stored value

        @param id The ID (from DataParticleKey) for the parameter to return
        @raises NotImplementedException If there is an invalid id
        """
        if DataParticleKey.has(id):
            return self.contents[id]
        else:
            raise NotImplementedException("Value %s not available in particle!", id)

    def data_particle_type(self):
        """
        Return the data particle type (aka stream name)
        @raise: NotImplementedException if _data_particle_type is not set
        """
        if self._data_particle_type is None:
            raise NotImplementedException("_data_particle_type not initialized")

        return self._data_particle_type

    def generate_dict(self):
        """
        Generate a simple dictionary of sensor data and timestamps, without
        going to JSON. This is useful for the times when JSON is not needed to
        go across an interface. There are times when particles are used
        internally to a component/process/module/etc.
        @retval A python dictionary with the proper timestamps and data values
        @throws InstrumentDriverException if there is a problem wtih the inputs
        """
        # verify preferred timestamp exists in the structure...
        if not self._check_preferred_timestamps():
            raise SampleException("Preferred timestamp not in particle!")

        # build response structure
        self._encoding_errors = []
        if self._values is None:
            self._values = self._build_parsed_values()
        result = self._build_base_structure()
        result[DataParticleKey.STREAM_NAME] = self.data_particle_type()
        result[DataParticleKey.VALUES] = self._values

        return result

    def generate(self, sorted=False):
        """
        Generates a JSON_parsed packet from a sample dictionary of sensor data and
        associates a timestamp with it

        @param sorted Returned sorted json dict, useful for testing, but slow,
           so dont do it unless it is important
        @return A JSON_raw string, properly structured with port agent time stamp
           and driver timestamp
        @throws InstrumentDriverException If there is a problem with the inputs
        """
        json_result = json.dumps(self.generate_dict(), sort_keys=sorted)
        return json_result

    def _build_parsed_values(self):
        """
        Build values of a parsed structure. Just the values are built so
        so that a child class can override this class, but call it with
        super() to get the base structure before modification

        @return the values tag for this data structure ready to JSONify
        @raises SampleException when parsed values can not be properly returned
        """
        raise SampleException("Parsed values block not overridden")

    def _build_base_structure(self):
        """
        Build the base/header information for an output structure.
        Follow on methods can then modify it by adding or editing values.

        @return A fresh copy of a core structure to be exported
        """
        result = dict(self.contents)
        # clean out optional fields that were missing
        if not self.contents[DataParticleKey.PORT_TIMESTAMP]:
            del result[DataParticleKey.PORT_TIMESTAMP]
        if not self.contents[DataParticleKey.INTERNAL_TIMESTAMP]:
            del result[DataParticleKey.INTERNAL_TIMESTAMP]
        return result

    def _check_timestamp(self, timestamp):
        """
        Check to make sure the timestamp is reasonable

        @param timestamp An NTP4 formatted timestamp (64bit)
        @return True if timestamp is okay or None, False otherwise
        """
        if timestamp is None:
            return True
        if not isinstance(timestamp, float):
            return False

        # is it sufficiently in the future to be unreasonable?
        if timestamp > ntplib.system_to_ntp_time(time.time() + (86400 * 365)):
            return False
        else:
            return True

    def _check_preferred_timestamps(self):
        """
        Check to make sure the preferred timestamp indicated in the
        particle is actually listed, possibly adjusting to 2nd best
        if not there.

        @throws SampleException When there is a problem with the preferred
            timestamp in the sample.
        """
        if self.contents[DataParticleKey.PREFERRED_TIMESTAMP] is None:
            raise SampleException("Missing preferred timestamp, %s, in particle" %
                                  self.contents[DataParticleKey.PREFERRED_TIMESTAMP])

        # This should be handled downstream.  Don't want to not publish data because
        # the port agent stopped putting out timestamps
        # if self.contents[self.contents[DataParticleKey.PREFERRED_TIMESTAMP]] == None:
        #    raise SampleException("Preferred timestamp, %s, is not defined" %
        #                          self.contents[DataParticleKey.PREFERRED_TIMESTAMP])

        return True

    def _encode_value(self, name, value, encoding_function, value_range=None):
        """
        Encode a value using the encoding function, if it fails store the error in a queue

        :param value_range  tuple containing min/max numerical values or min/max lengths
        """
        encoded_val = None

        # noinspection PyBroadException
        # - custom encoding_function exceptions are not known a priori
        try:
            encoded_val = encoding_function(value)

        except ValueError as e:
            log.error('Unable to convert %s to %s.', encoded_val, encoding_function)
            self._encoding_errors.append({name: value})
        except Exception as e:
            log.error('Data particle error encoding. Name: %s Value: %s, Encoding: %s', name, value, encoding_function)
            self._encoding_errors.append({name: value})

        # optional range checking
        if value_range:
            try:
                vmin, vmax = value_range

            except ValueError as e:  # this only occurs as a programming error and should cause the parser to exit
                log.exception('_encode_value must have exactly two values for tuple argument value_range')
                raise ValueError(e)

            if encoding_function in [int, float]:
                if vmin and encoded_val < vmin:
                    log.error('Particle value (%s) below minimum threshold (%s < %s)', name, value, vmin)
                    self._encoding_errors.append({name: value})
                elif vmax and encoded_val > vmax:
                    log.error('Particle value (%s) exceeds maximum threshold (%s > %s)', name, value, vmax)
                    self._encoding_errors.append({name: value})
            elif hasattr(encoded_val, '__len__'):
                try:
                    if vmin and len(encoded_val) < vmin:
                        log.error('Particle value (%s) length below minimum threshold (%s < %s)',
                                  name, value, vmin)
                        self._encoding_errors.append({name: value})
                    elif vmax and len(encoded_val) > vmax:
                        log.error('Particle value (%s) length exceeds maximum threshold (%s > %s)',
                                  name, value, vmax)
                        self._encoding_errors.append({name: value})
                # in the unlikely event that a range was specified and the encoding object created a bogus len()
                # we'll just ignore the range check
                except TypeError:
                    log.warning('_encode_value received an encoding function (%s) that claimed to implement len() but '
                                'does not. Unable to apply range test to %s', encoding_function, name)

        return {DataParticleKey.VALUE_ID: name,
                DataParticleKey.VALUE: encoded_val}

    def get_encoding_errors(self):
        """
        Return the encoding errors list
        """
        return self._encoding_errors

class RawDataParticleKey(BaseEnum):
    PAYLOAD = "raw"
    LENGTH = "length"
    TYPE = "type"
    CHECKSUM = "checksum"

class RawDataParticle(DataParticle):
    """
    This class a common data particle for generating data particles of raw
    data.

    It essentially is a translation of the port agent packet
    """
    _data_particle_type = CommonDataParticleType.RAW

    def _build_parsed_values(self):
        """
        Build a particle out of a port agent packet.
        @returns A list that is ready to be added to the "values" tag before
           the structure is JSONified
        """

        port_agent_packet = self.raw_data
        if not isinstance(port_agent_packet, dict):
            raise SampleException("raw data not a dictionary")

        for param in ["raw", "length", "type", "checksum"]:
            if param not in port_agent_packet:
                raise SampleException("raw data not a complete port agent packet. missing %s" % param)

        payload = None
        length = None
        type = None
        checksum = None

        # Attempt to convert values
        try:
            payload = base64.b64encode(port_agent_packet.get("raw"))
        except TypeError:
            pass

        try:
            length = int(port_agent_packet.get("length"))
        except TypeError:
            pass

        try:
            type = int(port_agent_packet.get("type"))
        except TypeError:
            pass

        try:
            checksum = int(port_agent_packet.get("checksum"))
        except TypeError:
            pass

        result = [{
            DataParticleKey.VALUE_ID: RawDataParticleKey.PAYLOAD,
            DataParticleKey.VALUE: payload,
            DataParticleKey.BINARY: True},
            {
                DataParticleKey.VALUE_ID: RawDataParticleKey.LENGTH,
                DataParticleKey.VALUE: length},
            {
                DataParticleKey.VALUE_ID: RawDataParticleKey.TYPE,
                DataParticleKey.VALUE: type},
            {
                DataParticleKey.VALUE_ID: RawDataParticleKey.CHECKSUM,
                DataParticleKey.VALUE: checksum},
        ]

        return result
<|endoftext|>


<commit_before>from django.db.models import get_model
from django.template import Node, TemplateSyntaxError, Library, Variable
from django.utils.translation import ugettext as _
from tagging.models import TaggedItem
register = Library()

class TaggedGetRelatedNode(Node):
    def __init__(self, obj, queryset_or_model, context_var, **kwargs):
        self.obj = Variable(obj)
        self.queryset_or_model = queryset_or_model
        self.context_var = context_var
        self.kwargs = kwargs

    def render(self, context):
        try:
            param = get_model(*self.queryset_or_model.split('.'))
        except:
            param = Variable(self.queryset_or_model).resolve(context)

        if param is None:
            raise TemplateSyntaxError(_('tagged_get_related tag was given an invalid input: %s') % self.queryset_or_model)

        context[self.context_var] = TaggedItem.objects.get_related(self.obj.resolve(context), param, **self.kwargs)
        return ''

@register.tag
def tagged_get_related(parser, token):
    """
    {% tagged_get_related [obj] in ([appname]).[queryset_or_model] as [varname] with num=10 %}
    """
    kwargs = {}
    bits = token.contents.split()
    if len(bits) != 6 and len(bits) != 8:
        raise TemplateSyntaxError(_('%s tag requires either five or seven arguments') % bits[0])
    if bits[2] != 'in':
        raise TemplateSyntaxError(_("second argument to %s tag must be 'in'") % bits[0])
    if bits[4] != 'as':
        raise TemplateSyntaxError(_("fourth argument to %s tag must be 'as'") % bits[0])
    if len(bits) == 8:
        if bits[6] != 'with':
            raise TemplateSyntaxError(_("if given, fourth argument to %s tag must be 'with'") % bits[0])
        try:
            name, value = bits[7].split('=')
            if name == 'num':
                try:
                    kwargs[str(name)] = int(value)
                except ValueError:
                    raise TemplateSyntaxError(_("%(tag)s tag's '%(option)s' option was not a valid integer: '%(value)s'") % {
                        'tag': bits[0],
                        'option': name,
                        'value': value,
                    })
            else:
                raise TemplateSyntaxError(_("%(tag)s tag was given an invalid option: '%(option)s'") % {
                    'tag': bits[0],
                    'option': name,
                })
        except ValueError:
            raise TemplateSyntaxError(_("%(tag)s tag was given a badly formatted option: '%(option)s'") % {
                'tag': bits[0],
                'option': bits[i],
            })

    return TaggedGetRelatedNode(bits[1], bits[3], bits[5], **kwargs)

<commit_msg>fixed issue with draft posts appearing in related, works but still needs more effort for cleaner execution<commit_after>from django.db.models import get_model
from django.template import Node, TemplateSyntaxError, Library, Variable
from django.utils.translation import ugettext as _
from tagging.models import TaggedItem
register = Library()

class TaggedGetRelatedNode(Node):
    def __init__(self, obj, queryset_or_model, context_var, **kwargs):
        self.obj = Variable(obj)
        self.queryset_or_model = queryset_or_model
        self.context_var = context_var
        self.kwargs = kwargs

    def render(self, context):
        try:
            param = get_model(*self.queryset_or_model.split('.'))
        except:
            param = Variable(self.queryset_or_model).resolve(context)

        if param is None:
            raise TemplateSyntaxError(_('tagged_get_related tag was given an invalid input: %s') % self.queryset_or_model)

        objs = list(TaggedItem.objects.get_related(self.obj.resolve(context), param, **self.kwargs))
        for obj in objs:
            if hasattr(obj, 'status'):
                if obj.status != 2:
                    objs.remove(obj)

        context[self.context_var] = objs
        return ''

@register.tag
def tagged_get_related(parser, token):
    """
    {% tagged_get_related [obj] in ([appname]).[queryset_or_model] as [varname] with num=10 %}
    """
    kwargs = {}
    bits = token.contents.split()
    if len(bits) != 6 and len(bits) != 8:
        raise TemplateSyntaxError(_('%s tag requires either five or seven arguments') % bits[0])
    if bits[2] != 'in':
        raise TemplateSyntaxError(_("second argument to %s tag must be 'in'") % bits[0])
    if bits[4] != 'as':
        raise TemplateSyntaxError(_("fourth argument to %s tag must be 'as'") % bits[0])
    if len(bits) == 8:
        if bits[6] != 'with':
            raise TemplateSyntaxError(_("if given, fourth argument to %s tag must be 'with'") % bits[0])
        try:
            name, value = bits[7].split('=')
            if name == 'num':
                try:
                    kwargs[str(name)] = int(value)
                except ValueError:
                    raise TemplateSyntaxError(_("%(tag)s tag's '%(option)s' option was not a valid integer: '%(value)s'") % {
                        'tag': bits[0],
                        'option': name,
                        'value': value,
                    })
            else:
                raise TemplateSyntaxError(_("%(tag)s tag was given an invalid option: '%(option)s'") % {
                    'tag': bits[0],
                    'option': name,
                })
        except ValueError:
            raise TemplateSyntaxError(_("%(tag)s tag was given a badly formatted option: '%(option)s'") % {
                'tag': bits[0],
                'option': bits[i],
            })

    return TaggedGetRelatedNode(bits[1], bits[3], bits[5], **kwargs)

<|endoftext|>


<commit_before>#!/usr/bin/python

import sys

# import the USB and Time libraries into Python
import usb.core, usb.util, time

# do we want to track the estimated position?

# motors

motor_map_bits = {
        '5'  : { 'byte': 1, 'bits': [1,    2]}, # Base
        '4'  : { 'byte': 0, 'bits': [64, 128]}, # Shoulder Down
        '3'  : { 'byte': 0, 'bits': [16,  32]}, # Elbow
        '2'  : { 'byte': 0, 'bits': [4,    8]}, # Wrist
        '1'  : { 'byte': 0, 'bits': [1,    2]}, # Grab
        };

light_bits =  { 'byte': 2, 'bits': [1] };

"""Low level driver for the OWI Edge"""
class EdgeRaw:

   """ sets up class; initialises arm instance  """
   def __init__(self):
      # initialise essentials
      self._arm = usb.core.find(idVendor=0x1267, idProduct=0x0000)
      # check this
      if self._arm is None:
         raise ValueError("Arm not found")
      # and check this
      self.stop()

   """ stops the arm """
   def stop(self):
      self._arm.ctrl_transfer(0x40,6,0x100,0,[0, 0, 0],1000)

   """applies output bit set for the duration"""
   def output(self, duration, motor_bytes):
#      print('moving', motor_bytes, 'for ', duration)
      # Start the movement
      self._arm.ctrl_transfer(0x40,6,0x100, 0, motor_bytes, 1000)
      time.sleep(duration)
      # Stop the movement after waiting specified duration
      self.stop()

   """applies output motor set for the duration"""
   def output(self, duration, motors):
#      print('moving', motors, 'for ', duration)
      # Start the movement
      # {1: {'dir': 1}, 2: {'dir': -1}, 'D': 1.0}
      motor_bytes = [0, 0, 0]
      self._arm.ctrl_transfer(0x40,6,0x100,0,motor_bytes,1000)
      time.sleep(duration)
      # Stop the movement after waiting specified duration
      self.stop()

<commit_msg>fixing unintentional override<commit_after>#!/usr/bin/python

import sys

# import the USB and Time libraries into Python
import usb.core, usb.util, time

# do we want to track the estimated position?

# motors

motor_map_bits = {
        '5'  : { 'byte': 1, 'bits': [1,    2]}, # Base
        '4'  : { 'byte': 0, 'bits': [64, 128]}, # Shoulder Down
        '3'  : { 'byte': 0, 'bits': [16,  32]}, # Elbow
        '2'  : { 'byte': 0, 'bits': [4,    8]}, # Wrist
        '1'  : { 'byte': 0, 'bits': [1,    2]}, # Grab
        };

light_bits =  { 'byte': 2, 'bits': [1] };

"""Low level driver for the OWI Edge"""
class EdgeRaw:

   """ sets up class; initialises arm instance  """
   def __init__(self):
      # initialise essentials
      self._arm = usb.core.find(idVendor=0x1267, idProduct=0x0000)
      # check this
      if self._arm is None:
         raise ValueError("Arm not found")
      # and check this
      self.stop()

   """ stops the arm """
   def stop(self):
      self._arm.ctrl_transfer(0x40,6,0x100,0,[0, 0, 0],1000)

   """applies output bit set for the duration"""
   def output(self, duration, motor_bytes):
#      print('moving', motor_bytes, 'for ', duration)
      # Start the movement
      self._arm.ctrl_transfer(0x40,6,0x100, 0, motor_bytes, 1000)
      time.sleep(duration)
      # Stop the movement after waiting specified duration
      self.stop()

   """applies output motor set for the duration"""
   def set_motors(self, duration, motors):
#      print('moving', motors, 'for ', duration)
      # Start the movement
      # {1: {'dir': 1}, 2: {'dir': -1}, 'D': 1.0}
      # TODO(PMM) map the motor bytes to the output bytes
      motor_bytes = [0, 0, 0]
      self._arm.ctrl_transfer(0x40,6,0x100,0,motor_bytes,1000)
      time.sleep(motors['D'])
      # Stop the movement after waiting specified duration
      self.stop()

<|endoftext|>


<commit_before>import argparse

def getArgs():
    parser = argparse.ArgumentParser(description='Autoprover')
    parser.add_argument("file", type=open, help="a theorem to be proved")
    args = parser.parse_args()
    return args
<commit_msg>add more parameter in parser<commit_after>import argparse

def getArgs():
    parser = argparse.ArgumentParser(description='Autoprover')
    parser.add_argument("file", type=open, help="a theorem to be proved")
    parser.add_argument("-p", "--population", dest='population',
            type=int, help="population size, default is 1000")
    parser.add_argument("-g", "--max-generation", dest='maxGeneration',
            type=int, help="a maximun generation, defalut is 100")
    parser.add_argument("-m", "--mutate-rate", dest='mutateRate',
            type=float, help="a mutate rate in interval (0, 1), default is 0.25")
    parser.add_argument("-e", "--elite-rate", dest='eliteRate',
            type=float, help="the percentage of elitism, default is 0")
    parser.add_argument("-c", "--cross-rate", dest='crossRate',
            type=int, help="a cross rate, defalut is 0.5")
    parser.add_argument("-t", "--cross-type", dest='crossType',
            type=int, help="different type of cross, default is 0")
    parser.add_argument("-n", "--verify-num", dest='verifyNum',
            type=int, help="number of top proofs to verify in each generation, default is 50")
    args = parser.parse_args()
    return args
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright (C) 2010-2014 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import random
import fnmatch
import logging
import argparse

try:
    import requests
    HAVE_REQUESTS = True
except ImportError:
    HAVE_REQUESTS = False

logging.basicConfig()

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

from lib.cuckoo.common.colors import bold, green, red, yellow
from lib.cuckoo.common.objects import File
from lib.cuckoo.common.utils import to_unicode
from lib.cuckoo.core.database import Database

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("target", type=str, help="URL, path to the file or folder to analyze")
    parser.add_argument("--remote", type=str, action="store", default=None, help="Specify IP:port to a Cuckoo API server to submit remotely", required=False)
    parser.add_argument("--url", action="store_true", default=False, help="Specify whether the target is an URL", required=False)
    parser.add_argument("--package", type=str, action="store", default="", help="Specify an analysis package", required=False)
    parser.add_argument("--custom", type=str, action="store", default="", help="Specify any custom value", required=False)
    parser.add_argument("--timeout", type=int, action="store", default=0, help="Specify an analysis timeout", required=False)
    parser.add_argument("--options", type=str, action="store", default="", help="Specify options for the analysis package (e.g. \"name=value,name2=value2\")", required=False)
    parser.add_argument("--priority", type=int, action="store", default=1, help="Specify a priority for the analysis represented by an integer", required=False)
    parser.add_argument("--machine", type=str, action="store", default="", help="Specify the identifier of a machine you want to use", required=False)
    parser.add_argument("--platform", type=str, action="store", default="", help="Specify the operating system platform you want to use (windows/darwin/linux)", required=False)
    parser.add_argument("--memory", action="store_true", default=False, help="Enable to take a memory dump of the analysis machine", required=False)
    parser.add_argument("--enforce-timeout", action="store_true", default=False, help="Enable to force the analysis to run for the full timeout period", required=False)
    parser.add_argument("--clock", type=str, action="store", default=None, help="Set virtual machine clock", required=False)
    parser.add_argument("--tags", type=str, action="store", default=None, help="Specify tags identifier of a machine you want to use", required=False)
    parser.add_argument("--max", type=int, action="store", default=None, help="Maximum samples to add in a row", required=False)
    parser.add_argument("--pattern", type=str, action="store", default=None, help="Pattern of files to submit", required=False)
    parser.add_argument("--shuffle", action="store_true", default=False, help="Shuffle samples before submitting them", required=False)
    parser.add_argument("--unique", action="store_true", default=False, help="Only submit new samples, ignore duplicates", required=False)
    parser.add_argument("--quiet", action="store_true", default=False, help="Only print text on failure", required=False)

    try:
        args = parser.parse_args()
    except IOError as e:
        parser.error(e)
        return False

    # If the quiet flag has been set, then we also disable the "warning"
    # level of the logging module. (E.g., when pydeep has not been installed,
    # there will be a warning message, because Cuckoo can't resolve the
    # ssdeep hash of this particular sample.)
    if args.quiet:
        logging.disable(logging.WARNING)

    db = Database()

    target = to_unicode(args.target)

    if args.url:
        if args.remote:
            if not HAVE_REQUESTS:
                print(bold(red("Error")) + ": you need to install python-requests (`pip install requests`)")
                return False

            url = "http://{0}/tasks/create/url".format(args.remote)

            data = dict(
                url=target,
                package=args.package,
                timeout=args.timeout,
                options=args.options,
                priority=args.priority,
                machine=args.machine,
                platform=args.platform,
                memory=args.memory,
                enforce_timeout=args.enforce_timeout,
                custom=args.custom,
                tags=args.tags
            )

            try:
                response = requests.post(url, data=data)
            except Exception as e:
                print(bold(red("Error")) + ": unable to send URL: {0}".format(e))
                return False

            json = response.json()
            task_id = json["task_id"]
        else:
            task_id = db.add_url(target,
                                 package=args.package,
                                 timeout=args.timeout,
                                 options=args.options,
                                 priority=args.priority,
                                 machine=args.machine,
                                 platform=args.platform,
                                 custom=args.custom,
                                 memory=args.memory,
                                 enforce_timeout=args.enforce_timeout,
                                 clock=args.clock,
                                 tags=args.tags)

        if task_id:
            if not args.quiet:
                print(bold(green("Success")) + u": URL \"{0}\" added as task with ID {1}".format(target, task_id))
        else:
            print(bold(red("Error")) + ": adding task to database")
    else:
        # Get absolute path to deal with relative.
        path = to_unicode(os.path.abspath(target))

        if not os.path.exists(path):
            print(bold(red("Error")) + u": the specified file/folder does not exist at path \"{0}\"".format(path))
            return False

        files = []
        if os.path.isdir(path):
            for dirname, dirnames, filenames in os.walk(path):
                for file_name in filenames:
                    file_path = os.path.join(dirname, file_name)

                    if os.path.isfile(file_path):
                        if args.pattern:
                            if fnmatch.fnmatch(file_name, args.pattern):
                                files.append(to_unicode(file_path))
                        else:
                            files.append(to_unicode(file_path))
        else:
            files.append(path)

        if args.shuffle:
            random.shuffle(files)
        else:
            sorted(files)

        for file_path in files:
            if not File(file_path).get_size():
                if not args.quiet:
                    print(bold(yellow("Empty") + ": sample {0} (skipping file)".format(file_path)))

                continue

            if not args.max is None:
                # Break if the maximum number of samples has been reached.
                if not args.max:
                    break

                args.max -= 1

            if args.remote:
                if not HAVE_REQUESTS:
                    print(bold(red("Error")) + ": you need to install python-requests (`pip install requests`)")
                    return False

                url = "http://{0}/tasks/create/file".format(args.remote)

                files = dict(
                    file=open(file_path, "rb"),
                    filename=os.path.basename(file_path)
                )

                data = dict(
                    package=args.package,
                    timeout=args.timeout,
                    options=args.options,
                    priority=args.priority,
                    machine=args.machine,
                    platform=args.platform,
                    memory=args.memory,
                    enforce_timeout=args.enforce_timeout,
                    custom=args.custom,
                    tags=args.tags
                )

                try:
                    response = requests.post(url, files=files, data=data)
                except Exception as e:
                    print(bold(red("Error")) + ": unable to send file: {0}".format(e))
                    return False

                json = response.json()
                task_id = json["task_id"]
            else:
                if args.unique:
                    sha256 = File(file_path).get_sha256()
                    if not db.find_sample(sha256=sha256) is None:
                        msg = ": Sample {0} (skipping file)".format(file_path)
                        if not args.quiet:
                            print(bold(yellow("Duplicate")) + msg)
                        continue

                task_id = db.add_path(file_path=file_path,
                                      package=args.package,
                                      timeout=args.timeout,
                                      options=args.options,
                                      priority=args.priority,
                                      machine=args.machine,
                                      platform=args.platform,
                                      custom=args.custom,
                                      memory=args.memory,
                                      enforce_timeout=args.enforce_timeout,
                                      clock=args.clock,
                                      tags=args.tags)

            if task_id:
                if not args.quiet:
                    print(bold(green("Success")) + u": File \"{0}\" added as task with ID {1}".format(file_path, task_id))
            else:
                print(bold(red("Error")) + ": adding task to database")

if __name__ == "__main__":
    main()
<commit_msg>Ooops.. editor didn't save it<commit_after>#!/usr/bin/env python
# Copyright (C) 2010-2014 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import random
import fnmatch
import logging
import argparse

try:
    import requests
    HAVE_REQUESTS = True
except ImportError:
    HAVE_REQUESTS = False

logging.basicConfig()

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

from lib.cuckoo.common.colors import bold, green, red, yellow
from lib.cuckoo.common.objects import File
from lib.cuckoo.common.utils import to_unicode
from lib.cuckoo.core.database import Database

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("target", type=str, help="URL, path to the file or folder to analyze")
    parser.add_argument("--remote", type=str, action="store", default=None, help="Specify IP:port to a Cuckoo API server to submit remotely", required=False)
    parser.add_argument("--url", action="store_true", default=False, help="Specify whether the target is an URL", required=False)
    parser.add_argument("--package", type=str, action="store", default="", help="Specify an analysis package", required=False)
    parser.add_argument("--custom", type=str, action="store", default="", help="Specify any custom value", required=False)
    parser.add_argument("--timeout", type=int, action="store", default=0, help="Specify an analysis timeout", required=False)
    parser.add_argument("--options", type=str, action="store", default="", help="Specify options for the analysis package (e.g. \"name=value,name2=value2\")", required=False)
    parser.add_argument("--priority", type=int, action="store", default=1, help="Specify a priority for the analysis represented by an integer", required=False)
    parser.add_argument("--machine", type=str, action="store", default="", help="Specify the identifier of a machine you want to use", required=False)
    parser.add_argument("--platform", type=str, action="store", default="", help="Specify the operating system platform you want to use (windows/darwin/linux)", required=False)
    parser.add_argument("--memory", action="store_true", default=False, help="Enable to take a memory dump of the analysis machine", required=False)
    parser.add_argument("--enforce-timeout", action="store_true", default=False, help="Enable to force the analysis to run for the full timeout period", required=False)
    parser.add_argument("--clock", type=str, action="store", default=None, help="Set virtual machine clock", required=False)
    parser.add_argument("--tags", type=str, action="store", default=None, help="Specify tags identifier of a machine you want to use", required=False)
    parser.add_argument("--max", type=int, action="store", default=None, help="Maximum samples to add in a row", required=False)
    parser.add_argument("--pattern", type=str, action="store", default=None, help="Pattern of files to submit", required=False)
    parser.add_argument("--shuffle", action="store_true", default=False, help="Shuffle samples before submitting them", required=False)
    parser.add_argument("--unique", action="store_true", default=False, help="Only submit new samples, ignore duplicates", required=False)
    parser.add_argument("--quiet", action="store_true", default=False, help="Only print text on failure", required=False)

    try:
        args = parser.parse_args()
    except IOError as e:
        parser.error(e)
        return False

    # If the quiet flag has been set, then we also disable the "warning"
    # level of the logging module. (E.g., when pydeep has not been installed,
    # there will be a warning message, because Cuckoo can't resolve the
    # ssdeep hash of this particular sample.)
    if args.quiet:
        logging.disable(logging.WARNING)

    db = Database()

    target = to_unicode(args.target)

    if args.url:
        if args.remote:
            if not HAVE_REQUESTS:
                print(bold(red("Error")) + ": you need to install python-requests (`pip install requests`)")
                return False

            url = "http://{0}/tasks/create/url".format(args.remote)

            data = dict(
                url=target,
                package=args.package,
                timeout=args.timeout,
                options=args.options,
                priority=args.priority,
                machine=args.machine,
                platform=args.platform,
                memory=args.memory,
                enforce_timeout=args.enforce_timeout,
                custom=args.custom,
                tags=args.tags
            )

            try:
                response = requests.post(url, data=data)
            except Exception as e:
                print(bold(red("Error")) + ": unable to send URL: {0}".format(e))
                return False

            json = response.json()
            task_id = json["task_id"]
        else:
            task_id = db.add_url(target,
                                 package=args.package,
                                 timeout=args.timeout,
                                 options=args.options,
                                 priority=args.priority,
                                 machine=args.machine,
                                 platform=args.platform,
                                 custom=args.custom,
                                 memory=args.memory,
                                 enforce_timeout=args.enforce_timeout,
                                 clock=args.clock,
                                 tags=args.tags)

        if task_id:
            if not args.quiet:
                print(bold(green("Success")) + u": URL \"{0}\" added as task with ID {1}".format(target, task_id))
        else:
            print(bold(red("Error")) + ": adding task to database")
    else:
        # Get absolute path to deal with relative.
        path = to_unicode(os.path.abspath(target))

        if not os.path.exists(path):
            print(bold(red("Error")) + u": the specified file/folder does not exist at path \"{0}\"".format(path))
            return False

        files = []
        if os.path.isdir(path):
            for dirname, dirnames, filenames in os.walk(path):
                for file_name in filenames:
                    file_path = os.path.join(dirname, file_name)

                    if os.path.isfile(file_path):
                        if args.pattern:
                            if fnmatch.fnmatch(file_name, args.pattern):
                                files.append(to_unicode(file_path))
                        else:
                            files.append(to_unicode(file_path))
        else:
            files.append(path)

        if args.shuffle:
            random.shuffle(files)
        else:
            files = sorted(files)

        for file_path in files:
            if not File(file_path).get_size():
                if not args.quiet:
                    print(bold(yellow("Empty") + ": sample {0} (skipping file)".format(file_path)))

                continue

            if not args.max is None:
                # Break if the maximum number of samples has been reached.
                if not args.max:
                    break

                args.max -= 1

            if args.remote:
                if not HAVE_REQUESTS:
                    print(bold(red("Error")) + ": you need to install python-requests (`pip install requests`)")
                    return False

                url = "http://{0}/tasks/create/file".format(args.remote)

                files = dict(
                    file=open(file_path, "rb"),
                    filename=os.path.basename(file_path)
                )

                data = dict(
                    package=args.package,
                    timeout=args.timeout,
                    options=args.options,
                    priority=args.priority,
                    machine=args.machine,
                    platform=args.platform,
                    memory=args.memory,
                    enforce_timeout=args.enforce_timeout,
                    custom=args.custom,
                    tags=args.tags
                )

                try:
                    response = requests.post(url, files=files, data=data)
                except Exception as e:
                    print(bold(red("Error")) + ": unable to send file: {0}".format(e))
                    return False

                json = response.json()
                task_id = json["task_id"]
            else:
                if args.unique:
                    sha256 = File(file_path).get_sha256()
                    if not db.find_sample(sha256=sha256) is None:
                        msg = ": Sample {0} (skipping file)".format(file_path)
                        if not args.quiet:
                            print(bold(yellow("Duplicate")) + msg)
                        continue

                task_id = db.add_path(file_path=file_path,
                                      package=args.package,
                                      timeout=args.timeout,
                                      options=args.options,
                                      priority=args.priority,
                                      machine=args.machine,
                                      platform=args.platform,
                                      custom=args.custom,
                                      memory=args.memory,
                                      enforce_timeout=args.enforce_timeout,
                                      clock=args.clock,
                                      tags=args.tags)

            if task_id:
                if not args.quiet:
                    print(bold(green("Success")) + u": File \"{0}\" added as task with ID {1}".format(file_path, task_id))
            else:
                print(bold(red("Error")) + ": adding task to database")

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


<commit_before><commit_msg>comment out test in correlate.py<commit_after><|endoftext|>


<commit_before>from .map import map
from .run import run
from .each import each
from .some import some
from .race import race
from .once import once
from .wait import wait
from .curry import curry
from .wraps import wraps
from .apply import apply
from .defer import defer
from .every import every
from .until import until
from .times import times
from .gather import gather
from .repeat import repeat
from .filter import filter
from .filterfalse import filterfalse
from .reduce import reduce
from .whilst import whilst
from .series import series
from .partial import partial
from .timeout import timeout, TimeoutLimit
from .compose import compose
from .flat_map import flat_map
from .constant import constant
from .throttle import throttle
from .dropwhile import dropwhile
from .concurrent import ConcurrentExecutor, concurrent

__author__ = 'Tomas Aparicio'
__license__ = 'MIT'

# Current package version
__version__ = '0.1.7'

# Explicit symbols to export
__all__ = (
    'ConcurrentExecutor',
    'apply',
    'compose',
    'concurrent',
    'constant',
    'curry',
    'defer',
    'dropwhile',
    'each',
    'every',
    'filter',
    'filterfalse',
    'flat_map',
    'gather',
    'map',
    'once',
    'partial',
    'race',
    'reduce',
    'repeat',
    'run',
    'series',
    'some',
    'throttle',
    'timeout',
    'TimeoutLimit',
    'times',
    'until',
    'wait',
    'whilst',
    'wraps',
)
<commit_msg>Bump version: 0.1.7 → 0.1.8<commit_after>from .map import map
from .run import run
from .each import each
from .some import some
from .race import race
from .once import once
from .wait import wait
from .curry import curry
from .wraps import wraps
from .apply import apply
from .defer import defer
from .every import every
from .until import until
from .times import times
from .gather import gather
from .repeat import repeat
from .filter import filter
from .filterfalse import filterfalse
from .reduce import reduce
from .whilst import whilst
from .series import series
from .partial import partial
from .timeout import timeout, TimeoutLimit
from .compose import compose
from .flat_map import flat_map
from .constant import constant
from .throttle import throttle
from .dropwhile import dropwhile
from .concurrent import ConcurrentExecutor, concurrent

__author__ = 'Tomas Aparicio'
__license__ = 'MIT'

# Current package version
__version__ = '0.1.8'

# Explicit symbols to export
__all__ = (
    'ConcurrentExecutor',
    'apply',
    'compose',
    'concurrent',
    'constant',
    'curry',
    'defer',
    'dropwhile',
    'each',
    'every',
    'filter',
    'filterfalse',
    'flat_map',
    'gather',
    'map',
    'once',
    'partial',
    'race',
    'reduce',
    'repeat',
    'run',
    'series',
    'some',
    'throttle',
    'timeout',
    'TimeoutLimit',
    'times',
    'until',
    'wait',
    'whilst',
    'wraps',
)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

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

import mock
import testtools

from heatclient.common import template_utils

import shade
from shade import meta
from shade.tests import fakes
from shade.tests.unit import base

class TestStack(base.TestCase):

    def setUp(self):
        super(TestStack, self).setUp()
        self.cloud = shade.openstack_cloud(validate=False)

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_list_stacks(self, mock_heat):
        fake_stacks = [
            fakes.FakeStack('001', 'stack1'),
            fakes.FakeStack('002', 'stack2'),
        ]
        mock_heat.stacks.list.return_value = fake_stacks
        stacks = self.cloud.list_stacks()
        mock_heat.stacks.list.assert_called_once_with()
        self.assertEqual(meta.obj_list_to_dict(fake_stacks), stacks)

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_list_stacks_exception(self, mock_heat):
        mock_heat.stacks.list.side_effect = Exception()
        with testtools.ExpectedException(
            shade.OpenStackCloudException,
            "Error fetching stack list"
        ):
            self.cloud.list_stacks()

    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_delete_stack(self, mock_heat, mock_get):
        stack = {'id': 'stack_id', 'name': 'stack_name'}
        mock_get.return_value = stack
        self.assertTrue(self.cloud.delete_stack('stack_name'))
        mock_get.assert_called_once_with('stack_name')
        mock_heat.stacks.delete.assert_called_once_with(id=stack['id'])

    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_delete_stack_not_found(self, mock_heat, mock_get):
        mock_get.return_value = None
        self.assertFalse(self.cloud.delete_stack('stack_name'))
        mock_get.assert_called_once_with('stack_name')
        self.assertFalse(mock_heat.stacks.delete.called)

    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_delete_stack_exception(self, mock_heat, mock_get):
        stack = {'id': 'stack_id', 'name': 'stack_name'}
        mock_get.return_value = stack
        mock_heat.stacks.delete.side_effect = Exception()
        with testtools.ExpectedException(
            shade.OpenStackCloudException,
            "Failed to delete stack %s" % stack['id']
        ):
            self.cloud.delete_stack('stack_name')

    @mock.patch.object(template_utils, 'get_template_contents')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_create_stack(self, mock_heat, mock_template):
        mock_template.return_value = ({}, {})
        self.cloud.create_stack('stack_name')
        self.assertTrue(mock_template.called)
        mock_heat.stacks.create.assert_called_once_with(
            stack_name='stack_name',
            disable_rollback=False,
            parameters={},
            template={},
            files={}
        )

    @mock.patch.object(template_utils, 'get_template_contents')
    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_create_stack_wait(self, mock_heat, mock_get, mock_template):
        stack = {'id': 'stack_id', 'name': 'stack_name'}
        mock_template.return_value = ({}, {})
        mock_get.side_effect = iter([None, stack])
        ret = self.cloud.create_stack('stack_name', wait=True)
        self.assertTrue(mock_template.called)
        mock_heat.stacks.create.assert_called_once_with(
            stack_name='stack_name',
            disable_rollback=False,
            parameters={},
            template={},
            files={}
        )
        self.assertEqual(2, mock_get.call_count)
        self.assertEqual(stack, ret)
<commit_msg>Add tests for stack search API<commit_after># -*- coding: utf-8 -*-

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

import mock
import testtools

from heatclient.common import template_utils

import shade
from shade import meta
from shade.tests import fakes
from shade.tests.unit import base

class TestStack(base.TestCase):

    def setUp(self):
        super(TestStack, self).setUp()
        self.cloud = shade.openstack_cloud(validate=False)

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_list_stacks(self, mock_heat):
        fake_stacks = [
            fakes.FakeStack('001', 'stack1'),
            fakes.FakeStack('002', 'stack2'),
        ]
        mock_heat.stacks.list.return_value = fake_stacks
        stacks = self.cloud.list_stacks()
        mock_heat.stacks.list.assert_called_once_with()
        self.assertEqual(meta.obj_list_to_dict(fake_stacks), stacks)

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_list_stacks_exception(self, mock_heat):
        mock_heat.stacks.list.side_effect = Exception()
        with testtools.ExpectedException(
            shade.OpenStackCloudException,
            "Error fetching stack list"
        ):
            self.cloud.list_stacks()

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_search_stacks(self, mock_heat):
        fake_stacks = [
            fakes.FakeStack('001', 'stack1'),
            fakes.FakeStack('002', 'stack2'),
        ]
        mock_heat.stacks.list.return_value = fake_stacks
        stacks = self.cloud.search_stacks()
        mock_heat.stacks.list.assert_called_once_with()
        self.assertEqual(meta.obj_list_to_dict(fake_stacks), stacks)

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_search_stacks_filters(self, mock_heat):
        fake_stacks = [
            fakes.FakeStack('001', 'stack1', status='GOOD'),
            fakes.FakeStack('002', 'stack2', status='BAD'),
        ]
        mock_heat.stacks.list.return_value = fake_stacks
        filters = {'stack_status': 'GOOD'}
        stacks = self.cloud.search_stacks(filters=filters)
        mock_heat.stacks.list.assert_called_once_with()
        self.assertEqual(meta.obj_list_to_dict(fake_stacks[:1]), stacks)

    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_search_stacks_exception(self, mock_heat):
        mock_heat.stacks.list.side_effect = Exception()
        with testtools.ExpectedException(
            shade.OpenStackCloudException,
            "Error fetching stack list"
        ):
            self.cloud.search_stacks()

    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_delete_stack(self, mock_heat, mock_get):
        stack = {'id': 'stack_id', 'name': 'stack_name'}
        mock_get.return_value = stack
        self.assertTrue(self.cloud.delete_stack('stack_name'))
        mock_get.assert_called_once_with('stack_name')
        mock_heat.stacks.delete.assert_called_once_with(id=stack['id'])

    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_delete_stack_not_found(self, mock_heat, mock_get):
        mock_get.return_value = None
        self.assertFalse(self.cloud.delete_stack('stack_name'))
        mock_get.assert_called_once_with('stack_name')
        self.assertFalse(mock_heat.stacks.delete.called)

    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_delete_stack_exception(self, mock_heat, mock_get):
        stack = {'id': 'stack_id', 'name': 'stack_name'}
        mock_get.return_value = stack
        mock_heat.stacks.delete.side_effect = Exception()
        with testtools.ExpectedException(
            shade.OpenStackCloudException,
            "Failed to delete stack %s" % stack['id']
        ):
            self.cloud.delete_stack('stack_name')

    @mock.patch.object(template_utils, 'get_template_contents')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_create_stack(self, mock_heat, mock_template):
        mock_template.return_value = ({}, {})
        self.cloud.create_stack('stack_name')
        self.assertTrue(mock_template.called)
        mock_heat.stacks.create.assert_called_once_with(
            stack_name='stack_name',
            disable_rollback=False,
            parameters={},
            template={},
            files={}
        )

    @mock.patch.object(template_utils, 'get_template_contents')
    @mock.patch.object(shade.OpenStackCloud, 'get_stack')
    @mock.patch.object(shade.OpenStackCloud, 'heat_client')
    def test_create_stack_wait(self, mock_heat, mock_get, mock_template):
        stack = {'id': 'stack_id', 'name': 'stack_name'}
        mock_template.return_value = ({}, {})
        mock_get.side_effect = iter([None, stack])
        ret = self.cloud.create_stack('stack_name', wait=True)
        self.assertTrue(mock_template.called)
        mock_heat.stacks.create.assert_called_once_with(
            stack_name='stack_name',
            disable_rollback=False,
            parameters={},
            template={},
            files={}
        )
        self.assertEqual(2, mock_get.call_count)
        self.assertEqual(stack, ret)
<|endoftext|>