<commit_before>'''
Authors: Donnie Marino, Kostas Stamatiou, Nate Ricklin
Contact: dmarino@digitalglobe.com

Class to represent a workflow task
'''
import json, uuid

class InvalidInputPort(AttributeError):
    pass

class InvalidOutputPort(Exception):
    pass

class WorkflowError(Exception):
    pass

class Port:
    def __init__(self, name, type, required, description, value):
        self.name = name
        self.type = type
        self.description = description
        self.required = required
        self.value = value

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        out = ""
        out += "Port %s:" % self.name
        out += "\n\ttype: %s" % self.type
        out += "\n\tdescription: %s" % self.description
        out += "\n\trequired: %s" % self.required
        out += "\n\tValue: %s" % self.value
        return out

class Inputs(object):
    def __init__(self, task):
        self._task = task
        self._portnames = [p['name'] for p in task.input_ports]
        for p in task.input_ports:
            self.__setattr__(p['name'], Port(p['name'], p['type'], p['required'], p['description'], value=None))

    # allow setting task input values like this:
    # task.inputs.port_name = value
    def __setattr__(self, k, v):
        # special handling for setting task & portname:
        if k in ['_portnames', '_task']:
            object.__setattr__(self, k, v)

        # special handling for port names
        elif k in self._portnames and hasattr(self, k):
            port = self.__getattribute__(k)
            port.value = v

        # default for everything else
        else:
            object.__setattr__(self, k, v)

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        out = ""
        for input_port in self._task.input_ports:
            out += input_port['name'] + "\n"
        return out

class Task:

    def __init__(self, interface, task_type, **kwargs):
        '''Construct an instance of GBDX Task

         Args:
            optional keywords:
                name (string): The task name.
                input_port_descriptors (list): A list of the input port descriptors.
                output_port_descriptors (list): A list of the output port descriptors.
                container_descriptors (list): A list of the container descriptors.
                properties (dict): A dictionary of task properties.
        Returns:
            An instance of Task.

        '''

        self.input_data = []
        self.name = task_type + '_' + str(uuid.uuid4())
        self.id = None

        self.interface = interface
        self.type = task_type
        self.definition = self.interface.workflow.describe_task(task_type)
        self.domain = self.definition['containerDescriptors'][0]['properties'].get('domain','default')

        self.inputs = Inputs(self)

        # all the other kwargs are input port values or sources
        self.set(**kwargs)

    # get a reference to the output port
    def get_output(self, port_name):
        output_port_names = [p['name'] for p in self.output_ports]
        if port_name not in output_port_names:
            raise InvalidOutputPort('Invalid output port %s.  Valid output ports for task %s are: %s' % (port_name, self.type, output_port_names))

        return "source:" + self.name + ":" + port_name

    # set input ports source or value
    def set( self, **kwargs ):
        for port_name, port_value in kwargs.iteritems():
            self.inputs.__getattribute__(port_name).value = port_value
        # input_port_names = [p['name'] for p in self.input_ports]
        # for input_port in kwargs.keys():
        #     if input_port in input_port_names:
        #         self.input_data.append( {
        #                                     'name': input_port,
        #                                     'value': kwargs[input_port]
        #                                 })
        #     else:
        #         raise InvalidInputPort('Invalid input port %s.  Valid input ports for task %s are: %s' % (input_port, self.type, input_port_names))

    @property
    def input_ports(self):
        return self.definition['inputPortDescriptors']

    @input_ports.setter
    def input_ports(self, value):
        raise NotImplementedError("Cannot set input ports")

    @property
    def output_ports(self):
        return self.definition['outputPortDescriptors']

    @output_ports.setter
    def output_ports(self, value):
        raise NotImplementedError("Cannot set output ports")

    def generate_task_workflow_json(self):
        d = {
                    "name": self.name,
                    "outputs": [],
                    "inputs": [],
                    "taskType": self.type,
                    "containerDescriptors": [{"properties": {"domain": self.domain}}]
                }

        for input_port_name in self.inputs._portnames:
            input_port_value = self.inputs.__getattribute__(input_port_name).value
            if input_port_value == None:
                continue

            if input_port_value == False:
                input_port_value = 'false'
            if input_port_value == True:
                input_port_value = 'true'

            if input_port_value.startswith('source:'):
                # this port is linked from a previous output task
                d['inputs'].append({
                                    "name": input_port_name,
                                    "source": input_port_value.replace('source:','')
                                })
            else:
                d['inputs'].append({
                                    "name": input_port_name,
                                    "value": input_port_value
                                })

        for output_port in self.output_ports:
            output_port_name = output_port['name']
            d['outputs'].append(  {
                    "name": output_port_name
                } )

        return d

class Workflow:
    def __init__(self, interface, tasks, **kwargs):
        self.interface = interface
        self.name = kwargs.get('name', str(uuid.uuid4()) )
        self.id = None

        self.definition = self.workflow_skeleton()

        for task in tasks:
            self.definition['tasks'].append( task.generate_task_workflow_json() )

        self.tasks = tasks

    def workflow_skeleton(self):
        return {
            "tasks": [],
            "name": self.name
        }

    def execute(self):
        if not self.tasks:
            raise WorkflowError('Workflow contains no tasks, and cannot be executed.')
        self.id = self.interface.workflow.launch(self.definition)
        return self.id

    def cancel(self):
        if not self.id:
            raise WorkflowError('Workflow is not running.  Cannot cancel.')

        self.interface.workflow.cancel(self.id)

    @property
    def status(self):
        if not self.id:
            raise WorkflowError('Workflow is not running.  Cannot check status.')
        return self.interface.workflow.status(self.id)

    @status.setter
    def status(self, value):
        raise NotImplementedError("Cannot set workflow status, readonly.")

    @property
    def complete(self):
        if not self.id:
            return False
        return self.status['state'] == 'complete'

    @complete.setter
    def complete(self, value):
        raise NotImplementedError("Cannot set workflow complete, readonly.")

    @property
    def failed(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'failed'

    @failed.setter
    def failed(self, value):
        raise NotImplementedError("Cannot set workflow failed, readonly.")

    @property
    def canceled(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'canceled'

    @canceled.setter
    def canceled(self, value):
        raise NotImplementedError("Cannot set workflow canceled, readonly.")

    @property
    def succeeded(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'succeeded'

    @succeeded.setter
    def succeeded(self, value):
        raise NotImplementedError("Cannot set workflow succeeded, readonly.")

    @property
    def running(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'running'

    @running.setter
    def running(self, value):
        raise NotImplementedError("Cannot set workflow running, readonly.")

    @property
    def timedout(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'timedout'

    @timedout.setter
    def timedout(self, value):
        raise NotImplementedError("Cannot set workflow timedout, readonly.")

<commit_msg>refactor<commit_after>'''
Authors: Donnie Marino, Kostas Stamatiou, Nate Ricklin
Contact: dmarino@digitalglobe.com

Class to represent a workflow task
'''
import json, uuid

class InvalidInputPort(AttributeError):
    pass

class InvalidOutputPort(Exception):
    pass

class WorkflowError(Exception):
    pass

class Port:
    def __init__(self, name, type, required, description, value):
        self.name = name
        self.type = type
        self.description = description
        self.required = required
        self.value = value

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        out = ""
        out += "Port %s:" % self.name
        out += "\n\ttype: %s" % self.type
        out += "\n\tdescription: %s" % self.description
        out += "\n\trequired: %s" % self.required
        out += "\n\tValue: %s" % self.value
        return out

class PortList(object):
    def __init__(self, ports):
        self._portnames = [p['name'] for p in ports]
        for p in ports:
            self.__setattr__(p['name'], Port(p['name'], p['type'], p['required'], p['description'], value=None))

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        out = ""
        for input_port_name in self._portnames:
            out += input_port_name + "\n"
        return out

class Inputs(PortList):
    # allow setting task input values like this:
    # task.inputs.port_name = value
    def __setattr__(self, k, v):
        # special handling for setting task & portname:
        if k in ['_portnames']:
            object.__setattr__(self, k, v)

        # special handling for port names
        elif k in self._portnames and hasattr(self, k):
            port = self.__getattribute__(k)
            port.value = v

        # default for everything else
        else:
            object.__setattr__(self, k, v)

class Task:

    def __init__(self, interface, task_type, **kwargs):
        '''Construct an instance of GBDX Task

         Args:
            optional keywords:
                name (string): The task name.
                input_port_descriptors (list): A list of the input port descriptors.
                output_port_descriptors (list): A list of the output port descriptors.
                container_descriptors (list): A list of the container descriptors.
                properties (dict): A dictionary of task properties.
        Returns:
            An instance of Task.

        '''

        self.input_data = []
        self.name = task_type + '_' + str(uuid.uuid4())
        self.id = None

        self.interface = interface
        self.type = task_type
        self.definition = self.interface.workflow.describe_task(task_type)
        self.domain = self.definition['containerDescriptors'][0]['properties'].get('domain','default')

        self.inputs = Inputs(self.input_ports)

        # all the other kwargs are input port values or sources
        self.set(**kwargs)

    # get a reference to the output port
    def get_output(self, port_name):
        output_port_names = [p['name'] for p in self.output_ports]
        if port_name not in output_port_names:
            raise InvalidOutputPort('Invalid output port %s.  Valid output ports for task %s are: %s' % (port_name, self.type, output_port_names))

        return "source:" + self.name + ":" + port_name

    # set input ports source or value
    def set( self, **kwargs ):
        for port_name, port_value in kwargs.iteritems():
            self.inputs.__getattribute__(port_name).value = port_value
        # input_port_names = [p['name'] for p in self.input_ports]
        # for input_port in kwargs.keys():
        #     if input_port in input_port_names:
        #         self.input_data.append( {
        #                                     'name': input_port,
        #                                     'value': kwargs[input_port]
        #                                 })
        #     else:
        #         raise InvalidInputPort('Invalid input port %s.  Valid input ports for task %s are: %s' % (input_port, self.type, input_port_names))

    @property
    def input_ports(self):
        return self.definition['inputPortDescriptors']

    @input_ports.setter
    def input_ports(self, value):
        raise NotImplementedError("Cannot set input ports")

    @property
    def output_ports(self):
        return self.definition['outputPortDescriptors']

    @output_ports.setter
    def output_ports(self, value):
        raise NotImplementedError("Cannot set output ports")

    def generate_task_workflow_json(self):
        d = {
                    "name": self.name,
                    "outputs": [],
                    "inputs": [],
                    "taskType": self.type,
                    "containerDescriptors": [{"properties": {"domain": self.domain}}]
                }

        for input_port_name in self.inputs._portnames:
            input_port_value = self.inputs.__getattribute__(input_port_name).value
            if input_port_value == None:
                continue

            if input_port_value == False:
                input_port_value = 'false'
            if input_port_value == True:
                input_port_value = 'true'

            if input_port_value.startswith('source:'):
                # this port is linked from a previous output task
                d['inputs'].append({
                                    "name": input_port_name,
                                    "source": input_port_value.replace('source:','')
                                })
            else:
                d['inputs'].append({
                                    "name": input_port_name,
                                    "value": input_port_value
                                })

        for output_port in self.output_ports:
            output_port_name = output_port['name']
            d['outputs'].append(  {
                    "name": output_port_name
                } )

        return d

class Workflow:
    def __init__(self, interface, tasks, **kwargs):
        self.interface = interface
        self.name = kwargs.get('name', str(uuid.uuid4()) )
        self.id = None

        self.definition = self.workflow_skeleton()

        for task in tasks:
            self.definition['tasks'].append( task.generate_task_workflow_json() )

        self.tasks = tasks

    def workflow_skeleton(self):
        return {
            "tasks": [],
            "name": self.name
        }

    def execute(self):
        if not self.tasks:
            raise WorkflowError('Workflow contains no tasks, and cannot be executed.')
        self.id = self.interface.workflow.launch(self.definition)
        return self.id

    def cancel(self):
        if not self.id:
            raise WorkflowError('Workflow is not running.  Cannot cancel.')

        self.interface.workflow.cancel(self.id)

    @property
    def status(self):
        if not self.id:
            raise WorkflowError('Workflow is not running.  Cannot check status.')
        return self.interface.workflow.status(self.id)

    @status.setter
    def status(self, value):
        raise NotImplementedError("Cannot set workflow status, readonly.")

    @property
    def complete(self):
        if not self.id:
            return False
        return self.status['state'] == 'complete'

    @complete.setter
    def complete(self, value):
        raise NotImplementedError("Cannot set workflow complete, readonly.")

    @property
    def failed(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'failed'

    @failed.setter
    def failed(self, value):
        raise NotImplementedError("Cannot set workflow failed, readonly.")

    @property
    def canceled(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'canceled'

    @canceled.setter
    def canceled(self, value):
        raise NotImplementedError("Cannot set workflow canceled, readonly.")

    @property
    def succeeded(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'succeeded'

    @succeeded.setter
    def succeeded(self, value):
        raise NotImplementedError("Cannot set workflow succeeded, readonly.")

    @property
    def running(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'running'

    @running.setter
    def running(self, value):
        raise NotImplementedError("Cannot set workflow running, readonly.")

    @property
    def timedout(self):
        if not self.id:
            return False
        status = self.status
        return status['state'] == 'complete' and status['event'] == 'timedout'

    @timedout.setter
    def timedout(self, value):
        raise NotImplementedError("Cannot set workflow timedout, readonly.")

<|endoftext|>


<commit_before>"""
Software Collections list output - command ``scl --list``
=========================================================================

This module provides parser for list output of ``scl``.

Parser provided by this module is:

SoftwareCollectionsListInstalled - command ``scl --list``
------------------------------------------------------------------------------------

"""

from .. import Parser, parser
from . import get_active_lines
from insights.specs import Specs

@parser(Specs.software_collections_list)
class SoftwareCollectionsListInstalled(Parser):
    """
    An object for parsing the output of ``scl --list``.

    Sample input file::

        devtoolset-7
        httpd24
        python27
        rh-mysql57
        rh-nodejs8
        rh-php71
        rh-python36
        rh-ruby24

    Examples:
        >>> type(collections)
        <class 'insights.parsers.software_collections_list.SoftwareCollectionsListInstalled'>
        >>> len(collections.records)
        8
        >>> coll1 = collections.records[0]
        >>> coll1
        'devtoolset-7'
        >>> coll2 = collections.records[4]
        >>> coll2
        'rh-nodejs8'
        >>> collections.exists('rh-ruby24')
        True
        >>> collections.exists('rh-missing_colection')
        False

    """
    def parse_content(self, content):
        self.records = []
        for line in get_active_lines(content):
            self.records.append(line)

    def exists(self, collection_name):
        """
        Checks if the collection is installed on the system.

        Args:
           service_name (str): collections name

        Returns:
            bool: True if collections exists, False otherwise.
        """
        return collection_name in self.records
<commit_msg>Added a doc string for RHEL-6/7 (#1867)<commit_after>"""
Software Collections list output - command ``scl --list`` RHEL-6/7
==================================================================

This module provides parser for list output of ``scl``. This spec
is valid for ``RHEL-6/7`` only, ``-l|--list`` is ``deprecated`` in
``RHEL-8``. On ``RHEL-8`` same functionality can be achieved by
``scl list-collections`` spec.

Parser provided by this module is:

SoftwareCollectionsListInstalled - command ``scl --list``
---------------------------------------------------------

"""

from .. import Parser, parser
from . import get_active_lines
from insights.specs import Specs

@parser(Specs.software_collections_list)
class SoftwareCollectionsListInstalled(Parser):
    """
    An object for parsing the output of ``scl --list``.

    Sample input file::

        devtoolset-7
        httpd24
        python27
        rh-mysql57
        rh-nodejs8
        rh-php71
        rh-python36
        rh-ruby24

    Examples:
        >>> type(collections)
        <class 'insights.parsers.software_collections_list.SoftwareCollectionsListInstalled'>
        >>> len(collections.records)
        8
        >>> coll1 = collections.records[0]
        >>> coll1
        'devtoolset-7'
        >>> coll2 = collections.records[4]
        >>> coll2
        'rh-nodejs8'
        >>> collections.exists('rh-ruby24')
        True
        >>> collections.exists('rh-missing_colection')
        False

    """
    def parse_content(self, content):
        self.records = []
        for line in get_active_lines(content):
            self.records.append(line)

    def exists(self, collection_name):
        """
        Checks if the collection is installed on the system.

        Args:
           service_name (str): collections name

        Returns:
            bool: True if collections exists, False otherwise.
        """
        return collection_name in self.records
<|endoftext|>


<commit_before>import json
import pytest

from insights.core.dr import SkipComponent
from insights.core.spec_factory import DatasourceProvider, simple_file
from insights.specs.datasources.cloud_init import cloud_cfg, LocalSpecs

CLOUD_CFG = """
users:
  - name: demo
    ssh-authorized-keys:
      - key_one
      - key_two

network:
  version: 1
  config:
    - type: physical
      name: eth0
      subnets:
        - type: dhcp
        - type: dhcp6
""".strip()

CLOUD_CFG_NO_NETWORK = """
users:
  - name: demo
    ssh-authorized-keys:
      - key_one
      - key_two
""".strip()

CLOUD_CFG_BAD = """
users
  name -demo
  ssh-authorized-keys
      - key_one
      - key_two
""".strip()

CLOUD_CFG_JSON = {
    'version': 1,
    'config': [
        {
            'type': 'physical',
            'name': 'eth0',
            'subnets': [
                {'type': 'dhcp'},
                {'type': 'dhcp6'}
            ]
        }
    ]
}

RELATIVE_PATH = '/etc/cloud/cloud.cfg'

def test_cloud_cfg():
    simple_file.content = CLOUD_CFG.splitlines()
    broker = {LocalSpecs.cloud_cfg_input: simple_file}
    result = cloud_cfg(broker)
    assert result is not None
    assert isinstance(result, DatasourceProvider)
    expected = DatasourceProvider(content=json.dumps(CLOUD_CFG_JSON), relative_path=RELATIVE_PATH)
    assert result.content == expected.content
    assert result.relative_path == expected.relative_path

def test_cloud_cfg_bad():
    simple_file.content = CLOUD_CFG_BAD.splitlines()
    broker = {LocalSpecs.cloud_cfg_input: simple_file}
    with pytest.raises(SkipComponent) as e:
        cloud_cfg(broker)
    assert 'Unexpected exception' in str(e)

def test_cloud_cfg_no_network():
    simple_file.content = CLOUD_CFG_NO_NETWORK.splitlines()
    broker = {LocalSpecs.cloud_cfg_input: simple_file}
    with pytest.raises(SkipComponent) as e:
        cloud_cfg(broker)
    assert 'No network section in yaml' in str(e)
<commit_msg>Fix test for cloud_init datasource (#3125)<commit_after>import json
import pytest
from mock.mock import Mock

from insights.core.dr import SkipComponent
from insights.core.spec_factory import DatasourceProvider
from insights.specs.datasources.cloud_init import cloud_cfg, LocalSpecs

CLOUD_CFG = """
users:
  - name: demo
    ssh-authorized-keys:
      - key_one
      - key_two

network:
  version: 1
  config:
    - type: physical
      name: eth0
      subnets:
        - type: dhcp
        - type: dhcp6
""".strip()

CLOUD_CFG_NO_NETWORK = """
users:
  - name: demo
    ssh-authorized-keys:
      - key_one
      - key_two
""".strip()

CLOUD_CFG_BAD = """
users
  name -demo
  ssh-authorized-keys
      - key_one
      - key_two
""".strip()

CLOUD_CFG_JSON = {
    'version': 1,
    'config': [
        {
            'type': 'physical',
            'name': 'eth0',
            'subnets': [
                {'type': 'dhcp'},
                {'type': 'dhcp6'}
            ]
        }
    ]
}

RELATIVE_PATH = '/etc/cloud/cloud.cfg'

def test_cloud_cfg():
    cloud_init_file = Mock()
    cloud_init_file.content = CLOUD_CFG.splitlines()
    broker = {LocalSpecs.cloud_cfg_input: cloud_init_file}
    result = cloud_cfg(broker)
    assert result is not None
    assert isinstance(result, DatasourceProvider)
    expected = DatasourceProvider(content=json.dumps(CLOUD_CFG_JSON), relative_path=RELATIVE_PATH)
    assert result.content == expected.content
    assert result.relative_path == expected.relative_path

def test_cloud_cfg_bad():
    cloud_init_file = Mock()
    cloud_init_file.content = CLOUD_CFG_BAD.splitlines()
    broker = {LocalSpecs.cloud_cfg_input: cloud_init_file}
    with pytest.raises(SkipComponent) as e:
        cloud_cfg(broker)
    assert 'Unexpected exception' in str(e)

def test_cloud_cfg_no_network():
    cloud_init_file = Mock()
    cloud_init_file.content = CLOUD_CFG_NO_NETWORK.splitlines()
    broker = {LocalSpecs.cloud_cfg_input: cloud_init_file}
    with pytest.raises(SkipComponent) as e:
        cloud_cfg(broker)
    assert 'No network section in yaml' in str(e)
<|endoftext|>


<commit_before>from .devicebase import DeviceBase
from insteonplm.statechangesignal import StateChangeSignal
from insteonplm.constants import *

class DimmableLightingControl(DeviceBase):
    """Dimmable Lighting Controller 0x01

    INSTEON On/Off switch device class. Available device control options are:
        - light_on(onlevel=0xff)
        - light_on_fast(onlevel=0xff)
        - light_off()
        - light_off_fast()

    To monitor changes to the state of the device subscribe to the state monitor:
         - lightOnLevel.connect(callback)  (state='LightOnLevel')

    where callback defined as:
        - callback(self, device_id, state, state_value)
    """

    def __init__(self, plm, address, cat, subcat, product_key=None, description=None, model=None, groupbutton=0x01):
        DeviceBase.__init__(self, plm, address, cat, subcat, product_key, description, model, groupbutton)

        # Setting the default value of the light to 0 (i.e. Off)
        self.lightOnLevel = StateChangeSignal('LightOnLevel', self.light_status_request, 0x00)

        self._nextCommandIsStatus = False

        self._message_callbacks.add_message_callback(MESSAGE_STANDARD_MESSAGE_RECEIVED_0X50, COMMAND_LIGHT_ON_0X11_NONE, self._light_on_command_received)
        self._message_callbacks.add_message_callback(MESSAGE_STANDARD_MESSAGE_RECEIVED_0X50, COMMAND_LIGHT_OFF_0X13_0X00, self._light_off_command_received)
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, COMMAND_LIGHT_ON_0X11_NONE, self._light_on_command_received, MESSAGE_ACK)
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, COMMAND_LIGHT_OFF_0X13_0X00, self._light_off_command_received, MESSAGE_ACK)
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00, self._light_status_request_ack, MESSAGE_ACK)

    def light_on(self, onlevel=0xff):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_ON_0X11_NONE, onlevel)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_ON_0X11_NONE, onlevel, **userdata)

    def light_on_fast(self, onlevel=0xff):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_ON_FAST_0X12_NONE, onlevel)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_ON_FAST_0X12_NONE, onlevel, **userdata)

    def light_off(self):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_OFF_0X13_0X00)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_OFF_0X13_0X00, **userdata)

    def light_off_fast(self):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_OFF_FAST_0X14_0X00)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_OFF_FAST_0X14_0X00, **userdata)

    def light_status_request(self):
        self._plm.send_standard(self.address.hex, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00)

    def receive_message(self, msg):
        self.log.debug('Starting DimmableLightingControl.receive_message')
        self.log.debug('Next command is status: %r', self._nextCommandIsStatus)
        if self._nextCommandIsStatus:
            self._status_update_received(msg)
        else:
            super().receive_message(msg)
        self.log.debug('Ending DimmableLightingControl.receive_message')

    def get_operating_flags(self):
        return NotImplemented

    def set_operating_flags(self):
        return NotImplemented

    def light_manually_turned_off(self):
        return NotImplemented

    def light_manually_turned_On(self):
        return NotImplemented

    def _light_on_command_received(self, msg):
        self.log.debug('Starting _light_on_command_received')
        # Message handler for Standard (0x50) or Extended (0x51) message commands 0x11 Light On
        # Also handles Standard or Extended (0x62) Lights On (0x11) ACK
        # When any of these messages are received any state listeners are updated with the
        # current light on level (cmd2)
        if msg.code == MESSAGE_EXTENDED_MESSAGE_RECEIVED_0X51 or
          (msg.code == MESSAGE_SEND_STANDARD_MESSAGE_0X62 and msg.isextendedflag):
              group = msg.userdata[0]
            if group == self._groupbutton:
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
            else:
                device2 = self._plm.devices[self._get_device_id(0x02)]
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        else:
            self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending _light_on_command_received')

    def _light_off_command_received(self, msg):
        self.log.debug('Starting _light_off_command_received')
        if msg.code == MESSAGE_EXTENDED_MESSAGE_RECEIVED_0X51 or
          (msg.code == MESSAGE_SEND_STANDARD_MESSAGE_0X62 and msg.isextendedflag):
              group = msg.userdata[0]
            if group == self._groupbutton:
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, 0)
            else:
                device2 = self._plm.devices[self._get_device_id(0x02)]
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, 0)
        else:
            self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending _light_off_command_received')

    def _light_status_request_ack(self, msg):
        self.log.debug('Starting _light_status_request_ack')
        self._nextCommandIsStatus = True
        self.log.debug('Ending _light_status_request_ack')

    def _status_update_received(self, msg):
        self.log.debug('Starting _status_update_received')
        self._nextCommandIsStatus = False
        self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending _status_update_received')

class DimmableLightingControl_2475F(DimmableLightingControl):

    """FanLinc model 2475F Dimmable Lighting Control Device Class 0x01 subcat 0x2e

    Two separate INSTEON On/Off switch devices are created with ID
        - 'address': Top Outlet
        - 'address_2': Bottom Outlet

    INSTEON On/Off switch device class. Available device control options are:
        - light_on(onlevel=0xff)
        - light_on_fast(onlevel=0xff)
        - light_off()
        - light_off_fast()

    To monitor the state of the device subscribe to the state monitor:
         - lightOnLevel.connect(callback)

    where callback defined as:
        - callback(self, device_id, state, state_value)
    """

    def __init__(self, plm, address, cat, subcat, product_key=None, description=None, model=None, groupbutton=0x01):
        super().__init__(plm, address, cat, subcat, product_key, description, model, groupbutton)

        self._nextCommandIsFanStatus = False

        # 2475F has a custom COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00 where cmd1:0x19 and cmd2:0x03 to get the fan status
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, {'cmd1':0x19, 'cmd2':0x03},
                                                     self._fan_status_request_ack, MESSAGE_ACK)

    @classmethod
    def create(cls, plm, address, cat, subcat, product_key=None, description=None, model=None, groupbutton = 0x01):
        devices = []
        devices.append(DimmableLightingControl_2475F(plm, address, cat, subcat, product_key, description, model, 0x01))
        devices.append(DimmableLightingControl_2475F(plm, address, cat, subcat, product_key, description, model, 0x02))
        return devices

    def receive_message(self, msg):
        """
        PLM will dispatch commands to the first device in a class. If there are two devices, like the 2475F,
        the device needs to recognize that the message is for a different group than the first group
        and dispatch the message to the correct device.
        """

        self.log.debug('Starting DimmableLightingControl_2475F.receive_message')
        self.log.debug('Next command is fan status: %r', self._nextCommandIsFanStatus)
        if self._nextCommandIsFanStatus:
            self._fan_status_update_received(msg)
        elif msg.code == MESSAGE_EXTENDED_MESSAGE_RECEIVED_0X51:
            # I think byte 0 ('d1') of the extended message is always the group number for 0x01 and 0x02 devices
            if msg.userdata[0] == self._groupbutton:
                super().receive_message(msg)
            else:
                id = self._get_device_id(msg.userdata[0])
                device = self._plm.devices[id]
                if device is not None:
                    device.receive_message(msg)
        else:
            super().receive_message(msg)

        self.log.debug('Ending DimmableLightingControl_2475F.receive_message')

    def light_status_request(self):
        """
        Status request options are set in cmd2 as:
            0x00 Light control
            0x03 Fan control.
        """
        self.log.debug('Starting DimmableLightingControl_2475F.light_status_request')
        if self._groupbutton == 0x01:
            self.log.debug('Sending light status request')
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00)
        else:
            self.log.debug('Sending fan status request')
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00, 0x03)

        self.log.debug('Ending DimmableLightingControl_2475F.light_status_request')

    def _fan_status_request_ack(self, msg):
        self.log.debug('Starting DimmableLightingControl_2475F._fan_status_request_ack')
        self._nextCommandIsFanStatus = True
        self.log.debug('Ending DimmableLightingControl_2475F._fan_status_request_ack')

    def _fan_status_update_received(self, msg):
        """
        The following status values can be recieve:
            0x00 = Off
            0x01:0x7f = Low
            0x80:0xfe = Med
            0xff = High
        """
        self.log.debug('Starting DimmableLightingControl_2475F._fan_status_update_received')
        device2 = self._plm.devices[self._get_device_id(0x02)]
        self._nextCommandIsFanStatus = False
        device2.lightOnLevel.update(device2.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending DimmableLightingControl_2475F._fan_status_update_received')<commit_msg>Fix syntax error<commit_after>from .devicebase import DeviceBase
from insteonplm.statechangesignal import StateChangeSignal
from insteonplm.constants import *

class DimmableLightingControl(DeviceBase):
    """Dimmable Lighting Controller 0x01

    INSTEON On/Off switch device class. Available device control options are:
        - light_on(onlevel=0xff)
        - light_on_fast(onlevel=0xff)
        - light_off()
        - light_off_fast()

    To monitor changes to the state of the device subscribe to the state monitor:
         - lightOnLevel.connect(callback)  (state='LightOnLevel')

    where callback defined as:
        - callback(self, device_id, state, state_value)
    """

    def __init__(self, plm, address, cat, subcat, product_key=None, description=None, model=None, groupbutton=0x01):
        DeviceBase.__init__(self, plm, address, cat, subcat, product_key, description, model, groupbutton)

        # Setting the default value of the light to 0 (i.e. Off)
        self.lightOnLevel = StateChangeSignal('LightOnLevel', self.light_status_request, 0x00)

        self._nextCommandIsStatus = False

        self._message_callbacks.add_message_callback(MESSAGE_STANDARD_MESSAGE_RECEIVED_0X50, COMMAND_LIGHT_ON_0X11_NONE, self._light_on_command_received)
        self._message_callbacks.add_message_callback(MESSAGE_STANDARD_MESSAGE_RECEIVED_0X50, COMMAND_LIGHT_OFF_0X13_0X00, self._light_off_command_received)
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, COMMAND_LIGHT_ON_0X11_NONE, self._light_on_command_received, MESSAGE_ACK)
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, COMMAND_LIGHT_OFF_0X13_0X00, self._light_off_command_received, MESSAGE_ACK)
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00, self._light_status_request_ack, MESSAGE_ACK)

    def light_on(self, onlevel=0xff):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_ON_0X11_NONE, onlevel)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_ON_0X11_NONE, onlevel, **userdata)

    def light_on_fast(self, onlevel=0xff):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_ON_FAST_0X12_NONE, onlevel)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_ON_FAST_0X12_NONE, onlevel, **userdata)

    def light_off(self):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_OFF_0X13_0X00)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_OFF_0X13_0X00, **userdata)

    def light_off_fast(self):
        if self._groupbutton == 0x01:
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_OFF_FAST_0X14_0X00)
        else:
            userdata = {'d1':self._groupbutton}
            self._plm.send_extended(self._address.hex, COMMAND_LIGHT_OFF_FAST_0X14_0X00, **userdata)

    def light_status_request(self):
        self._plm.send_standard(self.address.hex, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00)

    def receive_message(self, msg):
        self.log.debug('Starting DimmableLightingControl.receive_message')
        self.log.debug('Next command is status: %r', self._nextCommandIsStatus)
        if self._nextCommandIsStatus:
            self._status_update_received(msg)
        else:
            super().receive_message(msg)
        self.log.debug('Ending DimmableLightingControl.receive_message')

    def get_operating_flags(self):
        return NotImplemented

    def set_operating_flags(self):
        return NotImplemented

    def light_manually_turned_off(self):
        return NotImplemented

    def light_manually_turned_On(self):
        return NotImplemented

    def _light_on_command_received(self, msg):
        self.log.debug('Starting _light_on_command_received')
        # Message handler for Standard (0x50) or Extended (0x51) message commands 0x11 Light On
        # Also handles Standard or Extended (0x62) Lights On (0x11) ACK
        # When any of these messages are received any state listeners are updated with the
        # current light on level (cmd2)
        if msg.code == MESSAGE_EXTENDED_MESSAGE_RECEIVED_0X51 or \
          (msg.code == MESSAGE_SEND_STANDARD_MESSAGE_0X62 and msg.isextendedflag):
              group = msg.userdata[0]
            if group == self._groupbutton:
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
            else:
                device2 = self._plm.devices[self._get_device_id(0x02)]
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        else:
            self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending _light_on_command_received')

    def _light_off_command_received(self, msg):
        self.log.debug('Starting _light_off_command_received')
        if msg.code == MESSAGE_EXTENDED_MESSAGE_RECEIVED_0X51 or \
          (msg.code == MESSAGE_SEND_STANDARD_MESSAGE_0X62 and msg.isextendedflag):
              group = msg.userdata[0]
            if group == self._groupbutton:
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, 0)
            else:
                device2 = self._plm.devices[self._get_device_id(0x02)]
                self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, 0)
        else:
            self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending _light_off_command_received')

    def _light_status_request_ack(self, msg):
        self.log.debug('Starting _light_status_request_ack')
        self._nextCommandIsStatus = True
        self.log.debug('Ending _light_status_request_ack')

    def _status_update_received(self, msg):
        self.log.debug('Starting _status_update_received')
        self._nextCommandIsStatus = False
        self.lightOnLevel.update(self.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending _status_update_received')

class DimmableLightingControl_2475F(DimmableLightingControl):

    """FanLinc model 2475F Dimmable Lighting Control Device Class 0x01 subcat 0x2e

    Two separate INSTEON On/Off switch devices are created with ID
        - 'address': Top Outlet
        - 'address_2': Bottom Outlet

    INSTEON On/Off switch device class. Available device control options are:
        - light_on(onlevel=0xff)
        - light_on_fast(onlevel=0xff)
        - light_off()
        - light_off_fast()

    To monitor the state of the device subscribe to the state monitor:
         - lightOnLevel.connect(callback)

    where callback defined as:
        - callback(self, device_id, state, state_value)
    """

    def __init__(self, plm, address, cat, subcat, product_key=None, description=None, model=None, groupbutton=0x01):
        super().__init__(plm, address, cat, subcat, product_key, description, model, groupbutton)

        self._nextCommandIsFanStatus = False

        # 2475F has a custom COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00 where cmd1:0x19 and cmd2:0x03 to get the fan status
        self._message_callbacks.add_message_callback(MESSAGE_SEND_STANDARD_MESSAGE_0X62, {'cmd1':0x19, 'cmd2':0x03},
                                                     self._fan_status_request_ack, MESSAGE_ACK)

    @classmethod
    def create(cls, plm, address, cat, subcat, product_key=None, description=None, model=None, groupbutton = 0x01):
        devices = []
        devices.append(DimmableLightingControl_2475F(plm, address, cat, subcat, product_key, description, model, 0x01))
        devices.append(DimmableLightingControl_2475F(plm, address, cat, subcat, product_key, description, model, 0x02))
        return devices

    def receive_message(self, msg):
        """
        PLM will dispatch commands to the first device in a class. If there are two devices, like the 2475F,
        the device needs to recognize that the message is for a different group than the first group
        and dispatch the message to the correct device.
        """

        self.log.debug('Starting DimmableLightingControl_2475F.receive_message')
        self.log.debug('Next command is fan status: %r', self._nextCommandIsFanStatus)
        if self._nextCommandIsFanStatus:
            self._fan_status_update_received(msg)
        elif msg.code == MESSAGE_EXTENDED_MESSAGE_RECEIVED_0X51:
            # I think byte 0 ('d1') of the extended message is always the group number for 0x01 and 0x02 devices
            if msg.userdata[0] == self._groupbutton:
                super().receive_message(msg)
            else:
                id = self._get_device_id(msg.userdata[0])
                device = self._plm.devices[id]
                if device is not None:
                    device.receive_message(msg)
        else:
            super().receive_message(msg)

        self.log.debug('Ending DimmableLightingControl_2475F.receive_message')

    def light_status_request(self):
        """
        Status request options are set in cmd2 as:
            0x00 Light control
            0x03 Fan control.
        """
        self.log.debug('Starting DimmableLightingControl_2475F.light_status_request')
        if self._groupbutton == 0x01:
            self.log.debug('Sending light status request')
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00)
        else:
            self.log.debug('Sending fan status request')
            self._plm.send_standard(self.address.hex, COMMAND_LIGHT_STATUS_REQUEST_0X19_0X00, 0x03)

        self.log.debug('Ending DimmableLightingControl_2475F.light_status_request')

    def _fan_status_request_ack(self, msg):
        self.log.debug('Starting DimmableLightingControl_2475F._fan_status_request_ack')
        self._nextCommandIsFanStatus = True
        self.log.debug('Ending DimmableLightingControl_2475F._fan_status_request_ack')

    def _fan_status_update_received(self, msg):
        """
        The following status values can be recieve:
            0x00 = Off
            0x01:0x7f = Low
            0x80:0xfe = Med
            0xff = High
        """
        self.log.debug('Starting DimmableLightingControl_2475F._fan_status_update_received')
        device2 = self._plm.devices[self._get_device_id(0x02)]
        self._nextCommandIsFanStatus = False
        device2.lightOnLevel.update(device2.id, self.lightOnLevel._stateName, msg.cmd2)
        self.log.debug('Ending DimmableLightingControl_2475F._fan_status_update_received')<|endoftext|>


<commit_before>''' Extractor classes based on pre-trained models. '''

import numpy as np
from abc import ABCMeta

from pliers.extractors.image import ImageExtractor
from pliers.extractors.base import Extractor, ExtractorResult
from pliers.filters.image import ImageResizingFilter
from pliers.utils import (attempt_to_import, verify_dependencies,
                         listify)

tf = attempt_to_import('tensorflow')
attempt_to_import('tensorflow.keras')
hub = attempt_to_import('tensorflow_hub')

class TFHubExtractor(Extractor, metaclass=ABCMeta):

    ''' A generic class for Tensorflow Hub extractors
    Args:
        url (str): url of the tensorflow-hub model to download
        task (str): model task/domain identifier
        labels (optional): list of labels (if relevant, e.g. for
            classification models). If specified, labels are used as
            feature names and override task argument.
        transform_fn (optional): function applied to the input before
            feeding it to the model (e.g. to change TF signature).
        kwargs (dict): arguments to TensorFlow Hub load method (see
            https://www.tensorflow.org/hub/api_docs/python/hub/load)'''

    _log_attributes = ('url',)

    def __init__(self, url, task=None, labels=None,
                transform_fn=None, **kwargs):
        verify_dependencies(['tensorflow'])
        self.model = hub.load(url, **kwargs)
        self._labels = labels
        self._task = task
        self.transform_fn = transform_fn
        super().__init__()

    def get_feature_names(self):
        if self._labels:
            return self._labels
        else:
            return listify(self._task)

    def _extract(self, stim):
        features = listify(self.get_feature_names())
        input = listify(stim.data)
        if self.transform_fn:
            input = self.transform_fn(input)
        output = self.model.signatures['serving_default'](input).numpy()
        return ExtractorResult(output, stim, self,
                               features=features)

class TFHubEmbeddingExtractor(TFHubExtractor):

    ''' Extracts embedding from TF Hub embedding models '''

    _task = 'embedding'

class TFHubClassificationExtractor(TFHubExtractor):

    ''' Extracts logits for TF Hub classification models '''

    _task = 'logits'

class TensorFlowKerasApplicationExtractor(ImageExtractor):

    ''' Labels objects in images using a pretrained Inception V3 architecture
    implemented in TensorFlow / Keras.

    Images must be RGB and be a certain shape. Different model architectures
    may require different shapes, and images will be resized (with some
    distortion) if the shape of the image is different.

    Args:
        architecture (str): model architecture to use. One of 'vgg16', 'vgg19',
            'resnet50', 'inception_resnetv2', 'inceptionv3', 'xception',
            'densenet121', 'densenet169', 'nasnetlarge', or 'nasnetmobile'.
        weights (str): URL to download pre-trained weights. If None (default),
            uses the pre-trained weights trained on ImageNet used in Keras
            Applications.
        num_predictions (int): Number of top predicted labels to retain for
            each image.
     '''

    _log_attributes = ('architecture', 'weights', 'num_predictions')
    VERSION = '1.0'

    def __init__(self,
                 architecture='inceptionv3',
                 weights=None,
                 num_predictions=5):
        verify_dependencies(['tensorflow'])
        verify_dependencies(['tensorflow.keras'])
        super().__init__()

        # Model name: (model module, model function, required shape).
        apps = tf.keras.applications
        model_mapping = {
            'vgg16': (apps.vgg16, apps.vgg16.VGG16, (224, 224, 3)),
            'vgg19': (apps.vgg19, apps.VGG19, (224, 224, 3)),
            'resnet50': (apps.resnet50, apps.resnet50.ResNet50, (224, 224, 3)),
            'inception_resnetv2': (
                apps.inception_resnet_v2,
                apps.inception_resnet_v2.InceptionResNetV2, (299, 299, 3)),
            'inceptionv3': (
                apps.inception_v3, apps.inception_v3.InceptionV3,
                (299, 299, 3)),
            'xception': (apps.xception, apps.xception.Xception, (299, 299, 3)),
            'densenet121': (
                apps.densenet, apps.densenet.DenseNet121, (224, 224, 3)),
            'densenet169': (
                apps.densenet, apps.densenet.DenseNet169, (224, 224, 3)),
            'densenet201': (
                apps.densenet, apps.densenet.DenseNet201, (224, 224, 3)),
            'nasnetlarge': (
                apps.nasnet, apps.nasnet.NASNetLarge, (331, 331, 3)),
            'nasnetmobile': (
                apps.nasnet, apps.nasnet.NASNetMobile, (224, 224, 3)),
        }
        if weights is None:
            weights = 'imagenet'
        if architecture.lower() not in model_mapping.keys():
            raise ValueError(
                "Unknown architecture '{}'. Available arhitectures are '{}'."
                .format(architecture, "', '".join(model_mapping.keys())))

        self.architecture = architecture.lower()
        self.weights = weights
        self.num_predictions = num_predictions

        # The preprocessing and decoding functions are in the module.
        self._model_module = model_mapping[self.architecture][0]
        # Instantiating the model also downloads the weights to a cache.
        self.model = model_mapping[self.architecture][1](weights=self.weights)
        self._required_shape = model_mapping[self.architecture][2]

    def _extract(self, stim):
        x = stim.data
        if x.ndim != 3:
            raise ValueError(
                "Stim data must have rank 3 but got rank {}".format(x.ndim))
        if x.shape != self._required_shape:
            resizer = ImageResizingFilter(size=self._required_shape[:-1])
            x = resizer.transform(stim).data
        # Add batch dimension.
        x = x[None]
        # Normalize the features.
        x = self._model_module.preprocess_input(x)
        preds = self.model.predict(x, batch_size=1)

        # This produces a nested list. There is one sub-list per sample in the
        # batch, and each sub-list contains `self.num_predictions` tuples with
        # `(ID, human-readable-label, probability)`.
        decoded = self._model_module.decode_predictions(
            preds, top=self.num_predictions)

        # We assume there is only one sample in the batch.
        decoded = decoded[0]
        values = [t[2] for t in decoded]
        features = [t[1] for t in decoded]

        return ExtractorResult([values], stim, self, features=features)
<commit_msg>add comment<commit_after>''' Extractor classes based on pre-trained models. '''

import numpy as np
from abc import ABCMeta

from pliers.extractors.image import ImageExtractor
from pliers.extractors.base import Extractor, ExtractorResult
from pliers.filters.image import ImageResizingFilter
from pliers.utils import (attempt_to_import, verify_dependencies,
                         listify)

tf = attempt_to_import('tensorflow')
attempt_to_import('tensorflow.keras')
hub = attempt_to_import('tensorflow_hub')

class TFHubExtractor(Extractor, metaclass=ABCMeta):

    ''' A generic class for Tensorflow Hub extractors
    Args:
        url (str): url of the tensorflow-hub model to download
        task (str): model task/domain identifier
        labels (optional): list of labels (if relevant, e.g. for
            classification models). If specified, labels are used as
            feature names and override task argument.
        transform_fn (optional): function applied to the input before
            feeding it to the model (e.g. to change TF signature).
        kwargs (dict): arguments to TensorFlow Hub load method (see
            https://www.tensorflow.org/hub/api_docs/python/hub/load)'''

    _log_attributes = ('url',)

    def __init__(self, url, task=None, labels=None,
                transform_fn=None, **kwargs):
        verify_dependencies(['tensorflow'])
        self.model = hub.load(url, **kwargs)
        self._labels = labels
        self._task = task
        self.transform_fn = transform_fn
        super().__init__()

    def get_feature_names(self):
        if self._labels:
            return self._labels
        else:
            return listify(self._task)

    def _extract(self, stim):
        features = listify(self.get_feature_names())
        input = listify(stim.data)
        if self.transform_fn:
            input = self.transform_fn(input)
        output = self.model.signatures['serving_default'](input).numpy() # or self.model(input).numpy()
        return ExtractorResult(output, stim, self,
                               features=features)

class TFHubEmbeddingExtractor(TFHubExtractor):

    ''' Extracts embedding from TF Hub embedding models '''

    _task = 'embedding'

class TFHubClassificationExtractor(TFHubExtractor):

    ''' Extracts logits for TF Hub classification models '''

    _task = 'logits'

class TensorFlowKerasApplicationExtractor(ImageExtractor):

    ''' Labels objects in images using a pretrained Inception V3 architecture
    implemented in TensorFlow / Keras.

    Images must be RGB and be a certain shape. Different model architectures
    may require different shapes, and images will be resized (with some
    distortion) if the shape of the image is different.

    Args:
        architecture (str): model architecture to use. One of 'vgg16', 'vgg19',
            'resnet50', 'inception_resnetv2', 'inceptionv3', 'xception',
            'densenet121', 'densenet169', 'nasnetlarge', or 'nasnetmobile'.
        weights (str): URL to download pre-trained weights. If None (default),
            uses the pre-trained weights trained on ImageNet used in Keras
            Applications.
        num_predictions (int): Number of top predicted labels to retain for
            each image.
     '''

    _log_attributes = ('architecture', 'weights', 'num_predictions')
    VERSION = '1.0'

    def __init__(self,
                 architecture='inceptionv3',
                 weights=None,
                 num_predictions=5):
        verify_dependencies(['tensorflow'])
        verify_dependencies(['tensorflow.keras'])
        super().__init__()

        # Model name: (model module, model function, required shape).
        apps = tf.keras.applications
        model_mapping = {
            'vgg16': (apps.vgg16, apps.vgg16.VGG16, (224, 224, 3)),
            'vgg19': (apps.vgg19, apps.VGG19, (224, 224, 3)),
            'resnet50': (apps.resnet50, apps.resnet50.ResNet50, (224, 224, 3)),
            'inception_resnetv2': (
                apps.inception_resnet_v2,
                apps.inception_resnet_v2.InceptionResNetV2, (299, 299, 3)),
            'inceptionv3': (
                apps.inception_v3, apps.inception_v3.InceptionV3,
                (299, 299, 3)),
            'xception': (apps.xception, apps.xception.Xception, (299, 299, 3)),
            'densenet121': (
                apps.densenet, apps.densenet.DenseNet121, (224, 224, 3)),
            'densenet169': (
                apps.densenet, apps.densenet.DenseNet169, (224, 224, 3)),
            'densenet201': (
                apps.densenet, apps.densenet.DenseNet201, (224, 224, 3)),
            'nasnetlarge': (
                apps.nasnet, apps.nasnet.NASNetLarge, (331, 331, 3)),
            'nasnetmobile': (
                apps.nasnet, apps.nasnet.NASNetMobile, (224, 224, 3)),
        }
        if weights is None:
            weights = 'imagenet'
        if architecture.lower() not in model_mapping.keys():
            raise ValueError(
                "Unknown architecture '{}'. Available arhitectures are '{}'."
                .format(architecture, "', '".join(model_mapping.keys())))

        self.architecture = architecture.lower()
        self.weights = weights
        self.num_predictions = num_predictions

        # The preprocessing and decoding functions are in the module.
        self._model_module = model_mapping[self.architecture][0]
        # Instantiating the model also downloads the weights to a cache.
        self.model = model_mapping[self.architecture][1](weights=self.weights)
        self._required_shape = model_mapping[self.architecture][2]

    def _extract(self, stim):
        x = stim.data
        if x.ndim != 3:
            raise ValueError(
                "Stim data must have rank 3 but got rank {}".format(x.ndim))
        if x.shape != self._required_shape:
            resizer = ImageResizingFilter(size=self._required_shape[:-1])
            x = resizer.transform(stim).data
        # Add batch dimension.
        x = x[None]
        # Normalize the features.
        x = self._model_module.preprocess_input(x)
        preds = self.model.predict(x, batch_size=1)

        # This produces a nested list. There is one sub-list per sample in the
        # batch, and each sub-list contains `self.num_predictions` tuples with
        # `(ID, human-readable-label, probability)`.
        decoded = self._model_module.decode_predictions(
            preds, top=self.num_predictions)

        # We assume there is only one sample in the batch.
        decoded = decoded[0]
        values = [t[2] for t in decoded]
        features = [t[1] for t in decoded]

        return ExtractorResult([values], stim, self, features=features)
<|endoftext|>


<commit_before>"""Vim specific logic"""
import os
import time

import vim

import msg
import shared as G
import utils
import protocol

class View(object):
    """editors representation of the buffer"""

    def __init__(self, vim_buf, buf):
        self.vim_buf = vim_buf
        self.buf = buf

    def __repr__(self):
        return '%s %s %s' % (self.native_id, self.buf['id'], self.buf['path'])

    def __str__(self):
        return repr(self)

    def _offset_to_vim(self, offset):
        current_offset = 0
        for line_num, line in enumerate(self.vim_buf):
            next_offset = len(line) + 1
            if current_offset + next_offset > offset:
                break
            current_offset += next_offset
        col = offset - current_offset
        return line_num + 1, col + 1

    @property
    def native_id(self):
        return self.vim_buf.number

    def is_loading(self):
        return False

    def get_text(self):
        text = '\n'.join(self.vim_buf[:])
        return text.decode('utf-8')

    def set_text(self, text):
        msg.debug('\n\nabout to patch %s %s' % (str(self), self.vim_buf.name))
        try:
            msg.debug("now buf is loadedish? %s" % vim.eval('bufloaded(%s)' % self.native_id))
            self.vim_buf[:] = text.encode('utf-8').split('\n')
        except Exception as e:
            msg.error("couldn't apply patches because: %s!\nThe unencoded text was: %s" % (str(e), text))
            raise

    def apply_patches(self, buf, patches):
        cursor_offset = self.get_cursor_offset()
        msg.debug('cursor offset is %s bytes' % cursor_offset)

        self.set_text(patches[0])

        for patch in patches[2]:
            offset = patch[0]
            length = patch[1]
            patch_text = patch[2]
            if cursor_offset > offset:
                new_offset = len(patch_text) - length
                cursor_offset += new_offset

        self.set_cursor_position(cursor_offset)

    def focus(self):
        vim.command(':edit! %s' % self.vim_buf.name)
        vim.command(':filetype detect')

    def set_cursor_position(self, offset):
        line_num, col = self._offset_to_vim(offset)
        command = 'setpos(".", [%s, %s, %s, %s])' % (self.native_id, line_num, col, 0)
        msg.debug("setting pos: %s" % command)
        # rv = int(vim.eval(command))
        # if rv != 0:
        #     msg.debug('SHIIIIIIIIT %s' % rv)

    def get_cursor_position(self):
        """ [bufnum, lnum, col, off] """
        return vim.eval('getpos(".")')

    def get_cursor_offset(self):
        return int(vim.eval('line2byte(line("."))+col(".")')) - 2

    def get_selections(self):
        cursor = self.get_cursor_offset()
        return [[cursor, cursor]]

    def clear_selections(self):
        msg.debug('clearing selections for view %s' % self.vim_buf.name)

    def highlight(self, ranges, user_id):
        msg.debug('highlighting ranges %s' % (ranges))

    def rename(self, name):
        msg.debug('renaming %s to %s' % (self.vim_buf.name, name))
        current = vim.current.buffer
        text = self.get_text()
        old_name = self.vim_buf.name
        old_number = self.native_id
        with open(name, 'wb') as fd:
            fd.write(text.encode('utf-8'))
        vim.command('edit! %s' % name)
        self.vim_buf = vim.current.buffer
        vim.command('edit! %s' % current.name)
        vim.command('bdelete! %s' % old_number)
        try:
            utils.rm(old_name)
        except Exception as e:
            msg.debug("couldn't delete %s... maybe thats OK?" % str(e))

    def save(self):
        vim.command(':s!')

class Protocol(protocol.BaseProtocol):
    """understands vim"""
    CLIENT = 'VIM'

    def on_room_info(self, room_info):
        super(Protocol, self).on_room_info(room_info)
        vim.command(':Explore %s' % G.PROJECT_PATH)

    def maybe_selection_changed(self, vim_buf, is_ping):
        buf = self.get_buf(vim_buf)
        if not buf:
            msg.debug('no buffer found for view %s' % vim_buf.number)
            return
        view = self.get_view(buf['id'])
        msg.debug("selection changed: %s %s %s" % (vim_buf.number, buf['id'], view))
        self.SELECTION_CHANGED.append([view, is_ping])

    def maybe_buffer_changed(self, vim_buf):
        text = vim_buf[:]
        buf = self.get_buf(vim_buf)
        if not buf:
            return
        if buf['buf'] != text:
            self.BUFS_CHANGED.append(buf['id'])

    def get_vim_buf_by_path(self, p):
        for vim_buf in vim.buffers:
            if vim_buf.name and p == utils.to_rel_path(vim_buf.name):
                return vim_buf
        return None

    def get_view(self, buf_id):
        buf = self.FLOO_BUFS.get(buf_id)
        if not buf:
            return None

        vb = self.get_vim_buf_by_path(buf['path'])
        if not vb:
            return None

        if vim.eval('bufloaded(%s)' % vb.number) == '0':
            return None

        return View(vb, buf)

    def create_view(self, buf):
        path = self.save_buf(buf)
        vb = self.get_vim_buf_by_path(buf['path'])
        if vb:
            return View(vb, buf)

        vim.command(':edit! %s' % path)
        vb = self.get_vim_buf_by_path(buf['path'])
        if vb is None:
            msg.debug('vim buffer is none even though we tried to open it: %s' % path)
            return
        return View(vb, buf)

    def get_buf(self, vim_buf):
        """None- no sharing, False- should be but isn't """
        if vim_buf.name is None:
            msg.debug('get:buf buffer has no filename')
            return None

        if not self.is_shared(vim_buf.name):
            msg.debug('get_buf: %s is not shared' % vim_buf.name)
            return None

        rel_path = utils.to_rel_path(vim_buf.name)
        for buf_id, buf in self.FLOO_BUFS.iteritems():
            if rel_path == buf['path']:
                return buf

        msg.debug('get_buf: no buf has path %s' % rel_path)
        return False

    def save_buf(self, buf):
        path = utils.get_full_path(buf['path'])
        utils.mkdir(os.path.split(path)[0])
        with open(path, 'wb') as fd:
            fd.write(buf['buf'].encode('utf-8'))
        return path

    def chat(self, username, timestamp, message, self_msg=False):
        pass
        # envelope = msg.MSG(message, timestamp, username)
        # if not self_msg:
        #     self.chat_deck.appendleft(envelope)
        # envelope.display()

    def on_msg(self, data):
        timestamp = data.get('time') or time.time()
        msg.log('[%s] <%s> %s' % (time.ctime(timestamp), data.get('username', ''), data.get('data', '')))

    def update_view(self, buf, view=None):
        msg.debug('updating view for buf %s' % buf['id'])
        view = view or self.get_view(buf['id'])
        if not view:
            msg.log('view for buf %s not found. not updating' % buf['id'])
            return

        # visible_region = view.visible_region()
        # viewport_position = view.viewport_position()
        # region = sublime.Region(0, view.size())
        # # deep copy
        # selections = [x for x in view.sel()]
        self.MODIFIED_EVENTS.put(1)
        view.set_text(buf['buf'])
        # try:
        #     edit = view.begin_edit()
        #     view.replace(edit, region, buf['buf'])
        # except Exception as e:
        #     msg.error('Exception updating view: %s' % e)
        # finally:
        #     view.end_edit(edit)
        # sublime.set_timeout(lambda: view.set_viewport_position(viewport_position, False), 0)
        # view.sel().clear()
        # view.show(visible_region, False)
        # for sel in selections:
        #     view.sel().add(sel)
        # if 'patch' in G.PERMS:
        #     view.set_read_only(False)
        # else:
        #     view.set_status('Floobits', 'You don\'t have write permission. Buffer is read-only.')
        #     view.set_read_only(True)
<commit_msg>now trying to set the cursor again, but not on inactive buffers<commit_after>"""Vim specific logic"""
import os
import time

import vim

import msg
import shared as G
import utils
import protocol

class View(object):
    """editors representation of the buffer"""

    def __init__(self, vim_buf, buf):
        self.vim_buf = vim_buf
        self.buf = buf

    def __repr__(self):
        return '%s %s %s' % (self.native_id, self.buf['id'], self.buf['path'])

    def __str__(self):
        return repr(self)

    def _offset_to_vim(self, offset):
        current_offset = 0
        for line_num, line in enumerate(self.vim_buf):
            next_offset = len(line) + 1
            if current_offset + next_offset > offset:
                break
            current_offset += next_offset
        col = offset - current_offset
        return line_num + 1, col + 1

    @property
    def native_id(self):
        return self.vim_buf.number

    def is_loading(self):
        return False

    def get_text(self):
        text = '\n'.join(self.vim_buf[:])
        return text.decode('utf-8')

    def set_text(self, text):
        msg.debug('\n\nabout to patch %s %s' % (str(self), self.vim_buf.name))
        try:
            msg.debug("now buf is loadedish? %s" % vim.eval('bufloaded(%s)' % self.native_id))
            self.vim_buf[:] = text.encode('utf-8').split('\n')
        except Exception as e:
            msg.error("couldn't apply patches because: %s!\nThe unencoded text was: %s" % (str(e), text))
            raise

    def apply_patches(self, buf, patches):
        cursor_offset = self.get_cursor_offset()
        msg.debug('cursor offset is %s bytes' % cursor_offset)

        self.set_text(patches[0])

        for patch in patches[2]:
            offset = patch[0]
            length = patch[1]
            patch_text = patch[2]
            if cursor_offset > offset:
                new_offset = len(patch_text) - length
                cursor_offset += new_offset

        self.set_cursor_position(cursor_offset)

    def focus(self):
        vim.command(':edit! %s' % self.vim_buf.name)
        vim.command(':filetype detect')

    def set_cursor_position(self, offset):
        line_num, col = self._offset_to_vim(offset)
        command = 'setpos(".", [%s, %s, %s, %s])' % (self.native_id, line_num, col, 0)
        msg.debug("setting pos: %s" % command)
        rv = int(vim.eval(command))
        if rv != 0:
            msg.debug('SHIIIIIIIIT %s' % rv)

    def get_cursor_position(self):
        """ [bufnum, lnum, col, off] """
        return vim.eval('getpos(".")')

    def get_cursor_offset(self):
        return int(vim.eval('line2byte(line("."))+col(".")')) - 2

    def get_selections(self):
        cursor = self.get_cursor_offset()
        return [[cursor, cursor]]

    def clear_selections(self):
        msg.debug('clearing selections for view %s' % self.vim_buf.name)

    def highlight(self, ranges, user_id):
        msg.debug('highlighting ranges %s' % (ranges))

    def rename(self, name):
        msg.debug('renaming %s to %s' % (self.vim_buf.name, name))
        current = vim.current.buffer
        text = self.get_text()
        old_name = self.vim_buf.name
        old_number = self.native_id
        with open(name, 'wb') as fd:
            fd.write(text.encode('utf-8'))
        vim.command('edit! %s' % name)
        self.vim_buf = vim.current.buffer
        vim.command('edit! %s' % current.name)
        vim.command('bdelete! %s' % old_number)
        try:
            utils.rm(old_name)
        except Exception as e:
            msg.debug("couldn't delete %s... maybe thats OK?" % str(e))

    def save(self):
        vim.command(':s!')

class Protocol(protocol.BaseProtocol):
    """understands vim"""
    CLIENT = 'VIM'

    def on_room_info(self, room_info):
        super(Protocol, self).on_room_info(room_info)
        vim.command(':Explore %s' % G.PROJECT_PATH)

    def maybe_selection_changed(self, vim_buf, is_ping):
        buf = self.get_buf(vim_buf)
        if not buf:
            msg.debug('no buffer found for view %s' % vim_buf.number)
            return
        view = self.get_view(buf['id'])
        msg.debug("selection changed: %s %s %s" % (vim_buf.number, buf['id'], view))
        self.SELECTION_CHANGED.append([view, is_ping])

    def maybe_buffer_changed(self, vim_buf):
        text = vim_buf[:]
        buf = self.get_buf(vim_buf)
        if not buf:
            return
        if buf['buf'] != text:
            self.BUFS_CHANGED.append(buf['id'])

    def get_vim_buf_by_path(self, p):
        for vim_buf in vim.buffers:
            if vim_buf.name and p == utils.to_rel_path(vim_buf.name):
                return vim_buf
        return None

    def get_view(self, buf_id):
        buf = self.FLOO_BUFS.get(buf_id)
        if not buf:
            return None

        vb = self.get_vim_buf_by_path(buf['path'])
        if not vb:
            return None

        if vim.eval('bufloaded(%s)' % vb.number) == '0':
            return None

        return View(vb, buf)

    def create_view(self, buf):
        path = self.save_buf(buf)
        vb = self.get_vim_buf_by_path(buf['path'])
        if vb:
            return View(vb, buf)

        vim.command(':edit! %s' % path)
        vb = self.get_vim_buf_by_path(buf['path'])
        if vb is None:
            msg.debug('vim buffer is none even though we tried to open it: %s' % path)
            return
        return View(vb, buf)

    def get_buf(self, vim_buf):
        """None- no sharing, False- should be but isn't """
        if vim_buf.name is None:
            msg.debug('get:buf buffer has no filename')
            return None

        if not self.is_shared(vim_buf.name):
            msg.debug('get_buf: %s is not shared' % vim_buf.name)
            return None

        rel_path = utils.to_rel_path(vim_buf.name)
        for buf_id, buf in self.FLOO_BUFS.iteritems():
            if rel_path == buf['path']:
                return buf

        msg.debug('get_buf: no buf has path %s' % rel_path)
        return False

    def save_buf(self, buf):
        path = utils.get_full_path(buf['path'])
        utils.mkdir(os.path.split(path)[0])
        with open(path, 'wb') as fd:
            fd.write(buf['buf'].encode('utf-8'))
        return path

    def chat(self, username, timestamp, message, self_msg=False):
        pass
        # envelope = msg.MSG(message, timestamp, username)
        # if not self_msg:
        #     self.chat_deck.appendleft(envelope)
        # envelope.display()

    def on_msg(self, data):
        timestamp = data.get('time') or time.time()
        msg.log('[%s] <%s> %s' % (time.ctime(timestamp), data.get('username', ''), data.get('data', '')))

    def update_view(self, buf, view=None):
        msg.debug('updating view for buf %s' % buf['id'])
        view = view or self.get_view(buf['id'])
        if not view:
            msg.log('view for buf %s not found. not updating' % buf['id'])
            return

        # visible_region = view.visible_region()
        # viewport_position = view.viewport_position()
        # region = sublime.Region(0, view.size())
        # # deep copy
        # selections = [x for x in view.sel()]
        self.MODIFIED_EVENTS.put(1)
        view.set_text(buf['buf'])
        # try:
        #     edit = view.begin_edit()
        #     view.replace(edit, region, buf['buf'])
        # except Exception as e:
        #     msg.error('Exception updating view: %s' % e)
        # finally:
        #     view.end_edit(edit)
        # sublime.set_timeout(lambda: view.set_viewport_position(viewport_position, False), 0)
        # view.sel().clear()
        # view.show(visible_region, False)
        # for sel in selections:
        #     view.sel().add(sel)
        # if 'patch' in G.PERMS:
        #     view.set_read_only(False)
        # else:
        #     view.set_status('Floobits', 'You don\'t have write permission. Buffer is read-only.')
        #     view.set_read_only(True)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#-------------------------------------------------------------------------------
# Name:         PluginHeartbleed.py
# Purpose:      Tests the target server for CVE-2014-0160.
#
# Author:       alban
#
# Copyright:    2014 SSLyze developers
#
#   SSLyze 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 2 of the License, or
#   (at your option) any later version.
#
#   SSLyze 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 SSLyze.  If not, see <http://www.gnu.org/licenses/>.
#-------------------------------------------------------------------------------

import socket, new
from xml.etree.ElementTree import Element

from plugins import PluginBase
from utils.SSLyzeSSLConnection import create_sslyze_connection, SSLHandshakeRejected
from nassl._nassl import OpenSSLError, WantX509LookupError, WantReadError
from nassl import TLSV1, TLSV1_1, TLSV1_2, SSLV23, SSLV3

class PluginHeartbleed(PluginBase.PluginBase):

    interface = PluginBase.PluginInterface("PluginHeartbleed",  "")
    interface.add_command(
        command="heartbleed",
        help=(
            "Tests the server(s) for the OpenSSL Heartbleed vulnerability."))

    def process_task(self, target, command, args):

        (host, ip, port, sslVersion) = target

        if sslVersion == SSLV23: # Could not determine the preferred  SSL version - client cert was required ?
            sslVersion = TLSV1 # Default to TLS 1.0

        sslConn = create_sslyze_connection(target, self._shared_settings)
        sslConn.sslVersion = sslVersion # Needed by the heartbleed payload

        # Awful hack #1: replace nassl.sslClient.do_handshake() with a heartbleed
        # checking SSL handshake so that all the SSLyze options
        # (startTLS, proxy, etc.) still work
        sslConn.do_handshake = new.instancemethod(do_handshake_with_heartbleed, sslConn, None)

        try: # Perform the SSL handshake
            sslConn.connect()
        except HeartbleedSent:
            # Awful hack #2: directly read the underlying network socket
            heartbleed = sslConn._sock.recv(16381)
            print len(heartbleed)
            print repr(heartbleed)
        finally:
            sslConn.close()

        # Text output
        if heartbleed:
            heartbleedTxt = 'VULNERABLE'
        else:
            heartbleedTxt = 'NOT vulnerable'

        cmdTitle = 'Heartbleed'
        txtOutput = [self.PLUGIN_TITLE_FORMAT(cmdTitle)]
        txtOutput.append(OUT_FORMAT("OpenSSL Heartbleed:", heartbleedTxt))

        # XML output
        xmlOutput = Element(command, title=cmdTitle)
        if heartbleed:
            xmlNode = Element('heartbleed', type="DEFLATE")
            xmlOutput.append(xmlNode)

        return PluginBase.PluginResult(txtOutput, xmlOutput)

def heartbleed_payload(sslVersion):
    # This heartbleed payload does not exploit the server
    # https://blog.mozilla.org/security/2014/04/12/testing-for-heartbleed-vulnerability-without-exploiting-the-server/

    SSL_VERSION_MAPPING = {
        SSLV3 :  '\x00', # Surprising that it works with SSL 3 which doesn't define TLS extensions
        TLSV1 :  '\x01',
        TLSV1_1: '\x02',
        TLSV1_2: '\x03'}

    payload = ('\x18'           # Record type - Heartbeat
        '\x03{0}'               # TLS version
        '\x00\x03'              # Record length
        '\x01'                  # Heartbeat type - Request
        '\x00\x00')             # Heartbeat length

    payload = ('\x18'           # Record type - Heartbeat
        '\x03{0}'               # TLS version
        '\x40\x00'              # Record length
        '\x01'                  # Heartbeat type - Request
        '\x3f\xfd')             # Heartbeat length

    payload += '\x01'*16381     # Heartbeat data

    payload += (                # Second Heartbeat request with no padding
        '\x18'                  # Record type - Heartbeat
        '\x03{0}'
        '\x00\x03\x01\x00\x00'
    )

    return payload.format(SSL_VERSION_MAPPING[sslVersion])

class HeartbleedSent(SSLHandshakeRejected):
    # Awful hack #3: Use an exception to hack the handshake's control flow in
    # a super obscure way
    pass

def do_handshake_with_heartbleed(self):
    # This is nassl's code for do_handshake() modified to send a heartbleed
    # payload that will reveal 1 byte of the server's memory
    # I copied nassl's code here so I could leave anything heartbleed-related
    # outside of the nassl code base
    try:
        if self._ssl.do_handshake() == 1:
            self._handshakeDone = True
            return True # Handshake was successful

    except WantReadError:
        # OpenSSL is expecting more data from the peer
        # Send available handshake data to the peer
        # In this heartbleed handshake we only send the client hello
        lenToRead = self._networkBio.pending()
        while lenToRead:
            # Get the data from the SSL engine
            handshakeDataOut = self._networkBio.read(lenToRead)
            # Send it to the peer
            self._sock.send(handshakeDataOut)
            lenToRead = self._networkBio.pending()

        # Send the heartbleed payload after the client hello
        self._sock.send(heartbleed_payload(self.sslVersion))

        # Recover the peer's encrypted response
        # In this heartbleed handshake we only receive the server hello
        handshakeDataIn = self._sock.recv(2048)
        if len(handshakeDataIn) == 0:
            raise IOError('Nassl SSL handshake failed: peer did not send data back.')
        # Pass the data to the SSL engine
        self._networkBio.write(handshakeDataIn)

        # Signal that we sent the heartbleed payload and just stop the handshake
        raise HeartbleedSent("")

    except WantX509LookupError:
        # Server asked for a client certificate and we didn't provide one
        # Heartbleed should work anyway
        self._sock.send(heartbleed_payload(self.sslVersion)) # The heartbleed payload
        raise HeartbleedSent("") # Signal that we sent the heartbleed payload

<commit_msg>Fix heartbleed txt output<commit_after>#!/usr/bin/env python
#-------------------------------------------------------------------------------
# Name:         PluginHeartbleed.py
# Purpose:      Tests the target server for CVE-2014-0160.
#
# Author:       alban
#
# Copyright:    2014 SSLyze developers
#
#   SSLyze 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 2 of the License, or
#   (at your option) any later version.
#
#   SSLyze 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 SSLyze.  If not, see <http://www.gnu.org/licenses/>.
#-------------------------------------------------------------------------------

import socket, new
from xml.etree.ElementTree import Element

from plugins import PluginBase
from utils.SSLyzeSSLConnection import create_sslyze_connection, SSLHandshakeRejected
from nassl._nassl import OpenSSLError, WantX509LookupError, WantReadError
from nassl import TLSV1, TLSV1_1, TLSV1_2, SSLV23, SSLV3

class PluginHeartbleed(PluginBase.PluginBase):

    interface = PluginBase.PluginInterface("PluginHeartbleed",  "")
    interface.add_command(
        command="heartbleed",
        help=(
            "Tests the server(s) for the OpenSSL Heartbleed vulnerability."))

    def process_task(self, target, command, args):

        OUT_FORMAT = '      {0:<35}{1}'.format
        (host, ip, port, sslVersion) = target

        if sslVersion == SSLV23: # Could not determine the preferred  SSL version - client cert was required ?
            sslVersion = TLSV1 # Default to TLS 1.0

        sslConn = create_sslyze_connection(target, self._shared_settings)
        sslConn.sslVersion = sslVersion # Needed by the heartbleed payload

        # Awful hack #1: replace nassl.sslClient.do_handshake() with a heartbleed
        # checking SSL handshake so that all the SSLyze options
        # (startTLS, proxy, etc.) still work
        sslConn.do_handshake = new.instancemethod(do_handshake_with_heartbleed, sslConn, None)

        heartbleed = None
        try: # Perform the SSL handshake
            sslConn.connect()
        except HeartbleedSent:
            # Awful hack #2: directly read the underlying network socket
            heartbleed = sslConn._sock.recv(16381)
            print len(heartbleed)
            print repr(heartbleed)
        finally:
            sslConn.close()

        # Text output
        if heartbleed is None:
            heartbleedTxt = 'Error'
        elif '\x01\x01\x01\x01\x01\x01\x01\x01\x01' in heartbleed:
            heartbleedTxt = 'VULNERABLE'
        else:
            heartbleedTxt = 'NOT vulnerable'

        print heartbleedTxt

        cmdTitle = 'Heartbleed'
        txtOutput = [self.PLUGIN_TITLE_FORMAT(cmdTitle)]
        txtOutput.append(OUT_FORMAT("OpenSSL Heartbleed:", heartbleedTxt))

        # XML output
        xmlOutput = Element(command, title=cmdTitle)
        if heartbleed:
            xmlNode = Element('heartbleed', type="DEFLATE") #TBD
            xmlOutput.append(xmlNode)

        return PluginBase.PluginResult(txtOutput, xmlOutput)

def heartbleed_payload(sslVersion):
    # This heartbleed payload does not exploit the server
    # https://blog.mozilla.org/security/2014/04/12/testing-for-heartbleed-vulnerability-without-exploiting-the-server/

    SSL_VERSION_MAPPING = {
        SSLV3 :  '\x00', # Surprising that it works with SSL 3 which doesn't define TLS extensions
        TLSV1 :  '\x01',
        TLSV1_1: '\x02',
        TLSV1_2: '\x03'}

    payload = ('\x18'           # Record type - Heartbeat
        '\x03{0}'               # TLS version
        '\x00\x03'              # Record length
        '\x01'                  # Heartbeat type - Request
        '\x00\x00')             # Heartbeat length

    payload = ('\x18'           # Record type - Heartbeat
        '\x03{0}'               # TLS version
        '\x40\x00'              # Record length
        '\x01'                  # Heartbeat type - Request
        '\x3f\xfd')             # Heartbeat length

    payload += '\x01'*16381     # Heartbeat data

    payload += (                # Second Heartbeat request with no padding
        '\x18'                  # Record type - Heartbeat
        '\x03{0}'
        '\x00\x03\x01\x00\x00'
    )

    return payload.format(SSL_VERSION_MAPPING[sslVersion])

class HeartbleedSent(SSLHandshakeRejected):
    # Awful hack #3: Use an exception to hack the handshake's control flow in
    # a super obscure way
    pass

def do_handshake_with_heartbleed(self):
    # This is nassl's code for do_handshake() modified to send a heartbleed
    # payload that will reveal 1 byte of the server's memory
    # I copied nassl's code here so I could leave anything heartbleed-related
    # outside of the nassl code base
    try:
        if self._ssl.do_handshake() == 1:
            self._handshakeDone = True
            return True # Handshake was successful

    except WantReadError:
        # OpenSSL is expecting more data from the peer
        # Send available handshake data to the peer
        # In this heartbleed handshake we only send the client hello
        lenToRead = self._networkBio.pending()
        while lenToRead:
            # Get the data from the SSL engine
            handshakeDataOut = self._networkBio.read(lenToRead)
            # Send it to the peer
            self._sock.send(handshakeDataOut)
            lenToRead = self._networkBio.pending()

        # Send the heartbleed payload after the client hello
        self._sock.send(heartbleed_payload(self.sslVersion))

        # Recover the peer's encrypted response
        # In this heartbleed handshake we only receive the server hello
        handshakeDataIn = self._sock.recv(2048)
        if len(handshakeDataIn) == 0:
            raise IOError('Nassl SSL handshake failed: peer did not send data back.')
        # Pass the data to the SSL engine
        self._networkBio.write(handshakeDataIn)

        # Signal that we sent the heartbleed payload and just stop the handshake
        raise HeartbleedSent("")

    except WantX509LookupError:
        # Server asked for a client certificate and we didn't provide one
        # Heartbleed should work anyway
        self._sock.send(heartbleed_payload(self.sslVersion)) # The heartbleed payload
        raise HeartbleedSent("") # Signal that we sent the heartbleed payload

<|endoftext|>