<commit_before># -*- coding: utf-8 -*-
'''
Device worker process
'''
from __future__ import absolute_import

# Import python stdlib
import os
import re
import time
import signal
import logging
import threading
from datetime import datetime

# Import thrid party libs
import zmq
import umsgpack

# Import napalm-logs pkgs
import napalm_logs.utils
from napalm_logs.proc import NapalmLogsProc
from napalm_logs.config import PUB_IPC_URL
from napalm_logs.config import DEV_IPC_URL_TPL
from napalm_logs.config import REPLACEMENTS
from napalm_logs.config import UNKNOWN_DEVICE_NAME
# exceptions
from napalm_logs.exceptions import NapalmLogsExit

log = logging.getLogger(__name__)

class NapalmLogsDeviceProc(NapalmLogsProc):
    '''
    Device sub-process class.
    '''
    def __init__(self, name, config, pipe, pub_pipe, publisher_opts):
        self._name = name
        self.pipe = pipe
        self._config = config
        self.pub_pipe = pub_pipe
        self.publisher_opts = publisher_opts
        self.__up = False
        self.compiled_messages = None
        self._compile_messages()

    def _exit_gracefully(self, signum, _):
        log.debug('Caught signal in %s device process', self._name)
        self.stop()

    def _setup_ipc(self):
        '''
        Subscribe to the right topic
        in the device IPC and publish to the
        publisher proxy.
        '''
        ctx = zmq.Context()
        # subscribe to device IPC
        self.sub = ctx.socket(zmq.PULL)
        # subscribe to the corresponding IPC pipe
        ipc_url = DEV_IPC_URL_TPL.format(os=self._name)
        self.sub.connect(ipc_url)
        # self.sub.setsockopt(zmq.SUBSCRIBE, '')
        # publish to the publisher IPC
        self.pub = ctx.socket(zmq.PUB)
        self.pub.connect(PUB_IPC_URL)

    def _compile_messages(self):
        '''
        Create a list of all OS messages and their compiled regexs
        '''
        self.compiled_messages = []
        if not self._config:
            return
        for message_dict in self._config.get('messages', {}):
            error = message_dict['error']
            tag = message_dict['tag']
            model = message_dict['model']
            match_on = message_dict.get('match_on', 'tag')
            if '__python_fun__' in message_dict:
                self.compiled_messages.append({
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'model': model,
                    '__python_fun__': message_dict['__python_fun__']
                })
                continue
            values = message_dict['values']
            replace = message_dict['replace']
            line = message_dict['line']
            mapping = message_dict['mapping']
            # We will now figure out which position each value is in so we can use it with the match statement
            position = {}
            for key in values.keys():
                position[line.find('{' + key + '}')] = key
            sorted_position = {}
            for i, elem in enumerate(sorted(position.items())):
                sorted_position[elem[1]] = i + 1
            # Escape the line, then remove the escape for the curly bracets so they can be used when formatting
            escaped = re.escape(line).replace('\{', '{').replace('\}', '}')
            # Replace a whitespace with \s+
            escaped = escaped.replace('\ ', '\s+')
            self.compiled_messages.append(
                {
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'line': re.compile(escaped.format(**values)),
                    'positions': sorted_position,
                    'values': values,
                    'replace': replace,
                    'model': model,
                    'mapping': mapping
                }
            )

    def _parse(self, msg_dict):
        '''
        Parse a syslog message and check what OpenConfig object should
        be generated.
        '''
        error_present = False
        for message in self.compiled_messages:
            match_on = message['match_on']
            if message[match_on] != msg_dict[match_on]:
                continue
            if '__python_fun__' in message:
                return {
                    'model': message['model'],
                    'error': message['error'],
                    '__python_fun__': message['__python_fun__']
                }
            error_present = True
            match = message['line'].search(msg_dict['message'])
            if not match:
                continue
            positions = message.get('positions', {})
            values = message.get('values')
            ret = {
                'model': message['model'],
                'mapping': message['mapping'],
                'replace': message['replace'],
                'error': message['error']
            }
            for key in values.keys():
                # Check if the value needs to be replaced
                if message['replace'].get(key):
                    fun = REPLACEMENTS.get(message['replace'].get(key))
                    result = fun(match.group(positions.get(key)))
                else:
                    result = match.group(positions.get(key))
                ret[key] = result
            return ret
        if error_present is True:
            log.info('Configured regex did not match for os: %s tag %s', self._name, msg_dict.get('tag', ''))
        else:
            log.info('Syslog message not configured for os: %s tag %s', self._name, msg_dict.get('tag', ''))

    def _emit(self, **kwargs):
        '''
        Emit an OpenConfig object given a certain combination of
        fields mappeed in the config to the corresponding hierarchy.
        '''
        oc_dict = {}
        for mapping, result_key in kwargs['mapping']['variables'].items():
            result = kwargs[result_key]
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)
        for mapping, result in kwargs['mapping']['static'].items():
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)

        return oc_dict

    def _publish(self, obj):
        '''
        Publish the OC object.
        '''
        bin_obj = umsgpack.packb(obj)
        self.pub.send(bin_obj)

    def _format_time(self, time, date, prefix_id):
        # TODO can we work out the time format from the regex? Probably but this is a task for another day
        time_format = self._config['prefixes'][prefix_id].get('time_format', '')
        if not time or not date or not time_format:
            return int(datetime.now().strftime('%s'))
        # Most syslog do not include the year, so we will add the current year if we are not supplied with one
        if '%y' in time_format or '%Y' in time_format:
            timestamp = datetime.strptime('{} {}'.format(date, time), time_format)
        else:
            year = datetime.now().year
            timestamp = datetime.strptime('{} {} {}'.format(year, date, time), '%Y {}'.format(time_format))
        return int(timestamp.strftime('%s'))

    def start(self):
        '''
        Start the worker process.
        '''
        # self._setup_ipc()
        # Start suicide polling thread
        thread = threading.Thread(target=self._suicide_when_without_parent, args=(os.getppid(),))
        thread.start()
        signal.signal(signal.SIGTERM, self._exit_gracefully)
        self.__up = True
        while self.__up:
            # bin_obj = self.sub.recv()
            # msg_dict, address = umsgpack.unpackb(bin_obj, use_list=False)
            try:
                msg_dict, address = self.pipe.recv()
            except IOError as error:
                if self.__up is False:
                    return
                else:
                    msg = 'Received IOError on {} device pipe: {}'.format(self._name, error)
                    log.error(msg, exc_info=True)
                    raise NapalmLogsExit(msg)
            log.debug('%s: dequeued %s, received from %s', self._name, msg_dict, address)
            if self._name == UNKNOWN_DEVICE_NAME:
                # If running in the sub-process publishing messages for unknown OSs.
                # This will always send what receives, as-is.
                log.debug('Publishing the message as-is, no further processing possible.')
                to_publish = {
                    'ip': address,
                    'host': 'unknown',
                    'timestamp': int(time.time()),
                    'message_details': msg_dict,
                    'os': UNKNOWN_DEVICE_NAME,
                    'error': 'UNKNOWN',
                    'model_name': 'unknown'
                }
                log.debug('Queueing to be published:')
                log.debug(to_publish)
                self.pub_pipe.send(to_publish)
                continue
            # From here on, we're running in a regular OS sub-process.
            host = msg_dict.get('host')
            prefix_id = msg_dict.pop('__prefix_id__')
            if 'timestamp' in msg_dict:
                timestamp = msg_dict['timestamp']
            else:
                timestamp = self._format_time(msg_dict.get('time', ''),
                                              msg_dict.get('date', ''),
                                              prefix_id)
            facility = msg_dict.get('facility')
            severity = msg_dict.get('severity')

            kwargs = self._parse(msg_dict)
            if not kwargs:
                # Unable to identify what model to generate for the message in cause.
                if self.publisher_opts.get('send_raw'):
                    # But publish the message when the user requested to push raw messages.
                    to_publish = {
                        'ip': address,
                        'host': host,
                        'timestamp': timestamp,
                        'message_details': msg_dict,
                        'os': self._name,
                        'error': 'RAW',
                        'model_name': 'raw',
                        'facility': facility,
                        'severity': severity
                    }
                    log.debug('Queueing to be published:')
                    log.debug(to_publish)
                    self.pub_pipe.send(to_publish)
                continue
            try:
                if '__python_fun__' in kwargs:
                    log.debug('Using the Python parser to determine the YANG-equivalent object')
                    yang_obj = kwargs['__python_fun__'](msg_dict)
                else:
                    yang_obj = self._emit(**kwargs)
            except Exception:
                log.exception('Unexpected error when generating the OC object.', exc_info=True)
                continue
            log.debug('Generated OC object:')
            log.debug(yang_obj)
            error = kwargs.get('error')
            model_name = kwargs.get('model')
            to_publish = {
                'error': error,
                'host': host,
                'ip': address,
                'timestamp': timestamp,
                'yang_message': yang_obj,
                'message_details': msg_dict,
                'yang_model': model_name,
                'os': self._name,
                'facility': facility,
                'severity': severity
            }
            log.debug('Queueing to be published:')
            log.debug(to_publish)
            self.pub_pipe.send(to_publish)
            # self._publish(to_publish)

    def stop(self):
        '''
        Stop the worker process.
        '''
        log.info('Stopping %s device process', self._name)
        self.__up = False
        self.pipe.close()
        self.pub_pipe.close()
<commit_msg>Add more tests<commit_after># -*- coding: utf-8 -*-
'''
Device worker process
'''
from __future__ import absolute_import

# Import python stdlib
import os
import re
import time
import signal
import logging
import threading
from datetime import datetime

# Import thrid party libs
import zmq
import umsgpack

# Import napalm-logs pkgs
import napalm_logs.utils
from napalm_logs.proc import NapalmLogsProc
from napalm_logs.config import PUB_IPC_URL
from napalm_logs.config import DEV_IPC_URL_TPL
from napalm_logs.config import REPLACEMENTS
from napalm_logs.config import UNKNOWN_DEVICE_NAME
# exceptions
from napalm_logs.exceptions import NapalmLogsExit

log = logging.getLogger(__name__)

class NapalmLogsDeviceProc(NapalmLogsProc):
    '''
    Device sub-process class.
    '''
    def __init__(self, name, config, pipe, pub_pipe, publisher_opts):
        self._name = name
        self.pipe = pipe
        self._config = config
        self.pub_pipe = pub_pipe
        self.publisher_opts = publisher_opts
        self.__up = False
        self.compiled_messages = None
        self._compile_messages()

    def _exit_gracefully(self, signum, _):
        log.debug('Caught signal in %s device process', self._name)
        self.stop()

    def _setup_ipc(self):
        '''
        Subscribe to the right topic
        in the device IPC and publish to the
        publisher proxy.
        '''
        ctx = zmq.Context()
        # subscribe to device IPC
        self.sub = ctx.socket(zmq.PULL)
        # subscribe to the corresponding IPC pipe
        ipc_url = DEV_IPC_URL_TPL.format(os=self._name)
        self.sub.connect(ipc_url)
        # self.sub.setsockopt(zmq.SUBSCRIBE, '')
        # publish to the publisher IPC
        self.pub = ctx.socket(zmq.PUB)
        self.pub.connect(PUB_IPC_URL)

    def _compile_messages(self):
        '''
        Create a list of all OS messages and their compiled regexs
        '''
        self.compiled_messages = []
        if not self._config:
            return
        for message_dict in self._config.get('messages', {}):
            error = message_dict['error']
            tag = message_dict['tag']
            model = message_dict['model']
            match_on = message_dict.get('match_on', 'tag')
            if '__python_fun__' in message_dict:
                self.compiled_messages.append({
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'model': model,
                    '__python_fun__': message_dict['__python_fun__']
                })
                continue
            values = message_dict['values']
            replace = message_dict['replace']
            line = message_dict['line']
            mapping = message_dict['mapping']
            # We will now figure out which position each value is in so we can use it with the match statement
            position = {}
            for key in values.keys():
                position[line.find('{' + key + '}')] = key
            sorted_position = {}
            for i, elem in enumerate(sorted(position.items())):
                sorted_position[elem[1]] = i + 1
            # Escape the line, then remove the escape for the curly bracets so they can be used when formatting
            escaped = re.escape(line).replace('\{', '{').replace('\}', '}')
            # Replace a whitespace with \s+
            escaped = escaped.replace('\ ', '\s+')
            self.compiled_messages.append(
                {
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'line': re.compile(escaped.format(**values)),
                    'positions': sorted_position,
                    'values': values,
                    'replace': replace,
                    'model': model,
                    'mapping': mapping
                }
            )

    def _parse(self, msg_dict):
        '''
        Parse a syslog message and check what OpenConfig object should
        be generated.
        '''
        error_present = False
        for message in self.compiled_messages:
            match_on = message['match_on']
            if message[match_on] != msg_dict[match_on]:
                continue
            if '__python_fun__' in message:
                return {
                    'model': message['model'],
                    'error': message['error'],
                    '__python_fun__': message['__python_fun__']
                }
            error_present = True
            match = message['line'].search(msg_dict['message'])
            if not match:
                continue
            positions = message.get('positions', {})
            values = message.get('values')
            ret = {
                'model': message['model'],
                'mapping': message['mapping'],
                'replace': message['replace'],
                'error': message['error']
            }
            for key in values.keys():
                # Check if the value needs to be replaced
                if message['replace'].get(key):
                    fun = REPLACEMENTS.get(message['replace'].get(key))
                    result = fun(match.group(positions.get(key)))
                else:
                    result = match.group(positions.get(key))
                ret[key] = result
            return ret
        if error_present is True:
            log.info('Configured regex did not match for os: %s tag %s', self._name, msg_dict.get('tag', ''))
        else:
            log.info('Syslog message not configured for os: %s tag %s', self._name, msg_dict.get('tag', ''))

    def _emit(self, **kwargs):
        '''
        Emit an OpenConfig object given a certain combination of
        fields mappeed in the config to the corresponding hierarchy.
        '''
        oc_dict = {}
        for mapping, result_key in kwargs['mapping']['variables'].items():
            result = kwargs[result_key]
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)
        for mapping, result in kwargs['mapping']['static'].items():
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)

        return oc_dict

    def _publish(self, obj):
        '''
        Publish the OC object.
        '''
        bin_obj = umsgpack.packb(obj)
        self.pub.send(bin_obj)

    def _format_time(self, time, date, prefix_id):
        # TODO can we work out the time format from the regex? Probably but this is a task for another day
        time_format = self._config['prefixes'][prefix_id].get('time_format', '')
        if not time or not date or not time_format:
            return int(datetime.now().strftime('%s'))
        # Most syslog do not include the year, so we will add the current year if we are not supplied with one
        if '%y' in time_format or '%Y' in time_format:
            timestamp = datetime.strptime('{} {}'.format(date, time), time_format)
        else:
            year = datetime.now().year
            timestamp = datetime.strptime('{} {} {}'.format(year, date, time), '%Y {}'.format(time_format))
        return int(timestamp.strftime('%s'))

    def start(self):
        '''
        Start the worker process.
        '''
        # self._setup_ipc()
        # Start suicide polling thread
        thread = threading.Thread(target=self._suicide_when_without_parent, args=(os.getppid(),))
        thread.start()
        signal.signal(signal.SIGTERM, self._exit_gracefully)
        self.__up = True
        while self.__up:
            # bin_obj = self.sub.recv()
            # msg_dict, address = umsgpack.unpackb(bin_obj, use_list=False)
            try:
                msg_dict, address = self.pipe.recv()
            except IOError as error:
                if self.__up is False:
                    return
                else:
                    msg = 'Received IOError on {} device pipe: {}'.format(self._name, error)
                    log.error(msg, exc_info=True)
                    raise NapalmLogsExit(msg)
            log.debug('%s: dequeued %s, received from %s', self._name, msg_dict, address)
            if self._name == UNKNOWN_DEVICE_NAME:
                # If running in the sub-process publishing messages for unknown OSs.
                # This will always send what receives, as-is.
                log.debug('Publishing the message as-is, no further processing possible.')
                to_publish = {
                    'ip': address,
                    'host': 'unknown',
                    'timestamp': int(time.time()),
                    'message_details': msg_dict,
                    'os': UNKNOWN_DEVICE_NAME,
                    'error': 'UNKNOWN',
                    'model_name': 'unknown'
                }
                log.debug('Queueing to be published:')
                log.debug(to_publish)
                self.pub_pipe.send(to_publish)
                continue
            # From here on, we're running in a regular OS sub-process.
            host = msg_dict.get('host')
            prefix_id = msg_dict.pop('__prefix_id__')
            if 'timestamp' in msg_dict:
                timestamp = msg_dict.pop('timestamp')
            else:
                timestamp = self._format_time(msg_dict.get('time', ''),
                                              msg_dict.get('date', ''),
                                              prefix_id)
            facility = msg_dict.get('facility')
            severity = msg_dict.get('severity')

            kwargs = self._parse(msg_dict)
            if not kwargs:
                # Unable to identify what model to generate for the message in cause.
                if self.publisher_opts.get('send_raw'):
                    # But publish the message when the user requested to push raw messages.
                    to_publish = {
                        'ip': address,
                        'host': host,
                        'timestamp': timestamp,
                        'message_details': msg_dict,
                        'os': self._name,
                        'error': 'RAW',
                        'model_name': 'raw',
                        'facility': facility,
                        'severity': severity
                    }
                    log.debug('Queueing to be published:')
                    log.debug(to_publish)
                    self.pub_pipe.send(to_publish)
                continue
            try:
                if '__python_fun__' in kwargs:
                    log.debug('Using the Python parser to determine the YANG-equivalent object')
                    yang_obj = kwargs['__python_fun__'](msg_dict)
                else:
                    yang_obj = self._emit(**kwargs)
            except Exception:
                log.exception('Unexpected error when generating the OC object.', exc_info=True)
                continue
            log.debug('Generated OC object:')
            log.debug(yang_obj)
            error = kwargs.get('error')
            model_name = kwargs.get('model')
            to_publish = {
                'error': error,
                'host': host,
                'ip': address,
                'timestamp': timestamp,
                'yang_message': yang_obj,
                'message_details': msg_dict,
                'yang_model': model_name,
                'os': self._name,
                'facility': facility,
                'severity': severity
            }
            log.debug('Queueing to be published:')
            log.debug(to_publish)
            self.pub_pipe.send(to_publish)
            # self._publish(to_publish)

    def stop(self):
        '''
        Stop the worker process.
        '''
        log.info('Stopping %s device process', self._name)
        self.__up = False
        self.pipe.close()
        self.pub_pipe.close()
<|endoftext|>


<commit_before>class IMap:
    """
    Base class for IMap.
    IMap is a set of Conditional Independence (eg: "X is independent of Y given Z" where X, Y and Z
    are random variables) or Independence(eg: "X is independent of Y" where X and Y
    are random variables) assertions.

    Public Methods
    --------------
    add_assertions
    get_imap
    """
    def __init__(self, *assertions):
        """
        Initialize the IMap with Conditional Independence or Independence assertions.

        Parameters
        ----------
        assertions: Lists or Tuples
                Each assertion is a list or tuple of the form: [variable, independent_of and given]
                eg: assertion ['X', 'Y', 'Z'] would be X is independent of Y given Z.

        Examples
        --------
        Creating an IMap object with one independence assertion: Random Variable X is independent of Y
        >>> imap = IMap(['X', 'Y'])

        Creating an IMap object with three conditional independence assertions:
        First assertion is Random Variable X is independent of Y given Z.
        >>> imap = IMap(['X', 'Y', 'Z'],
        >>>             ['a', ['b', 'c'], 'd'],
        >>>             ['l', ['m', 'n'], 'o'])
        """
        self.imap = set()
        for assertion in assertions:
            self.imap.add(IndependenceAssertion(assertion[0], assertion[1], assertion[2]))

    def get_imap(self):
        """
        Returns the imap which is a set of IndependenceAssertion objects.

        See Also
        --------
        add_assertions

        Examples
        --------
        >>> imap = IMap(['X', 'Y', 'Z'])
        >>> imap.get_imap()
        """
        return self.imap

    def add_assertions(self, *assertions):
        """
        Adds assertions to imap.

        Parameters
        ----------
        assertions: Lists or Tuples
                Each assertion is a list or tuple of variable, independent_of and given.

        Examples
        --------
        >>> imap = IMap()
        >>> imap.add_assertions(['X', 'Y', 'Z'])
        >>> imap.add_assertions(['a', ['b', 'c'], 'd'])
        """
        for assertion in assertions:
            self.imap.add(IndependenceAssertion(assertion[0], assertion[1], assertion[2]))

class IndependenceAssertion:
    """
    Represents Conditional Independence or Independence assertion.

    Each assertion has 3 attributes: variable, independent_of, given.
    The attributes for (U || X, Y | Z) read as
                         ---
    Random Variable U is independent of X and Y given Z would be:

    variable = {U}
    independent_of = {X, Y}
    given = {Z}

    Public Methods
    --------------
    get_assertion
    set_assertion
    """
    def __init__(self, variable=[], independent_of=[], given=[]):
        """
        Initialize an IndependenceAssertion object with variable, independent_of and given attributes.

        Parameters
        ----------

             independent_of
                  ^
      variable   /   given
         ^      /     ^
         |     /      |
        (U || X, Y | Z) read as Random variable U is independent of X and Y given Z.
          ---

        variable: String or List
                Random Variable which is independent.

        independent_of: String or list of strings.
                Random Variables from which variable is independent

        given: String or list of strings.
                Random Variables given which variable is independent of independent_of.

        Examples
        --------
        >>> assertion = IndependenceAssertion('U', 'X')
        >>> assertion = IndependenceAssertion('U', ['X', 'Y'])
        >>> assertion = IndependenceAssertion('U', ['X', 'Y'], 'Z')
        >>> assertion = IndependenceAssertion('U', ['X', 'Y'], ['Z', 'A'])
        """
        if variable and not independent_of:
            raise Exception.Needed('independent_of')
        if any([independent_of, given]) and not variable:
            raise Exception.Needed('variable')
        if given and not all([variable, independent_of]):
            raise Exception.Needed('variable' if not variable else 'independent_of')

        self.variable = set(self._return_list_if_str(variable))
        self.independent_of = set(self._return_list_if_str(independent_of))
        self.given = set(self._return_list_if_str(given))

    @staticmethod
    def _return_list_if_str(variable):
        """
        If variable is a string returns a list containing variable.
        Else returns variable itself.
        """
        if isinstance(variable, str):
            return [variable]
        else:
            return variable

    def get_assertion(self):
        """
        Returns a tuple of the attributes: variable, independent_of, given.

        See Also
        --------
        set_assertion

        Examples
        --------
        >>> asser = IndependenceAssertion('X', 'Y', 'Z')
        >>> asser.get_assertion()
        """
        return self.variable, self.independent_of, self.given

    def set_assertion(self, variable, independent_of, given=[]):
        """
        Sets the attributes variable, independent_of and given.

        Parameters
        ----------

             independent_of
                  ^
      variable   /   given
         ^      /     ^
         |     /      |
        (U || X, Y | Z) read as Random variable U is independent of X and Y given Z.
          ---

        variable: String or List
                Random Variable which is independent.

        independent_of: String or list of strings.
                Random Variables from which variable is independent

        given: String or list of strings.
                Random Variables given which variable is independent of independent_of.

        See Also
        --------
        get_assertion

        Example
        -------
        For a random variable U independent of X and Y given Z, the function should be called as
        >>> asser = IndependenceAssertion()
        >>> asser.set_assertion('U', ['X', 'Y'], 'Z')
        >>> asser.set_assertion('U', ['X', 'Y'], ['Z', 'A'])
        """
        self.__init__(variable, independent_of, given)<commit_msg>edit doc test<commit_after>class IMap:
    """
    Base class for IMap.
    IMap is a set of Conditional Independence (eg: "X is independent of Y given Z" where X, Y and Z
    are random variables) or Independence(eg: "X is independent of Y" where X and Y
    are random variables) assertions.

    Public Methods
    --------------
    add_assertions
    get_imap
    """
    def __init__(self, *assertions):
        """
        Initialize the IMap with Conditional Independence or Independence assertions.

        Parameters
        ----------
        assertions: Lists or Tuples
                Each assertion is a list or tuple of the form: [variable, independent_of and given]
                eg: assertion ['X', 'Y', 'Z'] would be X is independent of Y given Z.

        Examples
        --------
        Creating an IMap object with one independence assertion: Random Variable X is independent of Y
        >>> imap = IMap(['X', 'Y'])

        Creating an IMap object with three conditional independence assertions:
        First assertion is Random Variable X is independent of Y given Z.
        >>> imap = IMap(['X', 'Y', 'Z'],
        ...             ['a', ['b', 'c'], 'd'],
        ...             ['l', ['m', 'n'], 'o'])
        """
        self.imap = set()
        for assertion in assertions:
            self.imap.add(IndependenceAssertion(assertion[0], assertion[1], assertion[2]))

    def get_imap(self):
        """
        Returns the imap which is a set of IndependenceAssertion objects.

        See Also
        --------
        add_assertions

        Examples
        --------
        >>> imap = IMap(['X', 'Y', 'Z'])
        >>> imap.get_imap()
        """
        return self.imap

    def add_assertions(self, *assertions):
        """
        Adds assertions to imap.

        Parameters
        ----------
        assertions: Lists or Tuples
                Each assertion is a list or tuple of variable, independent_of and given.

        Examples
        --------
        >>> imap = IMap()
        >>> imap.add_assertions(['X', 'Y', 'Z'])
        >>> imap.add_assertions(['a', ['b', 'c'], 'd'])
        """
        for assertion in assertions:
            self.imap.add(IndependenceAssertion(assertion[0], assertion[1], assertion[2]))

class IndependenceAssertion:
    """
    Represents Conditional Independence or Independence assertion.

    Each assertion has 3 attributes: variable, independent_of, given.
    The attributes for (U || X, Y | Z) read as
                         ---
    Random Variable U is independent of X and Y given Z would be:

    variable = {U}
    independent_of = {X, Y}
    given = {Z}

    Public Methods
    --------------
    get_assertion
    set_assertion
    """
    def __init__(self, variable=[], independent_of=[], given=[]):
        """
        Initialize an IndependenceAssertion object with variable, independent_of and given attributes.

        Parameters
        ----------

             independent_of
                  ^
      variable   /   given
         ^      /     ^
         |     /      |
        (U || X, Y | Z) read as Random variable U is independent of X and Y given Z.
          ---

        variable: String or List
                Random Variable which is independent.

        independent_of: String or list of strings.
                Random Variables from which variable is independent

        given: String or list of strings.
                Random Variables given which variable is independent of independent_of.

        Examples
        --------
        >>> assertion = IndependenceAssertion('U', 'X')
        >>> assertion = IndependenceAssertion('U', ['X', 'Y'])
        >>> assertion = IndependenceAssertion('U', ['X', 'Y'], 'Z')
        >>> assertion = IndependenceAssertion('U', ['X', 'Y'], ['Z', 'A'])
        """
        if variable and not independent_of:
            raise Exception.Needed('independent_of')
        if any([independent_of, given]) and not variable:
            raise Exception.Needed('variable')
        if given and not all([variable, independent_of]):
            raise Exception.Needed('variable' if not variable else 'independent_of')

        self.variable = set(self._return_list_if_str(variable))
        self.independent_of = set(self._return_list_if_str(independent_of))
        self.given = set(self._return_list_if_str(given))

    @staticmethod
    def _return_list_if_str(variable):
        """
        If variable is a string returns a list containing variable.
        Else returns variable itself.
        """
        if isinstance(variable, str):
            return [variable]
        else:
            return variable

    def get_assertion(self):
        """
        Returns a tuple of the attributes: variable, independent_of, given.

        See Also
        --------
        set_assertion

        Examples
        --------
        >>> asser = IndependenceAssertion('X', 'Y', 'Z')
        >>> asser.get_assertion()
        """
        return self.variable, self.independent_of, self.given

    def set_assertion(self, variable, independent_of, given=[]):
        """
        Sets the attributes variable, independent_of and given.

        Parameters
        ----------

             independent_of
                  ^
      variable   /   given
         ^      /     ^
         |     /      |
        (U || X, Y | Z) read as Random variable U is independent of X and Y given Z.
          ---

        variable: String or List
                Random Variable which is independent.

        independent_of: String or list of strings.
                Random Variables from which variable is independent

        given: String or list of strings.
                Random Variables given which variable is independent of independent_of.

        See Also
        --------
        get_assertion

        Example
        -------
        For a random variable U independent of X and Y given Z, the function should be called as
        >>> asser = IndependenceAssertion()
        >>> asser.set_assertion('U', ['X', 'Y'], 'Z')
        >>> asser.set_assertion('U', ['X', 'Y'], ['Z', 'A'])
        """
        self.__init__(variable, independent_of, given)<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
Device worker process
'''
from __future__ import absolute_import

# Import python stdlib
import os
import re
import time
import signal
import logging
import threading
from datetime import datetime

# Import thrid party libs
import zmq
import umsgpack

# Import napalm-logs pkgs
import napalm_logs.utils
from napalm_logs.proc import NapalmLogsProc
from napalm_logs.config import PUB_IPC_URL
from napalm_logs.config import DEV_IPC_URL
from napalm_logs.config import UNKNOWN_DEVICE_NAME
# exceptions
from napalm_logs.exceptions import NapalmLogsExit

log = logging.getLogger(__name__)

class NapalmLogsDeviceProc(NapalmLogsProc):
    '''
    Device sub-process class.
    '''
    def __init__(self, name, config, pipe, pub_pipe, publisher_opts):
        self._name = name
        self.pipe = pipe
        self._config = config
        self.pub_pipe = pub_pipe
        self.publisher_opts = publisher_opts
        self.__up = False
        self.compiled_messages = None
        self._compile_messages()

    def _exit_gracefully(self, signum, _):
        log.debug('Caught signal in %s device process', self._name)
        self.stop()

    def _setup_ipc(self):
        '''
        Subscribe to the right topic
        in the device IPC and publish to the
        publisher proxy.
        '''
        self.ctx = zmq.Context()
        # subscribe to device IPC
        log.debug('Creating the dealer IPC for %s', self._name)
        self.sub = self.ctx.socket(zmq.DEALER)
        self.sub.setsockopt(zmq.IDENTITY, self._name)
        # subscribe to the corresponding IPC pipe
        self.sub.connect(DEV_IPC_URL)
        # self.sub.setsockopt(zmq.SUBSCRIBE, '')
        # publish to the publisher IPC
        self.pub = self.ctx.socket(zmq.PUSH)
        self.pub.connect(PUB_IPC_URL)

    def _compile_messages(self):
        '''
        Create a list of all OS messages and their compiled regexs
        '''
        self.compiled_messages = []
        if not self._config:
            return
        for message_dict in self._config.get('messages', {}):
            error = message_dict['error']
            tag = message_dict['tag']
            model = message_dict['model']
            match_on = message_dict.get('match_on', 'tag')
            if '__python_fun__' in message_dict:
                self.compiled_messages.append({
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'model': model,
                    '__python_fun__': message_dict['__python_fun__']
                })
                continue
            values = message_dict['values']
            line = message_dict['line']
            mapping = message_dict['mapping']
            # We will now figure out which position each value is in so we can use it with the match statement
            position = {}
            replace = {}
            for key in values.keys():
                if '|' in key:
                    new_key, replace[new_key] = key.replace(' ', '').split('|')
                    values[new_key] = values.pop(key)
                    key = new_key
                position[line.find('{' + key + '}')] = key
            sorted_position = {}
            for i, elem in enumerate(sorted(position.items())):
                sorted_position[elem[1]] = i + 1
            # Escape the line, then remove the escape for the curly bracets so they can be used when formatting
            escaped = re.escape(line).replace('\{', '{').replace('\}', '}')
            # Replace a whitespace with \s+
            escaped = escaped.replace('\ ', '\s+')
            self.compiled_messages.append(
                {
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'line': re.compile(escaped.format(**values)),
                    'positions': sorted_position,
                    'values': values,
                    'replace': replace,
                    'model': model,
                    'mapping': mapping
                }
            )
        log.debug('Compiled messages:')
        log.debug(self.compiled_messages)

    def _parse(self, msg_dict):
        '''
        Parse a syslog message and check what OpenConfig object should
        be generated.
        '''
        error_present = False
        # log.debug('Matching the message:')
        # log.debug(msg_dict)
        for message in self.compiled_messages:
            # log.debug('Matching using:')
            # log.debug(message)
            match_on = message['match_on']
            if match_on not in msg_dict:
                # log.debug('%s is not a valid key in the partially parsed dict', match_on)
                continue
            if message['tag'] != msg_dict[match_on]:
                continue
            if '__python_fun__' in message:
                return {
                    'model': message['model'],
                    'error': message['error'],
                    '__python_fun__': message['__python_fun__']
                }
            error_present = True
            match = message['line'].search(msg_dict['message'])
            if not match:
                continue
            positions = message.get('positions', {})
            values = message.get('values')
            ret = {
                'model': message['model'],
                'mapping': message['mapping'],
                'replace': message['replace'],
                'error': message['error']
            }
            for key in values.keys():
                # Check if the value needs to be replaced
                if key in message['replace']:
                    result = napalm_logs.utils.cast(match.group(positions.get(key)), message['replace'][key])
                else:
                    result = match.group(positions.get(key))
                ret[key] = result
            return ret
        if error_present is True:
            log.info('Configured regex did not match for os: %s tag %s', self._name, msg_dict.get('tag', ''))
        else:
            log.info('Syslog message not configured for os: %s tag %s', self._name, msg_dict.get('tag', ''))

    def _emit(self, **kwargs):
        '''
        Emit an OpenConfig object given a certain combination of
        fields mappeed in the config to the corresponding hierarchy.
        '''
        oc_dict = {}
        for mapping, result_key in kwargs['mapping']['variables'].items():
            result = kwargs[result_key]
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)
        for mapping, result in kwargs['mapping']['static'].items():
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)

        return oc_dict

    def _publish(self, obj):
        '''
        Publish the OC object.
        '''
        bin_obj = umsgpack.packb(obj)
        self.pub.send(bin_obj)

    def _format_time(self, time, date, timezone, prefix_id):
        # TODO can we work out the time format from the regex? Probably but this is a task for another day
        time_format = self._config['prefixes'][prefix_id].get('time_format', '')
        if not time or not date or not time_format:
            return int(datetime.now().strftime('%s'))
        # Most syslog do not include the year, so we will add the current year if we are not supplied with one
        if '%y' in time_format or '%Y' in time_format:
            parsed_time = datetime.strptime('{} {}'.format(date, time), time_format)
        else:
            year = datetime.now().year
            parsed_time = datetime.strptime('{} {} {}'.format(year, date, time), '%Y {}'.format(time_format))
        return int((parsed_time - datetime(1970, 1, 1)).total_seconds())

    def start(self):
        '''
        Start the worker process.
        '''
        self._setup_ipc()
        # Start suicide polling thread
        thread = threading.Thread(target=self._suicide_when_without_parent, args=(os.getppid(),))
        thread.start()
        signal.signal(signal.SIGTERM, self._exit_gracefully)
        self.__up = True
        while self.__up:
            # bin_obj = self.sub.recv()
            # msg_dict, address = umsgpack.unpackb(bin_obj, use_list=False)
            try:
                bin_obj = self.sub.recv()
                msg_dict, address = umsgpack.unpackb(bin_obj, use_list=False)
            except IOError as error:
                if self.__up is False:
                    return
                else:
                    msg = 'Received IOError on {} device pipe: {}'.format(self._name, error)
                    log.error(msg, exc_info=True)
                    raise NapalmLogsExit(msg)
            log.debug('%s: dequeued %s, received from %s', self._name, msg_dict, address)
            if self._name == UNKNOWN_DEVICE_NAME:
                # If running in the sub-process publishing messages for unknown OSs.
                # This will always send what receives, as-is.
                log.debug('Publishing the message as-is, no further processing possible.')
                to_publish = {
                    'ip': address,
                    'host': 'unknown',
                    'timestamp': int(time.time()),
                    'message_details': msg_dict,
                    'os': UNKNOWN_DEVICE_NAME,
                    'error': 'UNKNOWN',
                    'model_name': 'unknown'
                }
                log.debug('Queueing to be published:')
                log.debug(to_publish)
                # self.pub_pipe.send(to_publish)
                self.pub.send(umsgpack.packb(to_publish))
                continue
            # From here on, we're running in a regular OS sub-process.
            host = msg_dict.get('host')
            prefix_id = msg_dict.pop('__prefix_id__')
            if 'timestamp' in msg_dict:
                timestamp = msg_dict.pop('timestamp')
            else:
                timestamp = self._format_time(msg_dict.get('time', ''),
                                              msg_dict.get('date', ''),
                                              msg_dict.get('timeZone', 'UTC'),
                                              prefix_id)
            facility = msg_dict.get('facility')
            severity = msg_dict.get('severity')

            kwargs = self._parse(msg_dict)
            if not kwargs:
                # Unable to identify what model to generate for the message in cause.
                if self.publisher_opts.get('send_raw'):
                    # But publish the message when the user requested to push raw messages.
                    to_publish = {
                        'ip': address,
                        'host': host,
                        'timestamp': timestamp,
                        'message_details': msg_dict,
                        'os': self._name,
                        'error': 'RAW',
                        'model_name': 'raw',
                        'facility': facility,
                        'severity': severity
                    }
                    log.debug('Queueing to be published:')
                    log.debug(to_publish)
                    # self.pub_pipe.send(to_publish)
                    self.pub.send(umsgpack.packb(to_publish))
                continue
            try:
                if '__python_fun__' in kwargs:
                    log.debug('Using the Python parser to determine the YANG-equivalent object')
                    yang_obj = kwargs['__python_fun__'](msg_dict)
                else:
                    yang_obj = self._emit(**kwargs)
            except Exception:
                log.exception('Unexpected error when generating the OC object.', exc_info=True)
                continue
            log.debug('Generated OC object:')
            log.debug(yang_obj)
            error = kwargs.get('error')
            model_name = kwargs.get('model')
            to_publish = {
                'error': error,
                'host': host,
                'ip': address,
                'timestamp': timestamp,
                'yang_message': yang_obj,
                'message_details': msg_dict,
                'yang_model': model_name,
                'os': self._name,
                'facility': facility,
                'severity': severity
            }
            log.debug('Queueing to be published:')
            log.debug(to_publish)
            # self.pub_pipe.send(to_publish)
            self.pub.send(umsgpack.packb(to_publish))
            # self._publish(to_publish)

    def stop(self):
        '''
        Stop the worker process.
        '''
        log.info('Stopping %s device process', self._name)
        self.__up = False
        self.sub.close()
        self.pub.close()
        self.ctx.term()
        self.pipe.close()
        self.pub_pipe.close()
<commit_msg>Dealer identity as bytes<commit_after># -*- coding: utf-8 -*-
'''
Device worker process
'''
from __future__ import absolute_import

# Import python stdlib
import os
import re
import time
import signal
import logging
import threading
from datetime import datetime

# Import thrid party libs
import zmq
import umsgpack

# Import napalm-logs pkgs
import napalm_logs.utils
from napalm_logs.proc import NapalmLogsProc
from napalm_logs.config import PUB_IPC_URL
from napalm_logs.config import DEV_IPC_URL
from napalm_logs.config import UNKNOWN_DEVICE_NAME
# exceptions
from napalm_logs.exceptions import NapalmLogsExit

log = logging.getLogger(__name__)

class NapalmLogsDeviceProc(NapalmLogsProc):
    '''
    Device sub-process class.
    '''
    def __init__(self, name, config, pipe, pub_pipe, publisher_opts):
        self._name = name
        self.pipe = pipe
        self._config = config
        self.pub_pipe = pub_pipe
        self.publisher_opts = publisher_opts
        self.__up = False
        self.compiled_messages = None
        self._compile_messages()

    def _exit_gracefully(self, signum, _):
        log.debug('Caught signal in %s device process', self._name)
        self.stop()

    def _setup_ipc(self):
        '''
        Subscribe to the right topic
        in the device IPC and publish to the
        publisher proxy.
        '''
        self.ctx = zmq.Context()
        # subscribe to device IPC
        log.debug('Creating the dealer IPC for %s', self._name)
        self.sub = self.ctx.socket(zmq.DEALER)
        self.sub.setsockopt(zmq.IDENTITY, bytes(self._name).encode('utf-8'))
        # subscribe to the corresponding IPC pipe
        self.sub.connect(DEV_IPC_URL)
        # self.sub.setsockopt(zmq.SUBSCRIBE, '')
        # publish to the publisher IPC
        self.pub = self.ctx.socket(zmq.PUSH)
        self.pub.connect(PUB_IPC_URL)

    def _compile_messages(self):
        '''
        Create a list of all OS messages and their compiled regexs
        '''
        self.compiled_messages = []
        if not self._config:
            return
        for message_dict in self._config.get('messages', {}):
            error = message_dict['error']
            tag = message_dict['tag']
            model = message_dict['model']
            match_on = message_dict.get('match_on', 'tag')
            if '__python_fun__' in message_dict:
                self.compiled_messages.append({
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'model': model,
                    '__python_fun__': message_dict['__python_fun__']
                })
                continue
            values = message_dict['values']
            line = message_dict['line']
            mapping = message_dict['mapping']
            # We will now figure out which position each value is in so we can use it with the match statement
            position = {}
            replace = {}
            for key in values.keys():
                if '|' in key:
                    new_key, replace[new_key] = key.replace(' ', '').split('|')
                    values[new_key] = values.pop(key)
                    key = new_key
                position[line.find('{' + key + '}')] = key
            sorted_position = {}
            for i, elem in enumerate(sorted(position.items())):
                sorted_position[elem[1]] = i + 1
            # Escape the line, then remove the escape for the curly bracets so they can be used when formatting
            escaped = re.escape(line).replace('\{', '{').replace('\}', '}')
            # Replace a whitespace with \s+
            escaped = escaped.replace('\ ', '\s+')
            self.compiled_messages.append(
                {
                    'error': error,
                    'tag': tag,
                    'match_on': match_on,
                    'line': re.compile(escaped.format(**values)),
                    'positions': sorted_position,
                    'values': values,
                    'replace': replace,
                    'model': model,
                    'mapping': mapping
                }
            )
        log.debug('Compiled messages:')
        log.debug(self.compiled_messages)

    def _parse(self, msg_dict):
        '''
        Parse a syslog message and check what OpenConfig object should
        be generated.
        '''
        error_present = False
        # log.debug('Matching the message:')
        # log.debug(msg_dict)
        for message in self.compiled_messages:
            # log.debug('Matching using:')
            # log.debug(message)
            match_on = message['match_on']
            if match_on not in msg_dict:
                # log.debug('%s is not a valid key in the partially parsed dict', match_on)
                continue
            if message['tag'] != msg_dict[match_on]:
                continue
            if '__python_fun__' in message:
                return {
                    'model': message['model'],
                    'error': message['error'],
                    '__python_fun__': message['__python_fun__']
                }
            error_present = True
            match = message['line'].search(msg_dict['message'])
            if not match:
                continue
            positions = message.get('positions', {})
            values = message.get('values')
            ret = {
                'model': message['model'],
                'mapping': message['mapping'],
                'replace': message['replace'],
                'error': message['error']
            }
            for key in values.keys():
                # Check if the value needs to be replaced
                if key in message['replace']:
                    result = napalm_logs.utils.cast(match.group(positions.get(key)), message['replace'][key])
                else:
                    result = match.group(positions.get(key))
                ret[key] = result
            return ret
        if error_present is True:
            log.info('Configured regex did not match for os: %s tag %s', self._name, msg_dict.get('tag', ''))
        else:
            log.info('Syslog message not configured for os: %s tag %s', self._name, msg_dict.get('tag', ''))

    def _emit(self, **kwargs):
        '''
        Emit an OpenConfig object given a certain combination of
        fields mappeed in the config to the corresponding hierarchy.
        '''
        oc_dict = {}
        for mapping, result_key in kwargs['mapping']['variables'].items():
            result = kwargs[result_key]
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)
        for mapping, result in kwargs['mapping']['static'].items():
            oc_dict = napalm_logs.utils.setval(mapping.format(**kwargs), result, oc_dict)

        return oc_dict

    def _publish(self, obj):
        '''
        Publish the OC object.
        '''
        bin_obj = umsgpack.packb(obj)
        self.pub.send(bin_obj)

    def _format_time(self, time, date, timezone, prefix_id):
        # TODO can we work out the time format from the regex? Probably but this is a task for another day
        time_format = self._config['prefixes'][prefix_id].get('time_format', '')
        if not time or not date or not time_format:
            return int(datetime.now().strftime('%s'))
        # Most syslog do not include the year, so we will add the current year if we are not supplied with one
        if '%y' in time_format or '%Y' in time_format:
            parsed_time = datetime.strptime('{} {}'.format(date, time), time_format)
        else:
            year = datetime.now().year
            parsed_time = datetime.strptime('{} {} {}'.format(year, date, time), '%Y {}'.format(time_format))
        return int((parsed_time - datetime(1970, 1, 1)).total_seconds())

    def start(self):
        '''
        Start the worker process.
        '''
        self._setup_ipc()
        # Start suicide polling thread
        thread = threading.Thread(target=self._suicide_when_without_parent, args=(os.getppid(),))
        thread.start()
        signal.signal(signal.SIGTERM, self._exit_gracefully)
        self.__up = True
        while self.__up:
            # bin_obj = self.sub.recv()
            # msg_dict, address = umsgpack.unpackb(bin_obj, use_list=False)
            try:
                bin_obj = self.sub.recv()
                msg_dict, address = umsgpack.unpackb(bin_obj, use_list=False)
            except IOError as error:
                if self.__up is False:
                    return
                else:
                    msg = 'Received IOError on {} device pipe: {}'.format(self._name, error)
                    log.error(msg, exc_info=True)
                    raise NapalmLogsExit(msg)
            log.debug('%s: dequeued %s, received from %s', self._name, msg_dict, address)
            if self._name == UNKNOWN_DEVICE_NAME:
                # If running in the sub-process publishing messages for unknown OSs.
                # This will always send what receives, as-is.
                log.debug('Publishing the message as-is, no further processing possible.')
                to_publish = {
                    'ip': address,
                    'host': 'unknown',
                    'timestamp': int(time.time()),
                    'message_details': msg_dict,
                    'os': UNKNOWN_DEVICE_NAME,
                    'error': 'UNKNOWN',
                    'model_name': 'unknown'
                }
                log.debug('Queueing to be published:')
                log.debug(to_publish)
                # self.pub_pipe.send(to_publish)
                self.pub.send(umsgpack.packb(to_publish))
                continue
            # From here on, we're running in a regular OS sub-process.
            host = msg_dict.get('host')
            prefix_id = msg_dict.pop('__prefix_id__')
            if 'timestamp' in msg_dict:
                timestamp = msg_dict.pop('timestamp')
            else:
                timestamp = self._format_time(msg_dict.get('time', ''),
                                              msg_dict.get('date', ''),
                                              msg_dict.get('timeZone', 'UTC'),
                                              prefix_id)
            facility = msg_dict.get('facility')
            severity = msg_dict.get('severity')

            kwargs = self._parse(msg_dict)
            if not kwargs:
                # Unable to identify what model to generate for the message in cause.
                if self.publisher_opts.get('send_raw'):
                    # But publish the message when the user requested to push raw messages.
                    to_publish = {
                        'ip': address,
                        'host': host,
                        'timestamp': timestamp,
                        'message_details': msg_dict,
                        'os': self._name,
                        'error': 'RAW',
                        'model_name': 'raw',
                        'facility': facility,
                        'severity': severity
                    }
                    log.debug('Queueing to be published:')
                    log.debug(to_publish)
                    # self.pub_pipe.send(to_publish)
                    self.pub.send(umsgpack.packb(to_publish))
                continue
            try:
                if '__python_fun__' in kwargs:
                    log.debug('Using the Python parser to determine the YANG-equivalent object')
                    yang_obj = kwargs['__python_fun__'](msg_dict)
                else:
                    yang_obj = self._emit(**kwargs)
            except Exception:
                log.exception('Unexpected error when generating the OC object.', exc_info=True)
                continue
            log.debug('Generated OC object:')
            log.debug(yang_obj)
            error = kwargs.get('error')
            model_name = kwargs.get('model')
            to_publish = {
                'error': error,
                'host': host,
                'ip': address,
                'timestamp': timestamp,
                'yang_message': yang_obj,
                'message_details': msg_dict,
                'yang_model': model_name,
                'os': self._name,
                'facility': facility,
                'severity': severity
            }
            log.debug('Queueing to be published:')
            log.debug(to_publish)
            # self.pub_pipe.send(to_publish)
            self.pub.send(umsgpack.packb(to_publish))
            # self._publish(to_publish)

    def stop(self):
        '''
        Stop the worker process.
        '''
        log.info('Stopping %s device process', self._name)
        self.__up = False
        self.sub.close()
        self.pub.close()
        self.ctx.term()
        self.pipe.close()
        self.pub_pipe.close()
<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Module which provides classes to perform source grouping."""

from __future__ import division
import abc
import numpy as np
from astropy.table import Column, Table, vstack

__all__ = ['DAOGroup', 'GroupStarsBase']

class GroupStarsBase(object):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def group_stars(self, starlist):
        pass

class DAOGroup(GroupStarsBase):
    """
    This is class implements the DAOGROUP algorithm presented by
    Stetson (1987).

    The method ``group_stars`` divides an entire starlist into sets of
    distinct, self-contained groups of mutually overlapping stars.
    It accepts as input a list of stars and determines which stars are close
    enough to be capable of adversely influencing each others' profile fits.

    Attributes
    ----------
    crit_separation : float or int
        Distance, in units of pixels, such that any two stars separated by
        less than this distance will be placed in the same group.

    Notes
    -----
    Assuming the psf fwhm to be known, 'crit_separation' may be set to
    k*fwhm, for some positive real k.

    See
    ---
    `~daofind`

    References
    ----------
    [1] Stetson, Astronomical Society of the Pacific, Publications,
        (ISSN 0004-6280), vol. 99, March 1987, p. 191-222.
        Available at: http://adsabs.harvard.edu/abs/1987PASP...99..191S
    """

    def __init__(self, crit_separation):
        self.crit_separation = crit_separation

    def __call__(self, starlist):
        """
        Parameters
        ----------
        starlist : `~astropy.table.Table`
            List of stars positions. Columns named as 'x_0' and 'y_0', which
            corresponds to the centroid coordinates of the sources, must be
            provided.

        Returns
        -------
        group_starlist : `~astropy.table.Table`
            ``starlist`` with an additional column named ``group_id`` whose
            unique values represent groups of mutually overlapping stars.
        """

        return self.group_stars(starlist)

    @property
    def crit_separation(self):
        return self._crit_separation

    @crit_separation.setter
    def crit_separation(self, crit_separation):
        if not isinstance(crit_separation, (float, int)):
            raise ValueError('crit_separation is expected to be either '+
                             'float or int. Received {}.'\
                             .format(type(crit_separation)))
        elif crit_separation < 0.0:
            raise ValueError('crit_separation is expected to be a positive '+
                             'real number. Got {}'.format(crit_separation))
        else:
            self._crit_separation = crit_separation

    def group_stars(self, starlist):
        cstarlist = starlist.copy()

        if 'id' not in cstarlist.colnames:
            cstarlist.add_column(Column(name='id',
                                        data=np.arange(len(cstarlist)) + 1))
        cstarlist.add_column(Column(name='group_id',\
                                data=np.zeros(len(cstarlist), dtype=np.int)))

        if not np.array_equal(cstarlist['id'], np.arange(len(cstarlist)) +1):
            raise ValueError('id colum must be an integer-valued ' +
                             'sequence starting from 1. ' +
                             'Got {}'.format(cstarlist['id']))

        n = 1
        while (cstarlist['group_id'] == 0).sum() > 0:
            init_star = cstarlist[np.where(cstarlist['group_id'] == 0)[0][0]]
            index = self.find_group(init_star,
                                    cstarlist[cstarlist['group_id'] == 0])
            cstarlist['group_id'][index-1] = n
            k = 1
            K = len(index)
            while k < K:
                init_star = cstarlist[cstarlist['id'] == index[k]]
                tmp_index = self.find_group(init_star,\
                                        cstarlist[cstarlist['group_id'] == 0])
                if len(tmp_index) > 0:
                    cstarlist['group_id'][tmp_index-1] = n
                    index = np.append(index, tmp_index)
                    K = len(index)
                k += 1
            n += 1

        return cstarlist

    def find_group(self, star, starlist):
        """
        Find the ids of those stars in ``starlist`` which are at a distance less
        than ``crit_separation`` from ``star``.

        Parameters
        ----------
        star : `~astropy.table.Row`
            Star which will be either the head of a cluster or an isolated one.
        starlist : `~astropy.table.Table`
            List of stars positions. Columns named as 'x_0' and 'y_0', which
            corresponds to the centroid coordinates of the sources, must be
            provided.

        Returns
        -------
        Array containing the ids of those stars which are at a
        distance less than `crit_separation` from `star`.
        """

        star_distance = np.hypot(star['x_0'] - starlist['x_0'],
                                 star['y_0'] - starlist['y_0'])
        distance_criteria = star_distance < self.crit_separation
        return np.asarray(starlist[distance_criteria]['id'])
<commit_msg>make sphinx happy (?)<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Module which provides classes to perform source grouping."""

from __future__ import division
import abc
import numpy as np
from astropy.table import Column, Table, vstack

__all__ = ['DAOGroup', 'GroupStarsBase']

class GroupStarsBase(object):
    __metaclass__ = abc.ABCMeta

    @abc.abstractmethod
    def group_stars(self, starlist):
        pass

class DAOGroup(GroupStarsBase):
    """
    This is class implements the DAOGROUP algorithm presented by
    Stetson (1987).

    The method ``group_stars`` divides an entire starlist into sets of
    distinct, self-contained groups of mutually overlapping stars.
    It accepts as input a list of stars and determines which stars are close
    enough to be capable of adversely influencing each others' profile fits.

    Attributes
    ----------
    crit_separation : float or int
        Distance, in units of pixels, such that any two stars separated by
        less than this distance will be placed in the same group.

    Notes
    -----
    Assuming the psf fwhm to be known, ``crit_separation`` may be set to
    k*fwhm, for some positive real k.

    See
    ---
    `~daofind`

    References
    ----------
    [1] Stetson, Astronomical Society of the Pacific, Publications,
        (ISSN 0004-6280), vol. 99, March 1987, p. 191-222.
        Available at: http://adsabs.harvard.edu/abs/1987PASP...99..191S
    """

    def __init__(self, crit_separation):
        self.crit_separation = crit_separation

    def __call__(self, starlist):
        """
        Parameters
        ----------
        starlist : `~astropy.table.Table`
            List of stars positions. Columns named as ``x_0`` and ``y_0``,
            which corresponds to the centroid coordinates of the sources,
            must be provided.

        Returns
        -------
        group_starlist : `~astropy.table.Table`
            ``starlist`` with an additional column named ``group_id`` whose
            unique values represent groups of mutually overlapping stars.
        """

        return self.group_stars(starlist)

    @property
    def crit_separation(self):
        return self._crit_separation

    @crit_separation.setter
    def crit_separation(self, crit_separation):
        if not isinstance(crit_separation, (float, int)):
            raise ValueError('crit_separation is expected to be either '+
                             'float or int. Received {}.'\
                             .format(type(crit_separation)))
        elif crit_separation < 0.0:
            raise ValueError('crit_separation is expected to be a positive '+
                             'real number. Got {}'.format(crit_separation))
        else:
            self._crit_separation = crit_separation

    def group_stars(self, starlist):
        cstarlist = starlist.copy()

        if 'id' not in cstarlist.colnames:
            cstarlist.add_column(Column(name='id',
                                        data=np.arange(len(cstarlist)) + 1))
        cstarlist.add_column(Column(name='group_id',\
                                data=np.zeros(len(cstarlist), dtype=np.int)))

        if not np.array_equal(cstarlist['id'], np.arange(len(cstarlist)) +1):
            raise ValueError('id colum must be an integer-valued ' +
                             'sequence starting from 1. ' +
                             'Got {}'.format(cstarlist['id']))

        n = 1
        while (cstarlist['group_id'] == 0).sum() > 0:
            init_star = cstarlist[np.where(cstarlist['group_id'] == 0)[0][0]]
            index = self.find_group(init_star,
                                    cstarlist[cstarlist['group_id'] == 0])
            cstarlist['group_id'][index-1] = n
            k = 1
            K = len(index)
            while k < K:
                init_star = cstarlist[cstarlist['id'] == index[k]]
                tmp_index = self.find_group(init_star,\
                                        cstarlist[cstarlist['group_id'] == 0])
                if len(tmp_index) > 0:
                    cstarlist['group_id'][tmp_index-1] = n
                    index = np.append(index, tmp_index)
                    K = len(index)
                k += 1
            n += 1

        return cstarlist

    def find_group(self, star, starlist):
        """
        Find the ids of those stars in ``starlist`` which are at a distance less
        than ``crit_separation`` from ``star``.

        Parameters
        ----------
        star : `~astropy.table.Row`
            Star which will be either the head of a cluster or an isolated one.
        starlist : `~astropy.table.Table`
            List of stars positions. Columns named as ``x_0`` and ``y_0``,
            which corresponds to the centroid coordinates of the sources,
            must be provided.

        Returns
        -------
        Array containing the ids of those stars which are at a distance less
        than ``crit_separation`` from ``star``.
        """

        star_distance = np.hypot(star['x_0'] - starlist['x_0'],
                                 star['y_0'] - starlist['y_0'])
        distance_criteria = star_distance < self.crit_separation
        return np.asarray(starlist[distance_criteria]['id'])
<|endoftext|>


<commit_before>import numpy
import matplotlib.pyplot as plt
import sys
import json
import argparse

# input:
#  alignment matrix - numpy array
#  shape (target tokens + eos, number of hidden source states = source tokens +eos)
# one line correpsonds to one decoding step producing one target token
# each line has the attention model weights corresponding to that decoding step
# each float on a line is the attention model weight for a corresponding source state.
# plot: a heat map of the alignment matrix
# x axis are the source tokens (alignment is to source hidden state that roughly corresponds to a source token)
# y axis are the target tokens

# http://stackoverflow.com/questions/14391959/heatmap-in-matplotlib-with-pcolor
def plot_head_map(mma, target_labels, source_labels):
  fig, ax = plt.subplots()
  heatmap = ax.pcolor(mma, cmap=plt.cm.Blues)

  # put the major ticks at the middle of each cell
  ax.set_xticks(numpy.arange(mma.shape[1])+0.5, minor=False)
  ax.set_yticks(numpy.arange(mma.shape[0])+0.5, minor=False)

  # without this I get some extra columns rows
  # http://stackoverflow.com/questions/31601351/why-does-this-matplotlib-heatmap-have-an-extra-blank-column
  ax.set_xlim(0, int(mma.shape[1]))
  ax.set_ylim(0, int(mma.shape[0]))

  # want a more natural, table-like display
  ax.invert_yaxis()
  ax.xaxis.tick_top()

  # source words -> column labels
  ax.set_xticklabels(source_labels, minor=False)
  # target words -> row labels
  ax.set_yticklabels(target_labels, minor=False)

  plt.xticks(rotation=45)

  #plt.tight_layout()
  plt.show()

# column labels -> target words
# row labels -> source words

def read_alignment_matrix(f):
  header = f.readline().strip().split('|||')
  if header[0] == '':
    return None, None, None, None
  sid = int(header[0].strip())
  # number of tokens in source and translation +1 for eos
  src_count, trg_count = map(int,header[-1].split())
  # source words
  source_labels = header[3].decode('UTF-8').split()
  source_labels.append('</s>')
  # target words
  target_labels = header[1].decode('UTF-8').split()
  target_labels.append('</s>')

  mm = []
  for r in range(trg_count):
    alignment = map(float,f.readline().strip().split())
    mm.append(alignment)
  mma = numpy.array(mm)
  return sid,mma, target_labels, source_labels

def read_plot_alignment_matrices(f, n):
  while(f):
    sid, mma, target_labels, source_labels = read_alignment_matrix(f)
    if mma is None:
      return
    if sid >n:
      return
    plot_head_map(mma, target_labels, source_labels)
    # empty line separating the matrices
    f.readline()

"""
Adding functions to read the json format.
"""

def read_plot_alignment_json(file, n):
    while (file):
        sid, mma, target_labels, source_labels = read_alignment_json(file)
        if mma is None:
            return
        if sid > n:
            return
        plot_head_map(mma, target_labels, source_labels)

def read_alignment_json(file):
    data = file.readline() ##one line containing the json object.
    if len(data.strip()) == 0:
        return None, None, None, None
    jdata = json.loads(data)
    ## messy json encodings... TODO: make this better
    jdata = json.loads(json.dumps(jdata).decode('unicode-escape').encode('utf8'))
    #print jdata
    sid = int(jdata["id"])
    mma = numpy.array(jdata["matrix"])
    ##target words
    target_labels = jdata["target_sent"].split()
    target_labels.append('</s>')
    ##source words
    source_labels = jdata["source_sent"].split()
    source_labels.append('</s>')
    return sid,mma, target_labels, source_labels

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    # '/Users/mnadejde/Documents/workspace/MTMA2016/models/wmt16_systems/en-de/test.alignment'
    parser.add_argument('--input', '-i', type=argparse.FileType('r'),
                            default='/Users/mnadejde/Documents/workspace/MTMA2016/models/wmt16_systems/ro-en/newstest2016-roen-src.ro.alignment', metavar='PATH',
                            help="Input file (default: standard input)")

    parser.add_argument('--json', '-j', required = False,action="store_true",
                            help="If this option is used, then read alignment matrix from a Json formatted file.")
    args = parser.parse_args()

    if args.json:
        read_plot_alignment_json(args.input, 10)   ##n is the maximum number of sentences to process.
    else:
        read_plot_alignment_matrices(args.input,10)
>>>>>>> origin/nematus-liucan
<commit_msg>fix plot_heatmap.py<commit_after>import numpy
import matplotlib.pyplot as plt
import sys
import json
import argparse

# input:
#  alignment matrix - numpy array
#  shape (target tokens + eos, number of hidden source states = source tokens +eos)
# one line correpsonds to one decoding step producing one target token
# each line has the attention model weights corresponding to that decoding step
# each float on a line is the attention model weight for a corresponding source state.
# plot: a heat map of the alignment matrix
# x axis are the source tokens (alignment is to source hidden state that roughly corresponds to a source token)
# y axis are the target tokens

# http://stackoverflow.com/questions/14391959/heatmap-in-matplotlib-with-pcolor
def plot_head_map(mma, target_labels, source_labels):
  fig, ax = plt.subplots()
  heatmap = ax.pcolor(mma, cmap=plt.cm.Blues)

  # put the major ticks at the middle of each cell
  ax.set_xticks(numpy.arange(mma.shape[1])+0.5, minor=False)
  ax.set_yticks(numpy.arange(mma.shape[0])+0.5, minor=False)

  # without this I get some extra columns rows
  # http://stackoverflow.com/questions/31601351/why-does-this-matplotlib-heatmap-have-an-extra-blank-column
  ax.set_xlim(0, int(mma.shape[1]))
  ax.set_ylim(0, int(mma.shape[0]))

  # want a more natural, table-like display
  ax.invert_yaxis()
  ax.xaxis.tick_top()

  # source words -> column labels
  ax.set_xticklabels(source_labels, minor=False)
  # target words -> row labels
  ax.set_yticklabels(target_labels, minor=False)

  plt.xticks(rotation=45)

  #plt.tight_layout()
  plt.show()

# column labels -> target words
# row labels -> source words

def read_alignment_matrix(f):
  header = f.readline().strip().split('|||')
  if header[0] == '':
    return None, None, None, None
  sid = int(header[0].strip())
  # number of tokens in source and translation +1 for eos
  src_count, trg_count = map(int,header[-1].split())
  # source words
  source_labels = header[3].decode('UTF-8').split()
  source_labels.append('</s>')
  # target words
  target_labels = header[1].decode('UTF-8').split()
  target_labels.append('</s>')

  mm = []
  for r in range(trg_count):
    alignment = map(float,f.readline().strip().split())
    mm.append(alignment)
  mma = numpy.array(mm)
  return sid,mma, target_labels, source_labels

def read_plot_alignment_matrices(f, n):
  while(f):
    sid, mma, target_labels, source_labels = read_alignment_matrix(f)
    if mma is None:
      return
    if sid >n:
      return
    plot_head_map(mma, target_labels, source_labels)
    # empty line separating the matrices
    f.readline()

"""
Adding functions to read the json format.
"""

def read_plot_alignment_json(file, n):
    while (file):
        sid, mma, target_labels, source_labels = read_alignment_json(file)
        if mma is None:
            return
        if sid > n:
            return
        plot_head_map(mma, target_labels, source_labels)

def read_alignment_json(file):
    data = file.readline() ##one line containing the json object.
    if len(data.strip()) == 0:
        return None, None, None, None
    jdata = json.loads(data)
    ## messy json encodings... TODO: make this better
    jdata = json.loads(json.dumps(jdata).decode('unicode-escape').encode('utf8'))
    #print jdata
    sid = int(jdata["id"])
    mma = numpy.array(jdata["matrix"])
    ##target words
    target_labels = jdata["target_sent"].split()
    target_labels.append('</s>')
    ##source words
    source_labels = jdata["source_sent"].split()
    source_labels.append('</s>')
    return sid,mma, target_labels, source_labels

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    # '/Users/mnadejde/Documents/workspace/MTMA2016/models/wmt16_systems/en-de/test.alignment'
    parser.add_argument('--input', '-i', type=argparse.FileType('r'),
                            default='/Users/mnadejde/Documents/workspace/MTMA2016/models/wmt16_systems/ro-en/newstest2016-roen-src.ro.alignment', metavar='PATH',
                            help="Input file (default: standard input)")

    parser.add_argument('--json', '-j', required = False,action="store_true",
                            help="If this option is used, then read alignment matrix from a Json formatted file.")
    args = parser.parse_args()

    if args.json:
        read_plot_alignment_json(args.input, 10)   ##n is the maximum number of sentences to process.
    else:
        read_plot_alignment_matrices(args.input,10)
<|endoftext|>


<commit_before>from phovea_server import ns
from phovea_server.config import view as config_view
import memcache
import logging
import gevent
import gevent.lock

__author__ = 'Samuel Gratzl'
app = ns.Namespace(__name__)
conf = config_view('phovea_clue')

mc = memcache.Client(conf.get('memcached'), debug=0)
mc_prefix = 'clue_'

_log = logging.getLogger(__name__)

class Screenshotter(object):
  def __init__(self):
    self._lock = gevent.lock.BoundedSemaphore(1)
    self._timeout = None
    self._driver = None
    pass

  def _timed_out(self):
    try:
      _log.info('quiting driver')
      self._driver.quit()
      _log.info('quitted driver')
    finally:
      self._driver = None

  def _get(self):
    from selenium import webdriver

    if self._timeout is not None:
      gevent.kill(self._timeout)
      self._timeout = None

    if self._driver is None:
      _log.info('create driver')
      options = webdriver.ChromeOptions()
      options.debugger_address = conf.chromeAddress
      self._driver = webdriver.Chrome(chrome_options=options)
      self._driver.implicitly_wait(30)  # wait at most 30 seconds
    return self._driver

  def _free(self):
    # schedule that the driver will be cleaned up if not used
    self._timeout = gevent.spawn_later(600, self._timed_out)

  def take(self, url, body=None):
    with self._lock:
      try:
        driver = self._get()
        _log.info('url %s', url)
        driver.get(url)

        if body is not None:
          try:
            body(driver)
          except Exception as e:
            _log.exception('cannot fullfil query %s', e)
        _log.info('take screenshot')
        obj = driver.get_screenshot_as_png()
        return obj
      finally:
        self._free()

screenshotter = Screenshotter()

def randomword(length):
  import random, string
  return ''.join(random.choice(string.lowercase) for i in range(length))

def generate_url(app, prov_id, state):
  # add a random parameter to force a proper reload
  base = '{s}/?clue_random={r}#clue_graph={g}&clue_state={n}&clue=P&clue_store=remote&clue_headless=Y'
  return base.format(s=conf.server, g=prov_id, n=state, r=randomword(5))

def generate_key(app, prov_id, state):
  return 'a={a},p={g},s={s}'.format(a=app, g=prov_id, s=state)

def generate_slide_url(app, prov_id, slide):
  base = '{s}/?clue_random={r}#clue_graph={g}&clue_slide={n}&clue=P&clue_store=remote&clue_headless=Y'
  return base.format(s=conf.server, g=prov_id, n=slide, r=randomword(5))

def generate_slide_key(app, prov_id, slide):
  return 'a={a},p={g},u={s}'.format(a=app, g=prov_id, s=slide)

@app.route('/dump/<path:page>')
def test(page):
  obj = screenshotter.take('http://' + page)
  return ns.Response(obj, mimetype='image/png')

def create_via_selenium(url, width, height):
  def eval_clue(driver):
    driver.find_element_by_css_selector('main')
    found = None
    tries = 0
    while not found and tries < 3:
      tries += 1
      for entry in driver.get_log('browser'):
        _log.info(entry)
      found = driver.find_element_by_css_selector('body.clue_jumped')

    if found:
      _log.info('found jumped flag: {}'.format(found.get_attribute('class')))
    else:
      _log.warn('cannnot find jumped flag after 3 x 30 seconds, give up and take a screenshot')

  return screenshotter.take(url, eval_clue)

def _create_screenshot_impl(app, prov_id, state, format, width=1920, height=1080, force=False):
  url = generate_url(app, prov_id, state)

  key = mc_prefix + url + 'w' + str(width) + 'h' + str(height)

  obj = mc.get(key) if not force else None
  if not obj:
    _log.info('requesting url %s', url)
    obj = create_via_selenium(url, width, height)
    mc.set(key, obj)
  return obj

def _create_preview_impl(app, prov_id, slide, format, width=1920, height=1080, force=False):
  url = generate_slide_url(app, prov_id, slide)

  key = mc_prefix + url + 'w' + str(width) + 'h' + str(height)

  obj = mc.get(key) if not force else None
  if not obj:
    _log.debug('requesting url %s', url)
    obj = create_via_selenium(url, width, height)
    mc.set(key, obj)
  return obj

def fix_format(format):
  return 'jpeg' if format == 'jpg' else format

@app.route('/screenshot/<app>/<prov_id>/<state>.<format>')
def create_screenshot(app, prov_id, state, format):
  width = ns.request.args.get('width', 1920)
  height = ns.request.args.get('height', 1080)
  force = ns.request.args.get('force', None) is not None

  s = _create_screenshot_impl(app, prov_id, state, format, width, height, force)
  return ns.Response(s, mimetype='image/' + fix_format(format))

def to_thumbnail(s, width, format):
  import PIL.Image
  import io

  b = io.BytesIO(s)
  img = PIL.Image.open(b)

  wpercent = (width / float(img.size[0]))
  height = int(float(img.size[1]) * float(wpercent))
  img.thumbnail((width, height), PIL.Image.ANTIALIAS)

  b = io.BytesIO()
  img.save(b, fix_format(format))
  b.seek(0)
  obj = b.read()
  return obj

@app.route('/thumbnail/<app>/<prov_id>/<state>.<format>')
def create_thumbnail(app, prov_id, state, format):
  format = fix_format(format)
  width = int(ns.request.args.get('width', 128))
  force = ns.request.args.get('force', None) is not None

  key = mc_prefix + generate_key(app, prov_id, state) + 't' + str(width)

  obj = mc.get(key)
  if not obj or force:
    s = _create_screenshot_impl(app, prov_id, state, format, force=force)
    obj = to_thumbnail(s, width, format)
    mc.set(key, obj)

  return ns.Response(obj, mimetype='image/' + format)

@app.route('/preview/<app>/<prov_id>/<slide>.<format>')
def create_preview(app, prov_id, slide, format):
  width = ns.request.args.get('width', 1920)
  height = ns.request.args.get('height', 1080)
  force = ns.request.args.get('force', None) is not None

  s = _create_preview_impl(app, prov_id, slide, format, width, height, force)
  return ns.Response(s, mimetype='image/' + fix_format(format))

@app.route('/preview_thumbnail/<app>/<prov_id>/<slide>.<format>')
def create_preview_thumbnail(app, prov_id, slide, format):
  format = fix_format(format)
  width = int(ns.request.args.get('width', 128))
  force = ns.request.args.get('force', None) is not None

  key = mc_prefix + generate_slide_key(app, prov_id, slide) + 't' + str(width)

  obj = mc.get(key)
  if not obj or force:
    s = _create_preview_impl(app, prov_id, slide, format, force=force)
    obj = to_thumbnail(s, width, format)
    mc.set(key, obj)

  return ns.Response(obj, mimetype='image/' + format)

def create():
  """
   entry point of this plugin
  """
  app.debug = True
  return app

if __name__ == '__main__':
  app.debug = True
  app.run(host='0.0.0.0')
<commit_msg>fix flake8<commit_after>from phovea_server import ns
from phovea_server.config import view as config_view
import memcache
import logging
import gevent
import gevent.lock

__author__ = 'Samuel Gratzl'
app = ns.Namespace(__name__)
conf = config_view('phovea_clue')

mc = memcache.Client(conf.get('memcached'), debug=0)
mc_prefix = 'clue_'

_log = logging.getLogger(__name__)

class Screenshotter(object):
  def __init__(self):
    self._lock = gevent.lock.BoundedSemaphore(1)
    self._timeout = None
    self._driver = None
    pass

  def _timed_out(self):
    try:
      _log.info('quiting driver')
      self._driver.quit()
      _log.info('quitted driver')
    finally:
      self._driver = None

  def _get(self):
    from selenium import webdriver

    if self._timeout is not None:
      gevent.kill(self._timeout)
      self._timeout = None

    if self._driver is None:
      _log.info('create driver')
      options = webdriver.ChromeOptions()
      options.debugger_address = conf.chromeAddress
      self._driver = webdriver.Chrome(chrome_options=options)
      self._driver.implicitly_wait(30)  # wait at most 30 seconds
    return self._driver

  def _free(self):
    # schedule that the driver will be cleaned up if not used
    self._timeout = gevent.spawn_later(600, self._timed_out)

  def take(self, url, body=None):
    with self._lock:
      try:
        driver = self._get()
        _log.info('url %s', url)
        driver.get(url)

        if body is not None:
          try:
            body(driver)
          except Exception as e:
            _log.exception('cannot fullfil query %s', e)
        _log.info('take screenshot')
        obj = driver.get_screenshot_as_png()
        return obj
      finally:
        self._free()

screenshotter = Screenshotter()

def randomword(length):
  import random
  import string
  return ''.join(random.choice(string.lowercase) for i in range(length))

def generate_url(app, prov_id, state):
  # add a random parameter to force a proper reload
  base = '{s}/?clue_random={r}#clue_graph={g}&clue_state={n}&clue=P&clue_store=remote&clue_headless=Y'
  return base.format(s=conf.server, g=prov_id, n=state, r=randomword(5))

def generate_key(app, prov_id, state):
  return 'a={a},p={g},s={s}'.format(a=app, g=prov_id, s=state)

def generate_slide_url(app, prov_id, slide):
  base = '{s}/?clue_random={r}#clue_graph={g}&clue_slide={n}&clue=P&clue_store=remote&clue_headless=Y'
  return base.format(s=conf.server, g=prov_id, n=slide, r=randomword(5))

def generate_slide_key(app, prov_id, slide):
  return 'a={a},p={g},u={s}'.format(a=app, g=prov_id, s=slide)

@app.route('/dump/<path:page>')
def test(page):
  obj = screenshotter.take('http://' + page)
  return ns.Response(obj, mimetype='image/png')

def create_via_selenium(url, width, height):
  def eval_clue(driver):
    driver.find_element_by_css_selector('main')
    found = None
    tries = 0
    while not found and tries < 3:
      tries += 1
      for entry in driver.get_log('browser'):
        _log.info(entry)
      found = driver.find_element_by_css_selector('body.clue_jumped')

    if found:
      _log.info('found jumped flag: {}'.format(found.get_attribute('class')))
    else:
      _log.warn('cannnot find jumped flag after 3 x 30 seconds, give up and take a screenshot')

  return screenshotter.take(url, eval_clue)

def _create_screenshot_impl(app, prov_id, state, format, width=1920, height=1080, force=False):
  url = generate_url(app, prov_id, state)

  key = mc_prefix + url + 'w' + str(width) + 'h' + str(height)

  obj = mc.get(key) if not force else None
  if not obj:
    _log.info('requesting url %s', url)
    obj = create_via_selenium(url, width, height)
    mc.set(key, obj)
  return obj

def _create_preview_impl(app, prov_id, slide, format, width=1920, height=1080, force=False):
  url = generate_slide_url(app, prov_id, slide)

  key = mc_prefix + url + 'w' + str(width) + 'h' + str(height)

  obj = mc.get(key) if not force else None
  if not obj:
    _log.debug('requesting url %s', url)
    obj = create_via_selenium(url, width, height)
    mc.set(key, obj)
  return obj

def fix_format(format):
  return 'jpeg' if format == 'jpg' else format

@app.route('/screenshot/<app>/<prov_id>/<state>.<format>')
def create_screenshot(app, prov_id, state, format):
  width = ns.request.args.get('width', 1920)
  height = ns.request.args.get('height', 1080)
  force = ns.request.args.get('force', None) is not None

  s = _create_screenshot_impl(app, prov_id, state, format, width, height, force)
  return ns.Response(s, mimetype='image/' + fix_format(format))

def to_thumbnail(s, width, format):
  import PIL.Image
  import io

  b = io.BytesIO(s)
  img = PIL.Image.open(b)

  wpercent = (width / float(img.size[0]))
  height = int(float(img.size[1]) * float(wpercent))
  img.thumbnail((width, height), PIL.Image.ANTIALIAS)

  b = io.BytesIO()
  img.save(b, fix_format(format))
  b.seek(0)
  obj = b.read()
  return obj

@app.route('/thumbnail/<app>/<prov_id>/<state>.<format>')
def create_thumbnail(app, prov_id, state, format):
  format = fix_format(format)
  width = int(ns.request.args.get('width', 128))
  force = ns.request.args.get('force', None) is not None

  key = mc_prefix + generate_key(app, prov_id, state) + 't' + str(width)

  obj = mc.get(key)
  if not obj or force:
    s = _create_screenshot_impl(app, prov_id, state, format, force=force)
    obj = to_thumbnail(s, width, format)
    mc.set(key, obj)

  return ns.Response(obj, mimetype='image/' + format)

@app.route('/preview/<app>/<prov_id>/<slide>.<format>')
def create_preview(app, prov_id, slide, format):
  width = ns.request.args.get('width', 1920)
  height = ns.request.args.get('height', 1080)
  force = ns.request.args.get('force', None) is not None

  s = _create_preview_impl(app, prov_id, slide, format, width, height, force)
  return ns.Response(s, mimetype='image/' + fix_format(format))

@app.route('/preview_thumbnail/<app>/<prov_id>/<slide>.<format>')
def create_preview_thumbnail(app, prov_id, slide, format):
  format = fix_format(format)
  width = int(ns.request.args.get('width', 128))
  force = ns.request.args.get('force', None) is not None

  key = mc_prefix + generate_slide_key(app, prov_id, slide) + 't' + str(width)

  obj = mc.get(key)
  if not obj or force:
    s = _create_preview_impl(app, prov_id, slide, format, force=force)
    obj = to_thumbnail(s, width, format)
    mc.set(key, obj)

  return ns.Response(obj, mimetype='image/' + format)

def create():
  """
   entry point of this plugin
  """
  app.debug = True
  return app

if __name__ == '__main__':
  app.debug = True
  app.run(host='0.0.0.0')
<|endoftext|>