<commit_before>"""
Message representations received from the panel through the `AlarmDecoder`_ (AD2)
devices.

* :py:class:`Message`: The standard and most common message received from a panel.
* :py:class:`ExpanderMessage`: Messages received from Relay or Zone expander modules.
* :py:class:`RFMessage`: Message received from an RF receiver module.
* :py:class:`LRRMessage`: Message received from a long-range radio module.

.. _AlarmDecoder: http://www.alarmdecoder.com

.. moduleauthor:: Scott Petersen <scott@nutech.com>
"""

import re
import datetime

from .util import InvalidMessageError
from .panels import MODES, ADEMCO, DSC

class BaseMessage(object):
    """
    Base class for messages.
    """

    raw = None
    """The raw message text"""

    timestamp = None
    """The timestamp of the message"""

    def __init__(self):
        """
        Constructor
        """
        self.timestamp = datetime.datetime.now()

    def __str__(self):
        """
        String conversion operator.
        """
        return self.raw

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time=self.timestamp,
            mesg=self.raw,
            **kwargs
        )

    def __repr__(self):
        """
        String representation.
        """
        return repr(self.dict())

class Message(BaseMessage):
    """
    Represents a message from the alarm panel.
    """

    ready = False
    """Indicates whether or not the panel is in a ready state."""
    armed_away = False
    """Indicates whether or not the panel is armed away."""
    armed_home = False
    """Indicates whether or not the panel is armed home."""
    backlight_on = False
    """Indicates whether or not the keypad backlight is on."""
    programming_mode = False
    """Indicates whether or not we're in programming mode."""
    beeps = -1
    """Number of beeps associated with a message."""
    zone_bypassed = False
    """Indicates whether or not a zone is bypassed."""
    ac_power = False
    """Indicates whether or not the panel is on AC power."""
    chime_on = False
    """Indicates whether or not the chime is enabled."""
    alarm_event_occurred = False
    """Indicates whether or not an alarm event has occurred."""
    alarm_sounding = False
    """Indicates whether or not an alarm is sounding."""
    battery_low = False
    """Indicates whether or not there is a low battery."""
    entry_delay_off = False
    """Indicates whether or not the entry delay is enabled."""
    fire_alarm = False
    """Indicates whether or not a fire alarm is sounding."""
    check_zone = False
    """Indicates whether or not there are zones that require attention."""
    perimeter_only = False
    """Indicates whether or not the perimeter is armed."""
    system_fault = False
    """Indicates whether a system fault has occurred."""
    panel_type = ADEMCO
    """Indicates which panel type was the source of this message."""
    numeric_code = None
    """The numeric code associated with the message."""
    text = None
    """The human-readable text to be displayed on the panel LCD."""
    cursor_location = -1
    """Current cursor location on the keypad."""
    mask = None
    """Address mask this message is intended for."""
    bitfield = None
    """The bitfield associated with this message."""
    panel_data = None
    """The panel data field associated with this message."""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        self._regex = re.compile('^(!KPM:){0,1}(\[[a-fA-F0-9\-]+\]),([a-fA-F0-9]+),(\[[a-fA-F0-9]+\]),(".+")$')

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parse the message from the device.

        :param data: message data
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        match = self._regex.match(data)

        if match is None:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

        header, self.bitfield, self.numeric_code, self.panel_data, alpha = match.group(1, 2, 3, 4, 5)
        self.mask = int(self.panel_data[3:3+8], 16)

        is_bit_set = lambda bit: not self.bitfield[bit] == "0"

        self.raw = data
        self.ready = is_bit_set(1)
        self.armed_away = is_bit_set(2)
        self.armed_home = is_bit_set(3)
        self.backlight_on = is_bit_set(4)
        self.programming_mode = is_bit_set(5)
        self.beeps = int(self.bitfield[6], 16)
        self.zone_bypassed = is_bit_set(7)
        self.ac_power = is_bit_set(8)
        self.chime_on = is_bit_set(9)
        self.alarm_event_occurred = is_bit_set(10)
        self.alarm_sounding = is_bit_set(11)
        self.battery_low = is_bit_set(12)
        self.entry_delay_off = is_bit_set(13)
        self.fire_alarm = is_bit_set(14)
        self.check_zone = is_bit_set(15)
        self.perimeter_only = is_bit_set(16)
        self.system_fault = is_bit_set(17)
        if self.bitfield[18] in MODES.keys():
            self.panel_type = MODES.keys()[MODES.values().index(self.bitfield[18])])
        # pos 20-21 - Unused.
        self.text = alpha.strip('"')

        if int(self.panel_data[19:21], 16) & 0x01 > 0:
            # Current cursor location on the alpha display.
            self.cursor_location = int(self.bitfield[21:23], 16)

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            bitfield              = self.bitfield,
            numeric_code          = self.numeric_code,
            panel_data            = self.panel_data,
            mask                  = self.mask,
            ready                 = self.ready,
            armed_away            = self.armed_away,
            armed_home            = self.armed_home,
            backlight_on          = self.backlight_on,
            programming_mode      = self.programming_mode,
            beeps                 = self.beeps,
            zone_bypassed         = self.zone_bypassed,
            ac_power              = self.ac_power,
            chime_on              = self.chime_on,
            alarm_event_occurred  = self.alarm_event_occurred,
            alarm_sounding        = self.alarm_sounding,
            battery_low           = self.battery_low,
            entry_delay_off       = self.entry_delay_off,
            fire_alarm            = self.fire_alarm,
            check_zone            = self.check_zone,
            perimeter_only        = self.perimeter_only,
            text                  = self.text,
            cursor_location       = self.cursor_location,
            **kwargs
        )

class ExpanderMessage(BaseMessage):
    """
    Represents a message from a zone or relay expansion module.
    """

    ZONE = 0
    """Flag indicating that the expander message relates to a Zone Expander."""
    RELAY = 1
    """Flag indicating that the expander message relates to a Relay Expander."""

    type = None
    """Expander message type: ExpanderMessage.ZONE or ExpanderMessage.RELAY"""
    address = -1
    """Address of expander"""
    channel = -1
    """Channel on the expander"""
    value = -1
    """Value associated with the message"""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parse the raw message from the device.

        :param data: message data
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        try:
            header, values = data.split(':')
            address, channel, value = values.split(',')

            self.raw = data
            self.address = int(address)
            self.channel = int(channel)
            self.value = int(value)

        except ValueError:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

        if header == '!EXP':
            self.type = ExpanderMessage.ZONE
        elif header == '!REL':
            self.type = ExpanderMessage.RELAY
        else:
            raise InvalidMessageError('Unknown expander message header: {0}'.format(data))

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            address               = self.address,
            channel               = self.channel,
            value                 = self.value,
            **kwargs
        )

class RFMessage(BaseMessage):
    """
    Represents a message from an RF receiver.
    """

    serial_number = None
    """Serial number of the RF device."""
    value = -1
    """Value associated with this message."""
    battery = False
    """Low battery indication"""
    supervision = False
    """Supervision required indication"""
    loop = [False for _ in range(4)]
    """Loop indicators"""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parses the raw message from the device.

        :param data: message data
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        try:
            self.raw = data

            _, values = data.split(':')
            self.serial_number, self.value = values.split(',')
            self.value = int(self.value, 16)

            is_bit_set = lambda b: self.value & (1 << (b - 1)) > 0

            # Bit 1 = unknown
            self.battery = is_bit_set(2)
            self.supervision = is_bit_set(3)
            # Bit 4 = unknown
            self.loop[2] = is_bit_set(5)
            self.loop[1] = is_bit_set(6)
            self.loop[3] = is_bit_set(7)
            self.loop[0] = is_bit_set(8)

        except ValueError:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            serial_number         = self.serial_number,
            value                 = self.value,
            battery               = self.battery,
            supervision           = self.supervision,
            **kwargs
        )

class LRRMessage(BaseMessage):
    """
    Represent a message from a Long Range Radio.
    """

    event_data = None
    """Data associated with the LRR message.  Usually user ID or zone."""
    partition = -1
    """The partition that this message applies to."""
    event_type = None
    """The type of the event that occurred."""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parses the raw message from the device.

        :param data: message data to parse
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        try:
            self.raw = data

            _, values = data.split(':')
            self.event_data, self.partition, self.event_type = values.split(',')

        except ValueError:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            event_data            = self.event_data,
            event_type            = self.event_type,
            partition             = self.partition,
            **kwargs
        )
<commit_msg>Added check for ADEMCO panel before pulling out the mask and cursor location.  Mask now defaults to 0xFFFFFFFF.<commit_after>"""
Message representations received from the panel through the `AlarmDecoder`_ (AD2)
devices.

* :py:class:`Message`: The standard and most common message received from a panel.
* :py:class:`ExpanderMessage`: Messages received from Relay or Zone expander modules.
* :py:class:`RFMessage`: Message received from an RF receiver module.
* :py:class:`LRRMessage`: Message received from a long-range radio module.

.. _AlarmDecoder: http://www.alarmdecoder.com

.. moduleauthor:: Scott Petersen <scott@nutech.com>
"""

import re
import datetime

from .util import InvalidMessageError
from .panels import MODES, ADEMCO, DSC

class BaseMessage(object):
    """
    Base class for messages.
    """

    raw = None
    """The raw message text"""

    timestamp = None
    """The timestamp of the message"""

    def __init__(self):
        """
        Constructor
        """
        self.timestamp = datetime.datetime.now()

    def __str__(self):
        """
        String conversion operator.
        """
        return self.raw

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time=self.timestamp,
            mesg=self.raw,
            **kwargs
        )

    def __repr__(self):
        """
        String representation.
        """
        return repr(self.dict())

class Message(BaseMessage):
    """
    Represents a message from the alarm panel.
    """

    ready = False
    """Indicates whether or not the panel is in a ready state."""
    armed_away = False
    """Indicates whether or not the panel is armed away."""
    armed_home = False
    """Indicates whether or not the panel is armed home."""
    backlight_on = False
    """Indicates whether or not the keypad backlight is on."""
    programming_mode = False
    """Indicates whether or not we're in programming mode."""
    beeps = -1
    """Number of beeps associated with a message."""
    zone_bypassed = False
    """Indicates whether or not a zone is bypassed."""
    ac_power = False
    """Indicates whether or not the panel is on AC power."""
    chime_on = False
    """Indicates whether or not the chime is enabled."""
    alarm_event_occurred = False
    """Indicates whether or not an alarm event has occurred."""
    alarm_sounding = False
    """Indicates whether or not an alarm is sounding."""
    battery_low = False
    """Indicates whether or not there is a low battery."""
    entry_delay_off = False
    """Indicates whether or not the entry delay is enabled."""
    fire_alarm = False
    """Indicates whether or not a fire alarm is sounding."""
    check_zone = False
    """Indicates whether or not there are zones that require attention."""
    perimeter_only = False
    """Indicates whether or not the perimeter is armed."""
    system_fault = False
    """Indicates whether a system fault has occurred."""
    panel_type = ADEMCO
    """Indicates which panel type was the source of this message."""
    numeric_code = None
    """The numeric code associated with the message."""
    text = None
    """The human-readable text to be displayed on the panel LCD."""
    cursor_location = -1
    """Current cursor location on the keypad."""
    mask = 0xFFFFFFFF
    """Address mask this message is intended for."""
    bitfield = None
    """The bitfield associated with this message."""
    panel_data = None
    """The panel data field associated with this message."""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        self._regex = re.compile('^(!KPM:){0,1}(\[[a-fA-F0-9\-]+\]),([a-fA-F0-9]+),(\[[a-fA-F0-9]+\]),(".+")$')

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parse the message from the device.

        :param data: message data
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        match = self._regex.match(data)

        if match is None:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

        header, self.bitfield, self.numeric_code, self.panel_data, alpha = match.group(1, 2, 3, 4, 5)

        is_bit_set = lambda bit: not self.bitfield[bit] == "0"

        self.raw = data
        self.ready = is_bit_set(1)
        self.armed_away = is_bit_set(2)
        self.armed_home = is_bit_set(3)
        self.backlight_on = is_bit_set(4)
        self.programming_mode = is_bit_set(5)
        self.beeps = int(self.bitfield[6], 16)
        self.zone_bypassed = is_bit_set(7)
        self.ac_power = is_bit_set(8)
        self.chime_on = is_bit_set(9)
        self.alarm_event_occurred = is_bit_set(10)
        self.alarm_sounding = is_bit_set(11)
        self.battery_low = is_bit_set(12)
        self.entry_delay_off = is_bit_set(13)
        self.fire_alarm = is_bit_set(14)
        self.check_zone = is_bit_set(15)
        self.perimeter_only = is_bit_set(16)
        self.system_fault = is_bit_set(17)
        if self.bitfield[18] in MODES.keys():
            self.panel_type = MODES.keys()[MODES.values().index(self.bitfield[18])])
        # pos 20-21 - Unused.
        self.text = alpha.strip('"')

        if self.panel_type == ADEMCO:
            self.mask = int(self.panel_data[3:3+8], 16)

            if int(self.panel_data[19:21], 16) & 0x01 > 0:
                # Current cursor location on the alpha display.
                self.cursor_location = int(self.bitfield[21:23], 16)

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            bitfield              = self.bitfield,
            numeric_code          = self.numeric_code,
            panel_data            = self.panel_data,
            mask                  = self.mask,
            ready                 = self.ready,
            armed_away            = self.armed_away,
            armed_home            = self.armed_home,
            backlight_on          = self.backlight_on,
            programming_mode      = self.programming_mode,
            beeps                 = self.beeps,
            zone_bypassed         = self.zone_bypassed,
            ac_power              = self.ac_power,
            chime_on              = self.chime_on,
            alarm_event_occurred  = self.alarm_event_occurred,
            alarm_sounding        = self.alarm_sounding,
            battery_low           = self.battery_low,
            entry_delay_off       = self.entry_delay_off,
            fire_alarm            = self.fire_alarm,
            check_zone            = self.check_zone,
            perimeter_only        = self.perimeter_only,
            text                  = self.text,
            cursor_location       = self.cursor_location,
            **kwargs
        )

class ExpanderMessage(BaseMessage):
    """
    Represents a message from a zone or relay expansion module.
    """

    ZONE = 0
    """Flag indicating that the expander message relates to a Zone Expander."""
    RELAY = 1
    """Flag indicating that the expander message relates to a Relay Expander."""

    type = None
    """Expander message type: ExpanderMessage.ZONE or ExpanderMessage.RELAY"""
    address = -1
    """Address of expander"""
    channel = -1
    """Channel on the expander"""
    value = -1
    """Value associated with the message"""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parse the raw message from the device.

        :param data: message data
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        try:
            header, values = data.split(':')
            address, channel, value = values.split(',')

            self.raw = data
            self.address = int(address)
            self.channel = int(channel)
            self.value = int(value)

        except ValueError:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

        if header == '!EXP':
            self.type = ExpanderMessage.ZONE
        elif header == '!REL':
            self.type = ExpanderMessage.RELAY
        else:
            raise InvalidMessageError('Unknown expander message header: {0}'.format(data))

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            address               = self.address,
            channel               = self.channel,
            value                 = self.value,
            **kwargs
        )

class RFMessage(BaseMessage):
    """
    Represents a message from an RF receiver.
    """

    serial_number = None
    """Serial number of the RF device."""
    value = -1
    """Value associated with this message."""
    battery = False
    """Low battery indication"""
    supervision = False
    """Supervision required indication"""
    loop = [False for _ in range(4)]
    """Loop indicators"""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parses the raw message from the device.

        :param data: message data
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        try:
            self.raw = data

            _, values = data.split(':')
            self.serial_number, self.value = values.split(',')
            self.value = int(self.value, 16)

            is_bit_set = lambda b: self.value & (1 << (b - 1)) > 0

            # Bit 1 = unknown
            self.battery = is_bit_set(2)
            self.supervision = is_bit_set(3)
            # Bit 4 = unknown
            self.loop[2] = is_bit_set(5)
            self.loop[1] = is_bit_set(6)
            self.loop[3] = is_bit_set(7)
            self.loop[0] = is_bit_set(8)

        except ValueError:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            serial_number         = self.serial_number,
            value                 = self.value,
            battery               = self.battery,
            supervision           = self.supervision,
            **kwargs
        )

class LRRMessage(BaseMessage):
    """
    Represent a message from a Long Range Radio.
    """

    event_data = None
    """Data associated with the LRR message.  Usually user ID or zone."""
    partition = -1
    """The partition that this message applies to."""
    event_type = None
    """The type of the event that occurred."""

    def __init__(self, data=None):
        """
        Constructor

        :param data: message data to parse
        :type data: string
        """
        BaseMessage.__init__(self)

        if data is not None:
            self._parse_message(data)

    def _parse_message(self, data):
        """
        Parses the raw message from the device.

        :param data: message data to parse
        :type data: string

        :raises: :py:class:`~alarmdecoder.util.InvalidMessageError`
        """
        try:
            self.raw = data

            _, values = data.split(':')
            self.event_data, self.partition, self.event_type = values.split(',')

        except ValueError:
            raise InvalidMessageError('Received invalid message: {0}'.format(data))

    def dict(self, **kwargs):
        """
        Dictionary representation.
        """
        return dict(
            time                  = self.timestamp,
            event_data            = self.event_data,
            event_type            = self.event_type,
            partition             = self.partition,
            **kwargs
        )
<|endoftext|>


<commit_before>import re
import uuid
import json

from twisted.python import log
from twisted.internet.defer import inlineCallbacks, returnValue, maybeDeferred
from smpp.pdu_builder import (BindTransceiver,
                                BindTransmitter,
                                BindReceiver,
                                DeliverSMResp,
                                SubmitSM,
                                SubmitMulti,
                                EnquireLink,
                                EnquireLinkResp,
                                QuerySM,
                                )
from smpp.pdu_inspector import (MultipartMessage,
                                detect_multipart,
                                multipart_key,
                                )
from vumi.message import TransportUserMessage
from vumi.transports.smpp import SmppTransport
from vumi.transports.smpp.clientserver.client import EsmeTransceiver, EsmeTransceiverFactory

from middlewares.custom_middleware_stack import CustomMiddlewareStack, useCustomMiddleware

@useCustomMiddleware
class PushTzSmppTransport(SmppTransport):

    regex_plus = re.compile("^\+")

    DELIVERY_REPORT_STATUS_MAPPING = {
        # Output values should map to themselves:
        'delivered': 'delivered',
        'failed': 'failed',
        'pending': 'pending',
        # SMPP `message_state` values:
        'ENROUTE': 'pending',
        'DELIVERED': 'delivered',
        'EXPIRED': 'failed',
        'DELETED': 'failed',
        'UNDELIVERABLE': 'failed',
        'ACCEPTED': 'delivered',
        'UNKNOWN': 'pending',
        'REJECTED': 'failed',
        # From the most common regex-extracted format:
        'DELIVRD': 'delivered',
        'REJECTD': 'failed',
        'UNDELIV': 'failed'}

    def delivery_status(self, stat):
        return self.DELIVERY_REPORT_STATUS_MAPPING.get(stat, 'pending')

    def make_factory(self):
        return PushTzEsmeTransceiverFactory(self.client_config,
                                            self.r_server,
                                            self.esme_callbacks)

    def handle_outbound_message(self, message):
        message['from_addr'] = message['from_addr'] if 'default_origin' not in self.config else self.config['default_origin']
        super(PushTzSmppTransport, self).handle_outbound_message(message)

    def deliver_sm(self, *args, **kwargs):
        if (not re.match(self.regex_plus, kwargs.get('source_addr'))):
            kwargs.update({'source_addr': ('+%s' % kwargs.get('source_addr'))})
        super(PushTzSmppTransport, self).deliver_sm(*args, **kwargs)

class PushTzEsmeTransceiverFactory(EsmeTransceiverFactory):

    def buildProtocol(self, addr):
        log.msg('Connected')
        self.esme = PushTzEsmeTransceiver(self.config, self.kvs, self.esme_callbacks)
        self.resetDelay()
        return self.esme

class PushTzEsmeTransceiver(EsmeTransceiver):

    def detect_error_delivery(self, pdu):
        if ('optional_parameters' in pdu['body']):
            for optional_parameter in pdu['body']['optional_parameters']:
                if optional_parameter['tag'] == 'network_error_code':
                    log.msg('NETWORK ERROR CODE %s' % optional_parameter['value'])
                    return True
        return False

    def handle_deliver_sm(self, pdu):
        #if ('optional_parameters' in pdu['body']):
            #for optional_parameter in pdu['body']['optional_parameters']:
                #if optional_parameter['tag'] == 'network_error_code':
                    #log.msg('NETWORK ERROR CODE %s' % optional_parameter['value'])
                    ##return
        #EsmeTransceiver.handle_deliver_sm(self, pdu)
        if self.state not in ['BOUND_RX', 'BOUND_TRX']:
            log.err('WARNING: Received deliver_sm in wrong state: %s' % (
                self.state))

        if pdu['header']['command_status'] == 'ESME_ROK':
            sequence_number = pdu['header']['sequence_number']
            message_id = str(uuid.uuid4())
            pdu_resp = DeliverSMResp(sequence_number,
                    **self.defaults)
            self.send_pdu(pdu_resp)
            pdu_params = pdu['body']['mandatory_parameters']
            delivery_report = self.config.delivery_report_re.search(
                    pdu_params['short_message'] or ''
                    )
            if delivery_report:
                self.esme_callbacks.delivery_report(
                        destination_addr=pdu_params['destination_addr'],
                        source_addr=pdu_params['source_addr'],
                        delivery_report=delivery_report.groupdict(),
                        )
            elif detect_multipart(pdu):
                redis_key = "%s#multi_%s" % (
                        self.r_prefix, multipart_key(detect_multipart(pdu)))
                log.msg("Redis multipart key: %s" % (redis_key))
                value = json.loads(self.r_server.get(redis_key) or 'null')
                log.msg("Retrieved value: %s" % (repr(value)))
                multi = MultipartMessage(value)
                multi.add_pdu(pdu)
                completed = multi.get_completed()
                if completed:
                    self.r_server.delete(redis_key)
                    log.msg("Reassembled Message: %s" % (completed['message']))
                    # and we can finally pass the whole message on
                    self.esme_callbacks.deliver_sm(
                            destination_addr=completed['to_msisdn'],
                            source_addr=completed['from_msisdn'],
                            short_message=completed['message'],
                            message_id=message_id,
                            )
                else:
                    self.r_server.set(redis_key, json.dumps(multi.get_array()))
            elif self.detect_error_delivery(pdu):
                return
            else:
                decoded_msg = self._decode_message(pdu_params['short_message'],
                                                   pdu_params['data_coding'])
                self.esme_callbacks.deliver_sm(
                        destination_addr=pdu_params['destination_addr'],
                        source_addr=pdu_params['source_addr'],
                        short_message=decoded_msg,
                        message_id=message_id,
                        )

    def submit_sm(self, **kwargs):
        if self.state not in ['BOUND_TX', 'BOUND_TRX']:
            log.err(('WARNING: submit_sm in wrong state: %s, '
                            'dropping message: %s' % (self.state, kwargs)))
            return 0
        else:
            sequence_number = self.get_seq()
            pdu = SubmitSM(sequence_number, **dict(self.defaults, **kwargs))
            self.get_next_seq()

            message = kwargs['short_message']
            if len(message) > 254:
                pdu.add_message_payload(''.join('%02x' % ord(c) for c in message))

            self.send_pdu(pdu)
            self.push_unacked(sequence_number)
            return sequence_number
<commit_msg>fix pushSMPP to extend the enhanced client<commit_after>import re
import uuid
import json

from twisted.python import log
from twisted.internet.defer import inlineCallbacks, returnValue, maybeDeferred
from smpp.pdu_builder import (BindTransceiver,
                                BindTransmitter,
                                BindReceiver,
                                DeliverSMResp,
                                SubmitSM,
                                SubmitMulti,
                                EnquireLink,
                                EnquireLinkResp,
                                QuerySM,
                                )
from smpp.pdu_inspector import (MultipartMessage,
                                detect_multipart,
                                multipart_key,
                                )
from vumi.message import TransportUserMessage
from vumi.transports.smpp.clientserver.client import EsmeTransceiver, EsmeTransceiverFactory

from middlewares.custom_middleware_stack import CustomMiddlewareStack, useCustomMiddleware
from transports.enhanced_smpp.enhanced_client import (EnhancedSmppTransport,
                                                      EnhancedEsmeTransceiver,
                                                      EnhancedEsmeTransceiverFactory)

@useCustomMiddleware
class PushTzSmppTransport(EnhancedSmppTransport):

    regex_plus = re.compile("^\+")

    def delivery_status(self, stat):
        return self.DELIVERY_REPORT_STATUS_MAPPING.get(stat, 'pending')

    def make_factory(self):
        return PushTzEsmeTransceiverFactory(self.client_config,
                                            self.r_server,
                                            self.esme_callbacks)

    def handle_outbound_message(self, message):
        message['from_addr'] = message['from_addr'] if 'default_origin' not in self.config else self.config['default_origin']
        super(PushTzSmppTransport, self).handle_outbound_message(message)

    def deliver_sm(self, *args, **kwargs):
        if (not re.match(self.regex_plus, kwargs.get('source_addr'))):
            kwargs.update({'source_addr': ('+%s' % kwargs.get('source_addr'))})
        super(PushTzSmppTransport, self).deliver_sm(*args, **kwargs)

class PushTzEsmeTransceiverFactory(EnhancedEsmeTransceiverFactory):

    def buildProtocol(self, addr):
        log.msg('Connected')
        self.esme = PushTzEsmeTransceiver(self.config, self.kvs, self.esme_callbacks)
        self.resetDelay()
        return self.esme

class PushTzEsmeTransceiver(EnhancedEsmeTransceiver):

    def detect_error_delivery(self, pdu):
        if ('optional_parameters' in pdu['body']):
            for optional_parameter in pdu['body']['optional_parameters']:
                if optional_parameter['tag'] == 'network_error_code':
                    log.msg('NETWORK ERROR CODE %s' % optional_parameter['value'])
                    return True
        return False

    def handle_deliver_sm(self, pdu):
        #if ('optional_parameters' in pdu['body']):
            #for optional_parameter in pdu['body']['optional_parameters']:
                #if optional_parameter['tag'] == 'network_error_code':
                    #log.msg('NETWORK ERROR CODE %s' % optional_parameter['value'])
                    ##return
        #EsmeTransceiver.handle_deliver_sm(self, pdu)
        if self.state not in ['BOUND_RX', 'BOUND_TRX']:
            log.err('WARNING: Received deliver_sm in wrong state: %s' % (
                self.state))

        if pdu['header']['command_status'] == 'ESME_ROK':
            sequence_number = pdu['header']['sequence_number']
            message_id = str(uuid.uuid4())
            pdu_resp = DeliverSMResp(sequence_number,
                    **self.defaults)
            self.send_pdu(pdu_resp)
            pdu_params = pdu['body']['mandatory_parameters']
            delivery_report = self.config.delivery_report_re.search(
                    pdu_params['short_message'] or ''
                    )
            if delivery_report:
                self.esme_callbacks.delivery_report(
                        destination_addr=pdu_params['destination_addr'],
                        source_addr=pdu_params['source_addr'],
                        delivery_report=delivery_report.groupdict(),
                        )
            elif detect_multipart(pdu):
                redis_key = "%s#multi_%s" % (
                        self.r_prefix, multipart_key(detect_multipart(pdu)))
                log.msg("Redis multipart key: %s" % (redis_key))
                value = json.loads(self.r_server.get(redis_key) or 'null')
                log.msg("Retrieved value: %s" % (repr(value)))
                multi = MultipartMessage(value)
                multi.add_pdu(pdu)
                completed = multi.get_completed()
                if completed:
                    self.r_server.delete(redis_key)
                    log.msg("Reassembled Message: %s" % (completed['message']))
                    # and we can finally pass the whole message on
                    self.esme_callbacks.deliver_sm(
                            destination_addr=completed['to_msisdn'],
                            source_addr=completed['from_msisdn'],
                            short_message=completed['message'],
                            message_id=message_id,
                            )
                else:
                    self.r_server.set(redis_key, json.dumps(multi.get_array()))
            elif self.detect_error_delivery(pdu):
                return
            else:
                decoded_msg = self._decode_message(pdu_params['short_message'],
                                                   pdu_params['data_coding'])
                self.esme_callbacks.deliver_sm(
                        destination_addr=pdu_params['destination_addr'],
                        source_addr=pdu_params['source_addr'],
                        short_message=decoded_msg,
                        message_id=message_id,
                        )
    <|endoftext|>


<commit_before>'''
Audio
=====

Load an audio sound and play it with::

    from kivy.core.audio import SoundLoader

    sound = SoundLoader.load('mytest.wav')
    if sound:
        print("Sound found at %s" % sound.source)
        print("Sound is %.3f seconds" % sound.length)
        sound.play()

You should not use the Sound class directly. The class returned by
:func:`SoundLoader.load` will be the best sound provider for that particular
file type, so it might return different Sound classes depending the file type.

Event dispatching and state changes
-----------------------------------

Audio is often processed in parallel to your code. This means you often need to
enter the Kivy :func:`eventloop <kivy.base.EventLoopBase>` in order to allow
events and state changes to be dispatched correctly.

You seldom need to worry about this as Kivy apps typically always
require this event loop for the GUI to remain responsive, but it is good to
keep this in mind when debugging or running in a
`REPL <https://en.wikipedia.org/wiki/Read%E2%80%93eval%E2%80%93print_loop>`_
(Read-eval-print loop).

.. versionchanged:: 1.8.0
    There are now 2 distinct Gstreamer implementations: one using Gi/Gst
    working for both Python 2+3 with Gstreamer 1.0, and one using PyGST
    working only for Python 2 + Gstreamer 0.10.

.. note::

    The core audio library does not support recording audio. If you require
    this functionality, please refer to the
    `audiostream <https://github.com/kivy/audiostream>`_ extension.

'''

__all__ = ('Sound', 'SoundLoader')

from kivy.logger import Logger
from kivy.event import EventDispatcher
from kivy.core import core_register_libs
from kivy.compat import PY2
from kivy.resources import resource_find
from kivy.properties import StringProperty, NumericProperty, OptionProperty, \
    AliasProperty, BooleanProperty, BoundedNumericProperty
from kivy.utils import platform
from kivy.setupconfig import USE_SDL2

from sys import float_info

class SoundLoader:
    '''Load a sound, using the best loader for the given file type.
    '''

    _classes = []

    @staticmethod
    def register(classobj):
        '''Register a new class to load the sound.'''
        Logger.debug('Audio: register %s' % classobj.__name__)
        SoundLoader._classes.append(classobj)

    @staticmethod
    def load(filename):
        '''Load a sound, and return a Sound() instance.'''
        rfn = resource_find(filename)
        if rfn is not None:
            filename = rfn
        ext = filename.split('.')[-1].lower()
        if '?' in ext:
            ext = ext.split('?')[0]
        for classobj in SoundLoader._classes:
            if ext in classobj.extensions():
                return classobj(source=filename)
        Logger.warning('Audio: Unable to find a loader for <%s>' %
                       filename)
        return None

class Sound(EventDispatcher):
    '''Represents a sound to play. This class is abstract, and cannot be used
    directly.

    Use SoundLoader to load a sound.

    :Events:
        `on_play`: None
            Fired when the sound is played.
        `on_stop`: None
            Fired when the sound is stopped.
    '''

    source = StringProperty(None)
    '''Filename / source of your audio file.

    .. versionadded:: 1.3.0

    :attr:`source` is a :class:`~kivy.properties.StringProperty` that defaults
    to None and is read-only. Use the :meth:`SoundLoader.load` for loading
    audio.
    '''

    volume = NumericProperty(1.)
    '''Volume, in the range 0-1. 1 means full volume, 0 means mute.

    .. versionadded:: 1.3.0

    :attr:`volume` is a :class:`~kivy.properties.NumericProperty` and defaults
    to 1.
    '''

    pitch = BoundedNumericProperty(1., min=float_info.epsilon)
    '''Pitch of a sound. 2 is an octave higher, .5 one below. This is only
    implemented for SDL2 audio provider yet.

    .. versionadded:: 1.9.2

    :attr:`pitch` is a :class:`~kivy.properties.NumericProperty` and defaults
    to 1.
    '''

    state = OptionProperty('stop', options=('stop', 'play'))
    '''State of the sound, one of 'stop' or 'play'.

    .. versionadded:: 1.3.0

    :attr:`state` is a read-only :class:`~kivy.properties.OptionProperty`.'''

    loop = BooleanProperty(False)
    '''Set to True if the sound should automatically loop when it finishes.

    .. versionadded:: 1.8.0

    :attr:`loop` is a :class:`~kivy.properties.BooleanProperty` and defaults to
    False.'''

    #
    # deprecated
    #
    def _get_status(self):
        return self.state
    status = AliasProperty(_get_status, None, bind=('state', ))
    '''
    .. deprecated:: 1.3.0
        Use :attr:`state` instead.
    '''

    def _get_filename(self):
        return self.source
    filename = AliasProperty(_get_filename, None, bind=('source', ))
    '''
    .. deprecated:: 1.3.0
        Use :attr:`source` instead.
    '''

    __events__ = ('on_play', 'on_stop')

    def on_source(self, instance, filename):
        self.unload()
        if filename is None:
            return
        self.load()

    def get_pos(self):
        '''
        Returns the current position of the audio file.
        Returns 0 if not playing.

        .. versionadded:: 1.4.1
        '''
        return 0

    def _get_length(self):
        return 0

    length = property(lambda self: self._get_length(),
                      doc='Get length of the sound (in seconds).')

    def load(self):
        '''Load the file into memory.'''
        pass

    def unload(self):
        '''Unload the file from memory.'''
        pass

    def play(self):
        '''Play the file.'''
        self.state = 'play'
        self.dispatch('on_play')

    def stop(self):
        '''Stop playback.'''
        self.state = 'stop'
        self.dispatch('on_stop')

    def seek(self, position):
        '''Go to the <position> (in seconds).'''
        pass

    def on_play(self):
        pass

    def on_stop(self):
        pass

# Little trick here, don't activate gstreamer on window
# seem to have lot of crackle or something...
audio_libs = []
if platform in ('macosx', 'ios'):
    audio_libs += [('avplayer', 'audio_avplayer')]
# from now on, prefer our gstplayer instead of gi/pygst.
try:
    from kivy.lib.gstplayer import GstPlayer  # NOQA
    audio_libs += [('gstplayer', 'audio_gstplayer')]
except ImportError:
    # audio_libs += [('gi', 'audio_gi')]
    if PY2:
        audio_libs += [('pygst', 'audio_pygst')]
audio_libs += [('ffpyplayer', 'audio_ffpyplayer')]
if USE_SDL2:
    audio_libs += [('sdl2', 'audio_sdl2')]
else:
    audio_libs += [('pygame', 'audio_pygame')]

libs_loaded = core_register_libs('audio', audio_libs)
<commit_msg>Added note to seek that sound must be playing. (#4875)<commit_after>'''
Audio
=====

Load an audio sound and play it with::

    from kivy.core.audio import SoundLoader

    sound = SoundLoader.load('mytest.wav')
    if sound:
        print("Sound found at %s" % sound.source)
        print("Sound is %.3f seconds" % sound.length)
        sound.play()

You should not use the Sound class directly. The class returned by
:func:`SoundLoader.load` will be the best sound provider for that particular
file type, so it might return different Sound classes depending the file type.

Event dispatching and state changes
-----------------------------------

Audio is often processed in parallel to your code. This means you often need to
enter the Kivy :func:`eventloop <kivy.base.EventLoopBase>` in order to allow
events and state changes to be dispatched correctly.

You seldom need to worry about this as Kivy apps typically always
require this event loop for the GUI to remain responsive, but it is good to
keep this in mind when debugging or running in a
`REPL <https://en.wikipedia.org/wiki/Read%E2%80%93eval%E2%80%93print_loop>`_
(Read-eval-print loop).

.. versionchanged:: 1.8.0
    There are now 2 distinct Gstreamer implementations: one using Gi/Gst
    working for both Python 2+3 with Gstreamer 1.0, and one using PyGST
    working only for Python 2 + Gstreamer 0.10.

.. note::

    The core audio library does not support recording audio. If you require
    this functionality, please refer to the
    `audiostream <https://github.com/kivy/audiostream>`_ extension.

'''

__all__ = ('Sound', 'SoundLoader')

from kivy.logger import Logger
from kivy.event import EventDispatcher
from kivy.core import core_register_libs
from kivy.compat import PY2
from kivy.resources import resource_find
from kivy.properties import StringProperty, NumericProperty, OptionProperty, \
    AliasProperty, BooleanProperty, BoundedNumericProperty
from kivy.utils import platform
from kivy.setupconfig import USE_SDL2

from sys import float_info

class SoundLoader:
    '''Load a sound, using the best loader for the given file type.
    '''

    _classes = []

    @staticmethod
    def register(classobj):
        '''Register a new class to load the sound.'''
        Logger.debug('Audio: register %s' % classobj.__name__)
        SoundLoader._classes.append(classobj)

    @staticmethod
    def load(filename):
        '''Load a sound, and return a Sound() instance.'''
        rfn = resource_find(filename)
        if rfn is not None:
            filename = rfn
        ext = filename.split('.')[-1].lower()
        if '?' in ext:
            ext = ext.split('?')[0]
        for classobj in SoundLoader._classes:
            if ext in classobj.extensions():
                return classobj(source=filename)
        Logger.warning('Audio: Unable to find a loader for <%s>' %
                       filename)
        return None

class Sound(EventDispatcher):
    '''Represents a sound to play. This class is abstract, and cannot be used
    directly.

    Use SoundLoader to load a sound.

    :Events:
        `on_play`: None
            Fired when the sound is played.
        `on_stop`: None
            Fired when the sound is stopped.
    '''

    source = StringProperty(None)
    '''Filename / source of your audio file.

    .. versionadded:: 1.3.0

    :attr:`source` is a :class:`~kivy.properties.StringProperty` that defaults
    to None and is read-only. Use the :meth:`SoundLoader.load` for loading
    audio.
    '''

    volume = NumericProperty(1.)
    '''Volume, in the range 0-1. 1 means full volume, 0 means mute.

    .. versionadded:: 1.3.0

    :attr:`volume` is a :class:`~kivy.properties.NumericProperty` and defaults
    to 1.
    '''

    pitch = BoundedNumericProperty(1., min=float_info.epsilon)
    '''Pitch of a sound. 2 is an octave higher, .5 one below. This is only
    implemented for SDL2 audio provider yet.

    .. versionadded:: 1.9.2

    :attr:`pitch` is a :class:`~kivy.properties.NumericProperty` and defaults
    to 1.
    '''

    state = OptionProperty('stop', options=('stop', 'play'))
    '''State of the sound, one of 'stop' or 'play'.

    .. versionadded:: 1.3.0

    :attr:`state` is a read-only :class:`~kivy.properties.OptionProperty`.'''

    loop = BooleanProperty(False)
    '''Set to True if the sound should automatically loop when it finishes.

    .. versionadded:: 1.8.0

    :attr:`loop` is a :class:`~kivy.properties.BooleanProperty` and defaults to
    False.'''

    #
    # deprecated
    #
    def _get_status(self):
        return self.state
    status = AliasProperty(_get_status, None, bind=('state', ))
    '''
    .. deprecated:: 1.3.0
        Use :attr:`state` instead.
    '''

    def _get_filename(self):
        return self.source
    filename = AliasProperty(_get_filename, None, bind=('source', ))
    '''
    .. deprecated:: 1.3.0
        Use :attr:`source` instead.
    '''

    __events__ = ('on_play', 'on_stop')

    def on_source(self, instance, filename):
        self.unload()
        if filename is None:
            return
        self.load()

    def get_pos(self):
        '''
        Returns the current position of the audio file.
        Returns 0 if not playing.

        .. versionadded:: 1.4.1
        '''
        return 0

    def _get_length(self):
        return 0

    length = property(lambda self: self._get_length(),
                      doc='Get length of the sound (in seconds).')

    def load(self):
        '''Load the file into memory.'''
        pass

    def unload(self):
        '''Unload the file from memory.'''
        pass

    def play(self):
        '''Play the file.'''
        self.state = 'play'
        self.dispatch('on_play')

    def stop(self):
        '''Stop playback.'''
        self.state = 'stop'
        self.dispatch('on_stop')

    def seek(self, position):
        '''Go to the <position> (in seconds).

        .. note::
            Most sound providers cannot seek when the audio is stopped.
            Play then seek.
        '''
        pass

    def on_play(self):
        pass

    def on_stop(self):
        pass

# Little trick here, don't activate gstreamer on window
# seem to have lot of crackle or something...
audio_libs = []
if platform in ('macosx', 'ios'):
    audio_libs += [('avplayer', 'audio_avplayer')]
# from now on, prefer our gstplayer instead of gi/pygst.
try:
    from kivy.lib.gstplayer import GstPlayer  # NOQA
    audio_libs += [('gstplayer', 'audio_gstplayer')]
except ImportError:
    # audio_libs += [('gi', 'audio_gi')]
    if PY2:
        audio_libs += [('pygst', 'audio_pygst')]
audio_libs += [('ffpyplayer', 'audio_ffpyplayer')]
if USE_SDL2:
    audio_libs += [('sdl2', 'audio_sdl2')]
else:
    audio_libs += [('pygame', 'audio_pygame')]

libs_loaded = core_register_libs('audio', audio_libs)
<|endoftext|>


<commit_before>import random
import math
from itertools import product
from itertools import chain
from thread_sum import ThreadSum

n = 128
N = 1000
d = 1339324948 #31 bits
#d = 265463185775234288092201526257649332892133932494809907679736989141210619872096735078276597605954826202148

#d = 13393249480990767973698914121061987209673507827659760595482620214891467806973397091277092174174393961305025200909026701058146674630543302845868229392185528

#q = 2**252 + 27742317777372353535851937790883648493 #ed25519

#temp = int("8335dc163bb124b65129c96fde933d8d723a70aadc873d6d54a7bb0d", 16)
q = 2^446 + 13818066809895115352007386748515426880336692474882178609894547503885 #Ed448-Goldilocks

print temp

window_size = 5
RANDOMIZED_BITS = 0
RANDOMIZE_V_0 = 0
r = []
v = []
alpha = []

def generate_v_values():
    for i in xrange(0, N):
        value = d + (alpha[i]*q)
        v.append(value)

def  generate_alpha_js():
    for i in xrange(1, N+1):
        al = r[i] - r[0]
        alpha.append(int(math.fabs(al)))

def  generate_r_js():
    for i in xrange(0, N+1):
        a = random.getrandbits(n)
        r.append(int(math.fabs(a)))

def bit_flip_fixed(bit_list, nr_of_changes):
    bit_list_t = bit_list[:]
    pos_list = []
    if len(bit_list) < nr_of_changes:
        raise Exception("Randomized bigger then d+(a*r)")
        print "Lenght: ", len(bit_list)
    print "FIXED: ", nr_of_changes
    #position = len(bit_list_t)-1
    for i in xrange(0, nr_of_changes):
        if bit_list_t[i] == 1:
            bit_list_t[i] = 0
        else:
            bit_list_t[i] = 1
        #position = position-1

    return bit_list

def bit_flip_random(bit_list, randomness):
    bit_list_t = bit_list[:]
    pos_list = []
    if len(bit_list) < randomness:
        raise Exception("Randomized bigger then d+(a*r)")
        print "Lenght: ", len(bit_list)
    for i in xrange(0, randomness):
        pos_bit_to_flip = random.randint(0, len(bit_list)-1)
        while(pos_bit_to_flip in pos_list):
            pos_bit_to_flip = random.randint(0, len(bit_list)-1)
        pos_list.append(pos_bit_to_flip)
        if bit_list_t[pos_bit_to_flip] == 1:
            bit_list_t[pos_bit_to_flip] = 0
        else:
            bit_list_t[pos_bit_to_flip] = 1

    return bit_list_t

def generate_v_values_with_bit_flip():
    value = d + (alpha[0]*q)
    bit_list = int_to_bin(value)
    print len(bit_list)
    bit_list_flipped = bit_flip_random(bit_list, RANDOMIZE_V_0)
    value_flipped = bin_to_int(bit_list_flipped)
    v.append(value_flipped)
    for i in xrange(1, N):
        value = d + (alpha[i]*q)
        bit_list = int_to_bin(value)
        #print len(bit_list)
        bit_list_flipped = bit_flip_random(bit_list, RANDOMIZED_BITS)
        value_flipped = bin_to_int(bit_list_flipped)
        v.append(value_flipped)

def int_to_bin(number):
    return [int(x) for x in bin(number)[2:]]

def bin_to_int(bit_list):
    output = 0
    for bit in bit_list:
        output = output * 2 + bit
    return output

def groupsof(n, xs):
    if len(xs) < n:
        return [xs]
    else:
        return chain([xs[0:n]], groupsof(n, xs[n:]))

def sum_all_ds(d_candidates, interval, mod_value):
    pairs = {}
    number_of_threads = 10
    ds = list(groupsof(len(d_candidates)/number_of_threads, d_candidates))
    threads = []
    for i in xrange(0, number_of_threads):
        threads.append(ThreadSum(i, ds[i], v, alpha, N, mod_value, interval))
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    for t in threads:
        key, d = t.return_result()
        try:
            if pairs[key] <> None:
                val = pairs[key]
                if val.count(1) > d.count(1):
                    pairs[key] = d
        except Exception as e:
            pairs[key] = d
    return min(pairs.keys()) , pairs

def test_d(to_test):
    """ Function to test the candidate to d. In our case, it is a comparasion
    with the original d. However, in a real case could be the ciphered text with the original
    and the candidate"""
    return (d==to_test)

def wide_widow_attack():
    generate_r_js()
    generate_alpha_js()
    generate_v_values_with_bit_flip()
    print "d = ",  int_to_bin(d), " len: ", len(int_to_bin(d))
    print "Starting...."
    w_prime = 0
    w = window_size
    d_prime = 0
    variations = []
    for i in product([0,1], repeat=window_size):
        variations.append(list(i))

    while(w < (n + window_size + window_size)):
        print "w: ", w
        print "w_prime: ", w_prime
        mod_value = 2**w
        d_prime = d_prime % mod_value
        d_prime_bin = int_to_bin(d_prime)
        to_iterate = []
        for variation in variations:
            to_iterate.append(variation+d_prime_bin)
        sum_d , d_candidate = sum_all_ds(to_iterate, w, mod_value)
        d_prime = bin_to_int(d_candidate[sum_d])
        print "sum: ", sum_d, " d_candidate = ", int_to_bin(d_prime)
        w_prime = w
        w = w + window_size

        if test_d(d_prime):
            break

    if (d == d_prime):
        print "FOUND KEY."
    else:
        print "SORRY"
    print "Finished."

wide_widow_attack()
<commit_msg> added number of attemps<commit_after>import random
import math
from itertools import product
from itertools import chain
from thread_sum import ThreadSum

n = 256
N = 250
#d = 1339324948 #31 bits
d = 265463185775234288092201526257649332892133932494809907679736989141210619872096735078276597605954826202148

#d = 13393249480990767973698914121061987209673507827659760595482620214891467806973397091277092174174393961305025200909026701058146674630543302845868229392185528

#q = 2**252 + 27742317777372353535851937790883648493 #ed25519

#temp = int("8335dc163bb124b65129c96fde933d8d723a70aadc873d6d54a7bb0d", 16)
q = 2^446 + 13818066809895115352007386748515426880336692474882178609894547503885 #Ed448-Goldilocks

window_size = 5
RANDOMIZED_BITS = 128
RANDOMIZE_V_0 = 0
r = []
v = []
alpha = []

def generate_v_values():
    for i in xrange(0, N):
        value = d + (alpha[i]*q)
        v.append(value)

def  generate_alpha_js():
    for i in xrange(1, N+1):
        al = r[i] - r[0]
        alpha.append(int(math.fabs(al)))

def  generate_r_js():
    for i in xrange(0, N+1):
        a = random.getrandbits(n)
        r.append(int(math.fabs(a)))

def bit_flip_fixed(bit_list, nr_of_changes):
    bit_list_t = bit_list[:]
    pos_list = []
    if len(bit_list) < nr_of_changes:
        raise Exception("Randomized bigger then d+(a*r)")
        print "Lenght: ", len(bit_list)
    print "FIXED: ", nr_of_changes
    #position = len(bit_list_t)-1
    for i in xrange(0, nr_of_changes):
        if bit_list_t[i] == 1:
            bit_list_t[i] = 0
        else:
            bit_list_t[i] = 1
        #position = position-1

    return bit_list

def bit_flip_random(bit_list, randomness):
    bit_list_t = bit_list[:]
    pos_list = []
    if len(bit_list) < randomness:
        raise Exception("Randomized bigger then d+(a*r)")
        print "Lenght: ", len(bit_list)
    for i in xrange(0, randomness):
        pos_bit_to_flip = random.randint(0, len(bit_list)-1)
        while(pos_bit_to_flip in pos_list):
            pos_bit_to_flip = random.randint(0, len(bit_list)-1)
        pos_list.append(pos_bit_to_flip)
        if bit_list_t[pos_bit_to_flip] == 1:
            bit_list_t[pos_bit_to_flip] = 0
        else:
            bit_list_t[pos_bit_to_flip] = 1

    return bit_list_t

def generate_v_values_with_bit_flip():
    value = d + (alpha[0]*q)
    bit_list = int_to_bin(value)
    print len(bit_list)
    bit_list_flipped = bit_flip_random(bit_list, RANDOMIZE_V_0)
    value_flipped = bin_to_int(bit_list_flipped)
    v.append(value_flipped)
    for i in xrange(1, N):
        value = d + (alpha[i]*q)
        bit_list = int_to_bin(value)
        #print len(bit_list)
        bit_list_flipped = bit_flip_random(bit_list, RANDOMIZED_BITS)
        value_flipped = bin_to_int(bit_list_flipped)
        v.append(value_flipped)

def int_to_bin(number):
    return [int(x) for x in bin(number)[2:]]

def bin_to_int(bit_list):
    output = 0
    for bit in bit_list:
        output = output * 2 + bit
    return output

def groupsof(n, xs):
    if len(xs) < n:
        return [xs]
    else:
        return chain([xs[0:n]], groupsof(n, xs[n:]))

def sum_all_ds(d_candidates, interval, mod_value):
    pairs = {}
    number_of_threads = 10
    ds = list(groupsof(len(d_candidates)/number_of_threads, d_candidates))
    threads = []
    for i in xrange(0, number_of_threads):
        threads.append(ThreadSum(i, ds[i], v, alpha, N, mod_value, interval))
    for t in threads:
        t.start()
    for t in threads:
        t.join()

    for t in threads:
        key, d = t.return_result()
        try:
            if pairs[key] <> None:
                val = pairs[key]
                if val.count(1) > d.count(1):
                    pairs[key] = d
        except Exception as e:
            pairs[key] = d
    return min(pairs.keys()) , pairs

def test_d(to_test):
    """ Function to test the candidate to d. In our case, it is a comparasion
    with the original d. However, in a real case could be the ciphered text with the original
    and the candidate"""
    return (d==to_test)

def wide_widow_attack():
    generate_r_js()
    generate_alpha_js()
    generate_v_values_with_bit_flip()
    print "d = ",  int_to_bin(d), " len: ", len(int_to_bin(d))
    print "Starting...."
    w_prime = 0
    w = window_size
    d_prime = 0
    variations = []
    for i in product([0,1], repeat=window_size):
        variations.append(list(i))
    while(w < (n + window_size + window_size)):
        print "w: ", w
        print "w_prime: ", w_prime
        mod_value = 2**w
        d_prime = d_prime % mod_value
        d_prime_bin = int_to_bin(d_prime)
        to_iterate = []
        for variation in variations:
            to_iterate.append(variation+d_prime_bin)
        sum_d , d_candidate = sum_all_ds(to_iterate, w, mod_value)
        d_prime = bin_to_int(d_candidate[sum_d])
        print "sum: ", sum_d, " d_candidate = ", int_to_bin(d_prime)
        w_prime = w
        w = w + window_size

        if test_d(d_prime):
            break

    if (d == d_prime):
        print "FOUND KEY."
        return True
    else:
        print "SORRY"
        return False
    print "Finished."

nr_of_attempts = 10
correct = 0
for i in xrange(0,nr_of_attempts):
    if wide_widow_attack():
        correct = correct + 1

print "Correct: ", correct, " of ", nr_of_attempts
<|endoftext|>


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

import angr
import driller.config as config

import argparse
import redis
import driller.tasks
import os
import subprocess
import signal
import sys
import multiprocessing
import cPickle as pickle
import time
import tempfile
import termcolor
import string

import logging

l = logging.getLogger("fuzz")
l.setLevel("INFO")

# global list of processes so we can kill them on SIGINT
procs = [ ]

start_time = 0

### EXIT HANDLERS

def kill_procs():
    map(lambda p: p.terminate(), procs)

def terminate(signal, frame):
    kill_procs()
    sys.exit(0)

### DICTIONARY CREATION

def hexescape(s):
    out = [ ]
    acceptable = string.letters + string.digits + " ."
    for c in s:
        if c not in acceptable:
            out.append("\\x%02x" % ord(c))
        else:
            out.append(c)

    return ''.join(out)

def create_dict(binary, outfile):
    b = angr.Project(binary)
    cfg = b.analyses.CFG(keep_input_state=True, enable_advanced_backward_slicing=True)

    string_references = [ ]
    for f in cfg.function_manager.functions.values():
        try:
            string_references.append(f.string_references())
        except ZeroDivisionError:
            pass

    string_references = sum(string_references, [])

    strings = [] if len(string_references) == 0 else zip(*string_references)[1]

    if len(strings) > 0:
        with open(outfile, 'wb') as f:
            for i, string in enumerate(strings):
                s = hexescape(string)
                f.write("driller_%d=\"%s\"\n" % (i, s))

        return True

    return False

### BEHAVIOR TESTING
def terminates_on_eof(qemu_dir, binary):
    l.info("attempting to detect if the binary doesn't terminate on EOF")

    # detect the binary type
    b = angr.Project(binary)
    ld_arch = b.loader.main_bin.arch
    ld_type = b.loader.main_bin.filetype

    if ld_type == "cgc":
        arch = "cgc"

    elif ld_type == "elf":
        if ld_arch == archinfo.arch_amd64.ArchAMD64:
            arch = "x86_64"
        if ld_arch == archinfo.arch_x86.ArchX86:
            arch = "i386"

    else:
        l.error("binary is of an unsupported architecture")
        raise NotImplementedError

    qemu_path = os.path.join(qemu_dir, "driller-qemu-%s" % arch)
    if not os.access(qemu_path, os.X_OK):
        l.error("either qemu does not exist in config.QEMU_DIR or it is not executable")
        return True

    # create a dumb test input
    fd, tinput = tempfile.mkstemp()
    os.close(fd)

    with open(tinput, 'wb') as f:
        f.write("fuzz")

    with open(tinput, 'rb') as i:
        with open('/dev/null', 'w') as o:
            args = [qemu_path, binary]
            p = subprocess.Popen(args, stdin=i, stdout=o)

            time.sleep(2) # generously give the binary two seconds to terminate

            if p.poll() is None: # it doesn't seem to terminate on EOF
                p.terminate()
                return False

    # good, it terminates on EOF, no monkey business
    return True

### AFL SPAWNERS

def start_afl_instance(afl_path, binary, in_dir, out_dir, fuzz_id, dictionary=None, memory="8G",
                driller=None, eof_termination=False):

    args = [afl_path]

    args += ["-i", in_dir]
    args += ["-o", out_dir]
    args += ["-m", memory]
    args += ["-Q"]
    if fuzz_id == 0:
        args += ["-M", "fuzzer-master"]
        outfile = "fuzzer-master.log"
    else:
        args += ["-S", "fuzzer-%d" % fuzz_id]
        outfile = "fuzzer-%d.log" % fuzz_id

    if dictionary is not None:
        args += ["-x", dictionary]

    if driller is not None:
        args += ["-D", driller]

    if eof_termination:
        args += ["-E"]

    args += ["--", binary]

    l.debug("execing: %s > %s" % (' '.join(args), outfile))

    # kind of gross hack to get the output directory
    work_dir = os.path.dirname(out_dir[:-1])

    outfile = os.path.join(work_dir, outfile)
    fp = open(outfile, "w")

    return subprocess.Popen(args, stdout=fp)

def start_afl(afl_path, binary_path, in_dir, out_dir, afl_count, driller_path, dictionary, eof_exit):
    global procs

    # spin up the master AFL instance
    master = start_afl_instance(afl_path,
                                binary_path,
                                in_dir,
                                out_dir,
                                fuzz_id=0, # the master fuzzer
                                dictionary=dictionary,
                                eof_termination=eof_exit)
    procs.append(master)

    if afl_count > 1:
        driller = start_afl_instance(afl_path,
                                     binary_path,
                                     in_dir,
                                     out_dir,
                                     1,
                                     driller=driller_path,
                                     dictionary=dictionary,
                                     eof_termination=eof_exit)
        procs.append(driller)

    else:
        l.warning("only one AFL instance was chosen to be spun up, driller will never be invoked")

    # only spins up an AFL instances if afl_count > 1
    for n in range(2, afl_count):
        slave = start_afl_instance(afl_path,
                                binary_path,
                                in_dir,
                                out_dir,
                                n,
                                dictionary=dictionary,
                                eof_termination=eof_exit)

        procs.append(slave)

### BACKEND HANDLERS

def clear_redis(identifier):
    redis_inst = redis.Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, db=config.REDIS_DB)

    # delete all the catalogue entries
    redis_inst.delete("%s-catalogue" % identifier)

    # delete all the traced entries
    redis_inst.delete("%s-traced" % identifier)

    # delete all the crash-found entry
    redis_inst.delete("%s-crash-found" % identifier)

def report_crash_found(identifier):
    redis_inst = redis.Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, db=config.REDIS_DB)

    # add True as a member
    redis_inst.sadd(identifier + "-crash-found", True)

def listen(queue_dir, channel):
    l.info("subscring to redis channel %s" % channel)
    l.info("new inputs will be placed into %s" % queue_dir)

    try:
        os.makedirs(queue_dir)
    except OSError:
        l.warning("could not create output directory '%s'" % queue_dir)

    redis_inst = redis.Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, db=config.REDIS_DB)
    p = redis_inst.pubsub()

    p.subscribe(channel)

    input_cnt = 0

    for msg in p.listen():
        if msg['type'] == 'message':
            real_msg = pickle.loads(msg['data'])
            out_filename = "driller-%d-%x-%x" % real_msg['meta']
            l.info("dumping new input to %s" % out_filename)
            afl_name = "id:%06d,src:%s" % (input_cnt, out_filename)
            out_file = os.path.join(queue_dir, afl_name)

            with open(out_file, 'wb') as ofp:
                ofp.write(real_msg['data'])

            input_cnt += 1

def start_redis_listener(identifier, out_dir):
    driller_queue_dir = os.path.join(out_dir, "driller", "queue")
    channel = "%s-generated" % identifier

    p = multiprocessing.Process(target=listen, args=(driller_queue_dir, channel,))
    p.start()
    return p

### STATS

def show_afl_stats(sync_dir):
    global start_time

    stats = {}
    checktime = time.time()

    driller_inputs = 0
    # collect stats into dictionary
    for fuzzer_dir in os.listdir(sync_dir):
        stat_path = os.path.join(sync_dir, fuzzer_dir, "fuzzer_stats")
        if os.path.isfile(stat_path):
            stats[fuzzer_dir] = {}

            stat_blob = open(stat_path, "rb").read()
            stat_lines = stat_blob.split("\n")[:-1]
            for stat in stat_lines:
                key, val = stat.split(":")
                stats[fuzzer_dir][key.strip()] = val

        else: # could be driller
            if fuzzer_dir == "driller":
                drilled_inputs = len(os.listdir(os.path.join(sync_dir, fuzzer_dir, "queue")))

    pending_total = 0
    pending_favs  = 0
    execs_done = 0
    total_time = 0
    alive_cnt = 0
    crashes = 0
    for fuzzer in stats:
        # is the fuzzer alive
        fuzz_pid = int(stats[fuzzer]['fuzzer_pid'])
        try:
            os.kill(fuzz_pid, 0)
            alive_cnt += 1
        except OSError:
            continue

        time_running   = checktime - int(stats[fuzzer]['start_time'])
        total_time    += time_running
        pending_total += int(stats[fuzzer]['pending_total'])
        pending_favs  += int(stats[fuzzer]['pending_favs'])
        execs_done    += int(stats[fuzzer]['execs_done'])
        crashes       += int(stats[fuzzer]['unique_crashes'])

    m, s = divmod(checktime - start_time, 60)
    h, m = divmod(m, 60)
    print "  Run time           : %d:%02d:%02d" % (h, m, s)
    print "  Fuzzers Alive      : %d alive" % alive_cnt
    print "  Pending paths      : %d faves, %d total" % (pending_favs, pending_total)
    print "  Pending per fuzzer : %d faves, %d total" % (pending_favs / alive_cnt, pending_total / alive_cnt)
    print "  Cumulative speed   : %d execs/sec" % int((execs_done * alive_cnt) / total_time)
    print "  Drilled inputs     : %d inputs" % drilled_inputs

    cstr = "%d" % crashes
    if crashes > 0:
        cstr = termcolor.colored(cstr, "red", attrs=["bold"])

    print "  Crashes            : %s crashes" % cstr
    print
    print "=" * 40

    return crashes

def start(binary_path, in_dir, out_dir, afl_count, work_dir=None):
    global procs
    global start_time

    base = os.path.dirname(__file__)

    work_dir      = "." if work_dir is None else work_dir
    # start time
    start_time    = time.time()
    # the path to AFL capable of calling driller
    afl_path      = os.path.join(base, "driller-afl-fuzz")
    # the AFL build path for afl-qemu-trace-*
    afl_path_var  = os.path.join(base, "build", "afl")
    # path to the drill script
    driller_path  = os.path.join(base, "drill.py")
    # driller-qemu
    qemu_dir      = os.path.join(base, config.QEMU_DIR)
    # redis channel id
    channel_id    = os.path.basename(binary_path)
    # afl dictionary
    dict_file     = os.path.join(work_dir, "%s.dict" % channel_id)

    l.debug("afl_path: %s" % afl_path)
    l.debug("driller_path: %s" % driller_path)
    l.debug("AFL_PATH_ENV: %s" % afl_path_var)
    l.debug("channel_id: %s" % channel_id)

    # clear redis database
    clear_redis(channel_id)

    # look for a dictionary, if one doesn't exist create it with angr
    if not os.path.isfile(dict_file):
        l.info("creating a dictionary of string references found in the binary")
        if not create_dict(binary_path, dict_file):
            l.warning("failed to create dictionary, this can really impede on AFL's progress")
            dict_file = None

    # set environment variable for the AFL_PATH
    os.environ['AFL_PATH'] = afl_path_var

    eof_exit = False
    # test if the binary terminates on EOF
    if not terminates_on_eof(qemu_dir, binary_path):
        l.warning("binary doesn't terminate on EOF! attempting to use hack to fix this")
        eof_exit = True

    # spin up the AFL guys
    start_afl(afl_path, binary_path, in_dir, out_dir, afl_count,
            dictionary=dict_file, driller_path=driller_path, eof_exit=eof_exit)

    # spin up the redis listener
    procs.append(start_redis_listener(channel_id, out_dir))

    # setup signal handler
    signal.signal(signal.SIGINT, terminate)

    crash_found = False
    while not crash_found:
        time.sleep(config.CRASH_CHECK_INTERVAL)
        crash_found = bool(show_afl_stats(out_dir))

    report_crash_found(channel_id)
    kill_procs()

    return 0

def main():

    parser = argparse.ArgumentParser(description="Driller")

    parser.add_argument("-b", dest="binary",
                        type=str,
                        metavar="<binary>",
                        help="binary executable to drill",
                        required=True)

    parser.add_argument("-i", dest="in_dir",
                        type=str,
                        metavar="<in_dir>",
                        help="input directory",
                        required=True)

    parser.add_argument("-o", dest="out_dir",
                        type=str,
                        metavar="<out_dir>",
                        help="output directory",
                        required=True)

    parser.add_argument("-n", dest="afl_count",
                        type=int,
                        metavar="<afl_count>" ,
                        help="number of AFL instances to use",
                        required=True)

    args = parser.parse_args()

    binary_path  = args.binary
    in_dir       = args.in_dir
    out_dir      = args.out_dir
    afl_count    = args.afl_count

    return start(binary_path, in_dir, out_dir, afl_count)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Remove temporary file after creating it for EOF exit detection<commit_after>#!/usr/bin/env python

import angr
import driller.config as config

import argparse
import redis
import driller.tasks
import os
import subprocess
import signal
import sys
import multiprocessing
import cPickle as pickle
import time
import tempfile
import termcolor
import string

import logging

l = logging.getLogger("fuzz")
l.setLevel("INFO")

# global list of processes so we can kill them on SIGINT
procs = [ ]

start_time = 0

### EXIT HANDLERS

def kill_procs():
    map(lambda p: p.terminate(), procs)

def terminate(signal, frame):
    kill_procs()
    sys.exit(0)

### DICTIONARY CREATION

def hexescape(s):
    out = [ ]
    acceptable = string.letters + string.digits + " ."
    for c in s:
        if c not in acceptable:
            out.append("\\x%02x" % ord(c))
        else:
            out.append(c)

    return ''.join(out)

def create_dict(binary, outfile):
    b = angr.Project(binary)
    cfg = b.analyses.CFG(keep_input_state=True, enable_advanced_backward_slicing=True)

    string_references = [ ]
    for f in cfg.function_manager.functions.values():
        try:
            string_references.append(f.string_references())
        except ZeroDivisionError:
            pass

    string_references = sum(string_references, [])

    strings = [] if len(string_references) == 0 else zip(*string_references)[1]

    if len(strings) > 0:
        with open(outfile, 'wb') as f:
            for i, string in enumerate(strings):
                s = hexescape(string)
                f.write("driller_%d=\"%s\"\n" % (i, s))

        return True

    return False

### BEHAVIOR TESTING
def terminates_on_eof(qemu_dir, binary):
    l.info("attempting to detect if the binary doesn't terminate on EOF")

    # detect the binary type
    b = angr.Project(binary)
    ld_arch = b.loader.main_bin.arch
    ld_type = b.loader.main_bin.filetype

    if ld_type == "cgc":
        arch = "cgc"

    elif ld_type == "elf":
        if ld_arch == archinfo.arch_amd64.ArchAMD64:
            arch = "x86_64"
        if ld_arch == archinfo.arch_x86.ArchX86:
            arch = "i386"

    else:
        l.error("binary is of an unsupported architecture")
        raise NotImplementedError

    qemu_path = os.path.join(qemu_dir, "driller-qemu-%s" % arch)
    if not os.access(qemu_path, os.X_OK):
        l.error("either qemu does not exist in config.QEMU_DIR or it is not executable")
        return True

    # create a dumb test input
    fd, tinput = tempfile.mkstemp()
    os.close(fd)

    with open(tinput, 'wb') as f:
        f.write("fuzz")

    with open(tinput, 'rb') as i:
        with open('/dev/null', 'w') as o:
            args = [qemu_path, binary]
            p = subprocess.Popen(args, stdin=i, stdout=o)

            time.sleep(2) # generously give the binary two seconds to terminate

            if p.poll() is None: # it doesn't seem to terminate on EOF
                p.terminate()
                os.remove(tinput)
                return False

    os.remove(tinput)
    # good, it terminates on EOF, no monkey business
    return True

### AFL SPAWNERS

def start_afl_instance(afl_path, binary, in_dir, out_dir, fuzz_id, dictionary=None, memory="8G",
                driller=None, eof_termination=False):

    args = [afl_path]

    args += ["-i", in_dir]
    args += ["-o", out_dir]
    args += ["-m", memory]
    args += ["-Q"]
    if fuzz_id == 0:
        args += ["-M", "fuzzer-master"]
        outfile = "fuzzer-master.log"
    else:
        args += ["-S", "fuzzer-%d" % fuzz_id]
        outfile = "fuzzer-%d.log" % fuzz_id

    if dictionary is not None:
        args += ["-x", dictionary]

    if driller is not None:
        args += ["-D", driller]

    if eof_termination:
        args += ["-E"]

    args += ["--", binary]

    l.debug("execing: %s > %s" % (' '.join(args), outfile))

    # kind of gross hack to get the output directory
    work_dir = os.path.dirname(out_dir[:-1])

    outfile = os.path.join(work_dir, outfile)
    fp = open(outfile, "w")

    return subprocess.Popen(args, stdout=fp)

def start_afl(afl_path, binary_path, in_dir, out_dir, afl_count, driller_path, dictionary, eof_exit):
    global procs

    # spin up the master AFL instance
    master = start_afl_instance(afl_path,
                                binary_path,
                                in_dir,
                                out_dir,
                                fuzz_id=0, # the master fuzzer
                                dictionary=dictionary,
                                eof_termination=eof_exit)
    procs.append(master)

    if afl_count > 1:
        driller = start_afl_instance(afl_path,
                                     binary_path,
                                     in_dir,
                                     out_dir,
                                     1,
                                     driller=driller_path,
                                     dictionary=dictionary,
                                     eof_termination=eof_exit)
        procs.append(driller)

    else:
        l.warning("only one AFL instance was chosen to be spun up, driller will never be invoked")

    # only spins up an AFL instances if afl_count > 1
    for n in range(2, afl_count):
        slave = start_afl_instance(afl_path,
                                binary_path,
                                in_dir,
                                out_dir,
                                n,
                                dictionary=dictionary,
                                eof_termination=eof_exit)

        procs.append(slave)

### BACKEND HANDLERS

def clear_redis(identifier):
    redis_inst = redis.Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, db=config.REDIS_DB)

    # delete all the catalogue entries
    redis_inst.delete("%s-catalogue" % identifier)

    # delete all the traced entries
    redis_inst.delete("%s-traced" % identifier)

    # delete all the crash-found entry
    redis_inst.delete("%s-crash-found" % identifier)

def report_crash_found(identifier):
    redis_inst = redis.Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, db=config.REDIS_DB)

    # add True as a member
    redis_inst.sadd(identifier + "-crash-found", True)

def listen(queue_dir, channel):
    l.info("subscring to redis channel %s" % channel)
    l.info("new inputs will be placed into %s" % queue_dir)

    try:
        os.makedirs(queue_dir)
    except OSError:
        l.warning("could not create output directory '%s'" % queue_dir)

    redis_inst = redis.Redis(host=config.REDIS_HOST, port=config.REDIS_PORT, db=config.REDIS_DB)
    p = redis_inst.pubsub()

    p.subscribe(channel)

    input_cnt = 0

    for msg in p.listen():
        if msg['type'] == 'message':
            real_msg = pickle.loads(msg['data'])
            out_filename = "driller-%d-%x-%x" % real_msg['meta']
            l.info("dumping new input to %s" % out_filename)
            afl_name = "id:%06d,src:%s" % (input_cnt, out_filename)
            out_file = os.path.join(queue_dir, afl_name)

            with open(out_file, 'wb') as ofp:
                ofp.write(real_msg['data'])

            input_cnt += 1

def start_redis_listener(identifier, out_dir):
    driller_queue_dir = os.path.join(out_dir, "driller", "queue")
    channel = "%s-generated" % identifier

    p = multiprocessing.Process(target=listen, args=(driller_queue_dir, channel,))
    p.start()
    return p

### STATS

def show_afl_stats(sync_dir):
    global start_time

    stats = {}
    checktime = time.time()

    driller_inputs = 0
    # collect stats into dictionary
    for fuzzer_dir in os.listdir(sync_dir):
        stat_path = os.path.join(sync_dir, fuzzer_dir, "fuzzer_stats")
        if os.path.isfile(stat_path):
            stats[fuzzer_dir] = {}

            stat_blob = open(stat_path, "rb").read()
            stat_lines = stat_blob.split("\n")[:-1]
            for stat in stat_lines:
                key, val = stat.split(":")
                stats[fuzzer_dir][key.strip()] = val

        else: # could be driller
            if fuzzer_dir == "driller":
                drilled_inputs = len(os.listdir(os.path.join(sync_dir, fuzzer_dir, "queue")))

    pending_total = 0
    pending_favs  = 0
    execs_done = 0
    total_time = 0
    alive_cnt = 0
    crashes = 0
    for fuzzer in stats:
        # is the fuzzer alive
        fuzz_pid = int(stats[fuzzer]['fuzzer_pid'])
        try:
            os.kill(fuzz_pid, 0)
            alive_cnt += 1
        except OSError:
            continue

        time_running   = checktime - int(stats[fuzzer]['start_time'])
        total_time    += time_running
        pending_total += int(stats[fuzzer]['pending_total'])
        pending_favs  += int(stats[fuzzer]['pending_favs'])
        execs_done    += int(stats[fuzzer]['execs_done'])
        crashes       += int(stats[fuzzer]['unique_crashes'])

    m, s = divmod(checktime - start_time, 60)
    h, m = divmod(m, 60)
    print "  Run time           : %d:%02d:%02d" % (h, m, s)
    print "  Fuzzers Alive      : %d alive" % alive_cnt
    print "  Pending paths      : %d faves, %d total" % (pending_favs, pending_total)
    print "  Pending per fuzzer : %d faves, %d total" % (pending_favs / alive_cnt, pending_total / alive_cnt)
    print "  Cumulative speed   : %d execs/sec" % int((execs_done * alive_cnt) / total_time)
    print "  Drilled inputs     : %d inputs" % drilled_inputs

    cstr = "%d" % crashes
    if crashes > 0:
        cstr = termcolor.colored(cstr, "red", attrs=["bold"])

    print "  Crashes            : %s crashes" % cstr
    print
    print "=" * 40

    return crashes

def start(binary_path, in_dir, out_dir, afl_count, work_dir=None):
    global procs
    global start_time

    base = os.path.dirname(__file__)

    work_dir      = "." if work_dir is None else work_dir
    # start time
    start_time    = time.time()
    # the path to AFL capable of calling driller
    afl_path      = os.path.join(base, "driller-afl-fuzz")
    # the AFL build path for afl-qemu-trace-*
    afl_path_var  = os.path.join(base, "build", "afl")
    # path to the drill script
    driller_path  = os.path.join(base, "drill.py")
    # driller-qemu
    qemu_dir      = os.path.join(base, config.QEMU_DIR)
    # redis channel id
    channel_id    = os.path.basename(binary_path)
    # afl dictionary
    dict_file     = os.path.join(work_dir, "%s.dict" % channel_id)

    l.debug("afl_path: %s" % afl_path)
    l.debug("driller_path: %s" % driller_path)
    l.debug("AFL_PATH_ENV: %s" % afl_path_var)
    l.debug("channel_id: %s" % channel_id)

    # clear redis database
    clear_redis(channel_id)

    # look for a dictionary, if one doesn't exist create it with angr
    if not os.path.isfile(dict_file):
        l.info("creating a dictionary of string references found in the binary")
        if not create_dict(binary_path, dict_file):
            l.warning("failed to create dictionary, this can really impede on AFL's progress")
            dict_file = None

    # set environment variable for the AFL_PATH
    os.environ['AFL_PATH'] = afl_path_var

    eof_exit = False
    # test if the binary terminates on EOF
    if not terminates_on_eof(qemu_dir, binary_path):
        l.warning("binary doesn't terminate on EOF! attempting to use hack to fix this")
        eof_exit = True

    # spin up the AFL guys
    start_afl(afl_path, binary_path, in_dir, out_dir, afl_count,
            dictionary=dict_file, driller_path=driller_path, eof_exit=eof_exit)

    # spin up the redis listener
    procs.append(start_redis_listener(channel_id, out_dir))

    # setup signal handler
    signal.signal(signal.SIGINT, terminate)

    crash_found = False
    while not crash_found:
        time.sleep(config.CRASH_CHECK_INTERVAL)
        crash_found = bool(show_afl_stats(out_dir))

    report_crash_found(channel_id)
    kill_procs()

    return 0

def main():

    parser = argparse.ArgumentParser(description="Driller")

    parser.add_argument("-b", dest="binary",
                        type=str,
                        metavar="<binary>",
                        help="binary executable to drill",
                        required=True)

    parser.add_argument("-i", dest="in_dir",
                        type=str,
                        metavar="<in_dir>",
                        help="input directory",
                        required=True)

    parser.add_argument("-o", dest="out_dir",
                        type=str,
                        metavar="<out_dir>",
                        help="output directory",
                        required=True)

    parser.add_argument("-n", dest="afl_count",
                        type=int,
                        metavar="<afl_count>" ,
                        help="number of AFL instances to use",
                        required=True)

    args = parser.parse_args()

    binary_path  = args.binary
    in_dir       = args.in_dir
    out_dir      = args.out_dir
    afl_count    = args.afl_count

    return start(binary_path, in_dir, out_dir, afl_count)

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