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

"""Workers for driving screen captures, perceptual diffs, and related work."""

import Queue
import logging
import sys
import threading

# Local Libraries
import gflags
FLAGS = gflags.FLAGS

gflags.DEFINE_float(
    'polltime', 1.0,
    'How long to sleep between polling for work from an input queue, '
    'a subprocess, or a waiting timer.')

class WorkItem(object):
    """Base work item that can be handled by a worker thread."""

    # Instance variables. May be overridden by a sub-class @property.
    error = None
    done = False

    # Set this to True for WorkItems that should never wait for their
    # return values.
    fire_and_forget = False

    def __init__(self):
        pass

    @staticmethod
    def _print_tree(obj):
        if isinstance(obj, dict):
            result = []
            for key, value in obj.iteritems():
                result.append("%s: %s" % (key, WorkItem._print_tree(value)))
            return '{%s}' % ', '.join(result)
        else:
            value_str = repr(obj)
            if len(value_str) > 100:
                return '%s...%s' % (value_str[:100], value_str[-1])
            else:
                return value_str

    def _get_dict_for_repr(self):
        return self.__dict__

    def __repr__(self):
        return '%s.%s(%s)' % (
            self.__class__.__module__,
            self.__class__.__name__,
            self._print_tree(self._get_dict_for_repr()))

    def check_result(self):
        # TODO: For WorkflowItems, remove generator.throw(*item.error) from
        # the stack trace since it's noise. General approach outlined here:
        # https://github.com/mitsuhiko/jinja2/blob/master/jinja2/debug.py
        if self.error:
            raise self.error[0], self.error[1], self.error[2]

class WorkerThread(threading.Thread):
    """Base worker thread that handles items one at a time."""

    def __init__(self, input_queue, output_queue):
        """Initializer.

        Args:
            input_queue: Queue this worker consumes work from.
            output_queue: Queue where this worker puts new work items, if any.
        """
        threading.Thread.__init__(self)
        self.daemon = True
        self.input_queue = input_queue
        self.output_queue = output_queue
        self.interrupted = False
        self.polltime = FLAGS.polltime

    def stop(self):
        """Stops the thread but does not join it."""
        if self.interrupted:
            return
        self.interrupted = True

    def run(self):
        while not self.interrupted:
            try:
                item = self.input_queue.get(True, self.polltime)
            except Queue.Empty:
                self.handle_nothing()
                continue

            try:
                next_item = self.handle_item(item)
            except Exception, e:
                item.error = sys.exc_info()
                logging.exception('%s error item=%r', self.worker_name, item)
                self.output_queue.put(item)
            else:
                logging.debug('%s processed item=%r', self.worker_name, item)
                if not isinstance(item, WorkflowItem):
                    item.done = True
                if next_item:
                    self.output_queue.put(next_item)
            finally:
                self.input_queue.task_done()

    @property
    def worker_name(self):
        return '%s:%s' % (self.__class__.__name__, self.ident)

    def handle_nothing(self):
        """Runs whenever there are no items in the queue."""
        pass

    def handle_item(self, item):
        """Handles a single item.

        Args:
            item: WorkItem to process.

        Returns:
            A WorkItem that should go on the output queue. If None, then
            the provided work item is considered finished and no
            additional work is needed.
        """
        raise NotImplemented

class WorkflowItem(WorkItem):
    """Work item for coordinating other work items.

    To use: Sub-class and override run(). Yield WorkItems you want processed
    as part of this workflow. Exceptions in child workflows will be reinjected
    into the run() generator at the yield point. Results will be available on
    the WorkItems returned by yield statements. Yield a list of WorkItems
    to do them in parallel. The first error encountered for the whole list
    will be raised if there's an exception.
    """

    # Allow these value to be assigned or overridden by a sub-class @property.
    result = None
    root = False

    def __init__(self, *args, **kwargs):
        WorkItem.__init__(self)
        self.args = args
        self.kwargs = kwargs

    def run(self, *args, **kwargs):
        raise NotImplemented

class WaitAny(object):
    """Return control to a workflow after any one of these items finishes.

    As soon as a single work item completes, the combined barrier will be
    fulfilled and control will return to the WorkflowItem. The return values
    will be WorkItem instances, the same ones passed into WaitAny. For
    WorkflowItems the return values will be WorkflowItems if the work is not
    finished yet, and the final return value once work is finished.
    """

    def __init__(self, items):
        """Initializer.

        Args:
            items: List of WorkItems to wait for.
        """
        self.items = items

class Barrier(list):
    """Barrier for running multiple WorkItems in parallel."""

    def __init__(self, workflow, generator, work):
        """Initializer.

        Args:
            workflow: WorkflowItem instance this is for.
            generator: Current state of the WorkflowItem's generator.
            work: Next set of work to do. May be a single WorkItem object or
                a list or tuple that contains a set of WorkItems to run in
                parallel.
        """
        list.__init__(self)
        self.workflow = workflow
        self.generator = generator

        if isinstance(work, (list, tuple)):
            self[:] = list(work)
            self.was_list = True
            self.wait_any = False
        elif isinstance(work, WaitAny):
            self[:] = list(work.items)
            self.was_list = True
            self.wait_any = True
        else:
            self[:] = [work]
            self.was_list = False
            self.wait_any = False

        for item in self:
            assert isinstance(item, WorkItem)

    @property
    def outstanding(self):
        """Returns whether or not this barrier has pending work."""
        # Allow the same WorkItem to be yielded multiple times but not
        # count towards blocking the barrier.
        done_count = 0
        for item in self:
            if not self.wait_any and item.fire_and_forget:
                # Only count fire_and_forget items as done if this is
                # *not* a WaitAny barrier. We only want to return control
                # to the caller when at least one of the blocking items
                # has completed.
                done_count += 1
            elif item.done:
                done_count += 1

        if self.wait_any and done_count > 0:
            return False

        if done_count == len(self):
            return False

        return True

    @property
    def error(self):
        """Returns the error for this barrier and all work items, if any."""
        # Copy the error from any failed item to be the error for the whole
        # barrier. The first error seen "wins". Also handles the case where
        # the WorkItems passed into the barrier have already completed and
        # been marked with errors.
        for item in self:
            if isinstance(item, WorkItem) and item.error:
                return item.error
        return None

    def get_item(self):
        """Returns the item to send back into the workflow generator."""
        if self.was_list:
            blocking_items = self[:]
            self[:] = []
            for item in blocking_items:
                if (isinstance(item, WorkflowItem) and
                        item.done and
                        not item.error):
                    self.append(item.result)
                else:
                    self.append(item)
            return self
        else:
            return self[0]

class Return(Exception):
    """Raised in WorkflowItem.run to return a result to the caller."""

    def __init__(self, result=None):
        """Initializer.

        Args:
            result: Result of a WorkflowItem, if any.
        """
        self.result = result

class WorkflowThread(WorkerThread):
    """Worker thread for running workflows."""

    def __init__(self, input_queue, output_queue):
        """Initializer.

        Args:
            input_queue: Queue this worker consumes work from. These should be
                WorkflowItems to process, or any WorkItems registered with this
                class using the register() method.
            output_queue: Queue where this worker puts finished work items,
                if any.
        """
        WorkerThread.__init__(self, input_queue, output_queue)
        self.pending = {}
        self.work_map = {}
        self.worker_threads = []
        self.register(WorkflowItem, input_queue)

    # TODO: Implement drain, to let all existing work finish but no new work
    # allowed at the top of the funnel.

    def start(self):
        """Starts the coordinator thread and all related worker threads."""
        assert not self.interrupted
        for thread in self.worker_threads:
            thread.start()
        WorkerThread.start(self)

    def stop(self):
        """Stops the coordinator thread and all related threads."""
        if self.interrupted:
            return
        for thread in self.worker_threads:
            thread.interrupted = True
        self.interrupted = True

    def join(self):
        """Joins the coordinator thread and all worker threads."""
        for thread in self.worker_threads:
            thread.join()
        WorkerThread.join(self)

    def wait_one(self):
        """Waits until this worker has finished one work item or died."""
        while True:
            try:
                item = self.output_queue.get(True, 1)
            except Queue.Empty:
                continue
            except KeyboardInterrupt:
                logging.debug('Exiting')
                return
            else:
                item.check_result()
                return

    def register(self, work_type, queue):
        """Registers where work for a specific type can be executed.

        Args:
            work_type: Sub-class of WorkItem to register.
            queue: Queue instance where WorkItems of the work_type should be
                enqueued when they are yielded by WorkflowItems being run by
                this worker.
        """
        self.work_map[work_type] = queue

    def enqueue_barrier(self, barrier):
        for item in barrier:
            if item.done:
                # Don't reenqueue items that are already done.
                continue

            # Try to find a queue by the work item's class, then by its
            # super class, and so on.
            next_types = [type(item)]
            while next_types:
                try_types = next_types[:]
                next_types[:] = []
                for current_type in try_types:
                    target_queue = self.work_map.get(current_type)
                    if target_queue:
                        next_types = []
                        break
                    next_types.extend(current_type.__bases__)

            assert target_queue, 'Could not find queue to handle %r' % item

            if not item.fire_and_forget:
                self.pending[item] = barrier
            target_queue.put(item)

    def dequeue_barrier(self, item):
        # This is a WorkItem from a worker thread that has finished and
        # needs to be reinjected into a WorkflowItem generator.
        barrier = self.pending.pop(item, None)
        if not barrier:
            # Item was already finished in another barrier, or was
            # fire-and-forget and never part of a barrier; ignore it.
            return None

        if barrier.outstanding and not barrier.error:
            # More work to do and no error seen. Keep waiting.
            return None

        for work in barrier:
            # The barrier has been fulfilled one way or another. Clear out any
            # other pending parts of the barrier so they don't trigger again.
            self.pending.pop(work, None)

        return barrier

    def handle_item(self, item):
        if isinstance(item, WorkflowItem) and not item.done:
            workflow = item
            try:
                generator = item.run(*item.args, **item.kwargs)
            except TypeError, e:
                raise TypeError('%s: item=%r', e, item)
            item = None
        else:
            barrier = self.dequeue_barrier(item)
            if not barrier:
                return
            item = barrier.get_item()
            workflow = barrier.workflow
            generator = barrier.generator

        while True:
            logging.debug('Transitioning workflow=%r, generator=%r, item=%r',
                          workflow, generator, item)
            try:
                try:
                    error = item is not None and item.error
                    if error:
                        next_item = generator.throw(*error)
                    elif isinstance(item, WorkflowItem) and item.done:
                        next_item = generator.send(item.result)
                    else:
                        next_item = generator.send(item)
                except StopIteration:
                    workflow.done = True
                except Return, e:
                    workflow.done = True
                    workflow.result = e.result
                except Exception, e:
                    logging.exception('%s error item=%r',
                                      self.worker_name, item)
                    workflow.done = True
                    workflow.error = sys.exc_info()
            finally:
                if workflow.done:
                    if workflow.root:
                        # Root workflow finished. This goes to the output
                        # queue so it can be received by the main thread.
                        return workflow
                    else:
                        # Sub-workflow finished. Reinject it into the
                        # workflow so a pending parent can catch it.
                        self.input_queue.put(workflow)
                        return

            barrier = Barrier(workflow, generator, next_item)
            self.enqueue_barrier(barrier)
            if barrier.outstanding:
                break

            # If a returned barrier has no oustanding parts, immediately
            # progress the workflow.
            item = barrier.get_item()

class PrintWorkflow(WorkflowItem):
    """Prints a message to stdout."""

    def run(self, message):
        yield []  # Make this into a generator
        print message

def get_coordinator():
    """Creates a coordinator and returns it."""
    workflow_queue = Queue.Queue()
    complete_queue = Queue.Queue()
    coordinator = WorkflowThread(workflow_queue, complete_queue)
    coordinator.register(WorkflowItem, workflow_queue)
    return coordinator
<commit_msg>Clearing up exception logging to help track down an issue<commit_after>#!/usr/bin/env python
# Copyright 2013 Brett Slatkin
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Workers for driving screen captures, perceptual diffs, and related work."""

import Queue
import logging
import sys
import threading

# Local Libraries
import gflags
FLAGS = gflags.FLAGS

gflags.DEFINE_float(
    'polltime', 1.0,
    'How long to sleep between polling for work from an input queue, '
    'a subprocess, or a waiting timer.')

class WorkItem(object):
    """Base work item that can be handled by a worker thread."""

    # Instance variables. May be overridden by a sub-class @property.
    error = None
    done = False

    # Set this to True for WorkItems that should never wait for their
    # return values.
    fire_and_forget = False

    def __init__(self):
        pass

    @staticmethod
    def _print_tree(obj):
        if isinstance(obj, dict):
            result = []
            for key, value in obj.iteritems():
                result.append('%s: %s' % (key, WorkItem._print_tree(value)))
            return '{%s}' % ', '.join(result)
        else:
            value_str = repr(obj)
            if len(value_str) > 100:
                return '%s...%s' % (value_str[:100], value_str[-1])
            else:
                return value_str

    def _get_dict_for_repr(self):
        return self.__dict__

    def __repr__(self):
        return '%s.%s(%s)' % (
            self.__class__.__module__,
            self.__class__.__name__,
            self._print_tree(self._get_dict_for_repr()))

    def check_result(self):
        # TODO: For WorkflowItems, remove generator.throw(*item.error) from
        # the stack trace since it's noise. General approach outlined here:
        # https://github.com/mitsuhiko/jinja2/blob/master/jinja2/debug.py
        if self.error:
            raise self.error[0], self.error[1], self.error[2]

class WorkerThread(threading.Thread):
    """Base worker thread that handles items one at a time."""

    def __init__(self, input_queue, output_queue):
        """Initializer.

        Args:
            input_queue: Queue this worker consumes work from.
            output_queue: Queue where this worker puts new work items, if any.
        """
        threading.Thread.__init__(self)
        self.daemon = True
        self.input_queue = input_queue
        self.output_queue = output_queue
        self.interrupted = False
        self.polltime = FLAGS.polltime

    def stop(self):
        """Stops the thread but does not join it."""
        if self.interrupted:
            return
        self.interrupted = True

    def run(self):
        while not self.interrupted:
            try:
                item = self.input_queue.get(True, self.polltime)
            except Queue.Empty:
                self.handle_nothing()
                continue

            try:
                next_item = self.handle_item(item)
            except Exception, e:
                item.error = sys.exc_info()
                logging.exception('%s error item=%r', self.worker_name, item)
                self.output_queue.put(item)
            else:
                logging.debug('%s processed item=%r', self.worker_name, item)
                if not isinstance(item, WorkflowItem):
                    item.done = True
                if next_item:
                    self.output_queue.put(next_item)
            finally:
                self.input_queue.task_done()

    @property
    def worker_name(self):
        return '%s:%s' % (self.__class__.__name__, self.ident)

    def handle_nothing(self):
        """Runs whenever there are no items in the queue."""
        pass

    def handle_item(self, item):
        """Handles a single item.

        Args:
            item: WorkItem to process.

        Returns:
            A WorkItem that should go on the output queue. If None, then
            the provided work item is considered finished and no
            additional work is needed.
        """
        raise NotImplemented

class WorkflowItem(WorkItem):
    """Work item for coordinating other work items.

    To use: Sub-class and override run(). Yield WorkItems you want processed
    as part of this workflow. Exceptions in child workflows will be reinjected
    into the run() generator at the yield point. Results will be available on
    the WorkItems returned by yield statements. Yield a list of WorkItems
    to do them in parallel. The first error encountered for the whole list
    will be raised if there's an exception.
    """

    # Allow these value to be assigned or overridden by a sub-class @property.
    result = None
    root = False

    def __init__(self, *args, **kwargs):
        WorkItem.__init__(self)
        self.args = args
        self.kwargs = kwargs

    def run(self, *args, **kwargs):
        raise NotImplemented

class WaitAny(object):
    """Return control to a workflow after any one of these items finishes.

    As soon as a single work item completes, the combined barrier will be
    fulfilled and control will return to the WorkflowItem. The return values
    will be WorkItem instances, the same ones passed into WaitAny. For
    WorkflowItems the return values will be WorkflowItems if the work is not
    finished yet, and the final return value once work is finished.
    """

    def __init__(self, items):
        """Initializer.

        Args:
            items: List of WorkItems to wait for.
        """
        self.items = items

class Barrier(list):
    """Barrier for running multiple WorkItems in parallel."""

    def __init__(self, workflow, generator, work):
        """Initializer.

        Args:
            workflow: WorkflowItem instance this is for.
            generator: Current state of the WorkflowItem's generator.
            work: Next set of work to do. May be a single WorkItem object or
                a list or tuple that contains a set of WorkItems to run in
                parallel.
        """
        list.__init__(self)
        self.workflow = workflow
        self.generator = generator

        if isinstance(work, (list, tuple)):
            self[:] = list(work)
            self.was_list = True
            self.wait_any = False
        elif isinstance(work, WaitAny):
            self[:] = list(work.items)
            self.was_list = True
            self.wait_any = True
        else:
            self[:] = [work]
            self.was_list = False
            self.wait_any = False

        for item in self:
            assert isinstance(item, WorkItem)

    @property
    def outstanding(self):
        """Returns whether or not this barrier has pending work."""
        # Allow the same WorkItem to be yielded multiple times but not
        # count towards blocking the barrier.
        done_count = 0
        for item in self:
            if not self.wait_any and item.fire_and_forget:
                # Only count fire_and_forget items as done if this is
                # *not* a WaitAny barrier. We only want to return control
                # to the caller when at least one of the blocking items
                # has completed.
                done_count += 1
            elif item.done:
                done_count += 1

        if self.wait_any and done_count > 0:
            return False

        if done_count == len(self):
            return False

        return True

    @property
    def error(self):
        """Returns the error for this barrier and all work items, if any."""
        # Copy the error from any failed item to be the error for the whole
        # barrier. The first error seen "wins". Also handles the case where
        # the WorkItems passed into the barrier have already completed and
        # been marked with errors.
        for item in self:
            if isinstance(item, WorkItem) and item.error:
                return item.error
        return None

    def get_item(self):
        """Returns the item to send back into the workflow generator."""
        if self.was_list:
            blocking_items = self[:]
            self[:] = []
            for item in blocking_items:
                if (isinstance(item, WorkflowItem) and
                        item.done and
                        not item.error):
                    self.append(item.result)
                else:
                    self.append(item)
            return self
        else:
            return self[0]

class Return(Exception):
    """Raised in WorkflowItem.run to return a result to the caller."""

    def __init__(self, result=None):
        """Initializer.

        Args:
            result: Result of a WorkflowItem, if any.
        """
        self.result = result

class WorkflowThread(WorkerThread):
    """Worker thread for running workflows."""

    def __init__(self, input_queue, output_queue):
        """Initializer.

        Args:
            input_queue: Queue this worker consumes work from. These should be
                WorkflowItems to process, or any WorkItems registered with this
                class using the register() method.
            output_queue: Queue where this worker puts finished work items,
                if any.
        """
        WorkerThread.__init__(self, input_queue, output_queue)
        self.pending = {}
        self.work_map = {}
        self.worker_threads = []
        self.register(WorkflowItem, input_queue)

    # TODO: Implement drain, to let all existing work finish but no new work
    # allowed at the top of the funnel.

    def start(self):
        """Starts the coordinator thread and all related worker threads."""
        assert not self.interrupted
        for thread in self.worker_threads:
            thread.start()
        WorkerThread.start(self)

    def stop(self):
        """Stops the coordinator thread and all related threads."""
        if self.interrupted:
            return
        for thread in self.worker_threads:
            thread.interrupted = True
        self.interrupted = True

    def join(self):
        """Joins the coordinator thread and all worker threads."""
        for thread in self.worker_threads:
            thread.join()
        WorkerThread.join(self)

    def wait_one(self):
        """Waits until this worker has finished one work item or died."""
        while True:
            try:
                item = self.output_queue.get(True, 1)
            except Queue.Empty:
                continue
            except KeyboardInterrupt:
                logging.debug('Exiting')
                return
            else:
                item.check_result()
                return

    def register(self, work_type, queue):
        """Registers where work for a specific type can be executed.

        Args:
            work_type: Sub-class of WorkItem to register.
            queue: Queue instance where WorkItems of the work_type should be
                enqueued when they are yielded by WorkflowItems being run by
                this worker.
        """
        self.work_map[work_type] = queue

    def enqueue_barrier(self, barrier):
        for item in barrier:
            if item.done:
                # Don't reenqueue items that are already done.
                continue

            # Try to find a queue by the work item's class, then by its
            # super class, and so on.
            next_types = [type(item)]
            while next_types:
                try_types = next_types[:]
                next_types[:] = []
                for current_type in try_types:
                    target_queue = self.work_map.get(current_type)
                    if target_queue:
                        next_types = []
                        break
                    next_types.extend(current_type.__bases__)

            assert target_queue, 'Could not find queue to handle %r' % item

            if not item.fire_and_forget:
                self.pending[item] = barrier
            target_queue.put(item)

    def dequeue_barrier(self, item):
        # This is a WorkItem from a worker thread that has finished and
        # needs to be reinjected into a WorkflowItem generator.
        barrier = self.pending.pop(item, None)
        if not barrier:
            # Item was already finished in another barrier, or was
            # fire-and-forget and never part of a barrier; ignore it.
            return None

        if barrier.outstanding and not barrier.error:
            # More work to do and no error seen. Keep waiting.
            return None

        for work in barrier:
            # The barrier has been fulfilled one way or another. Clear out any
            # other pending parts of the barrier so they don't trigger again.
            self.pending.pop(work, None)

        return barrier

    def handle_item(self, item):
        if isinstance(item, WorkflowItem) and not item.done:
            workflow = item
            try:
                generator = item.run(*item.args, **item.kwargs)
            except TypeError, e:
                raise TypeError('%s: item=%r', e, item)
            item = None
        else:
            barrier = self.dequeue_barrier(item)
            if not barrier:
                return
            item = barrier.get_item()
            workflow = barrier.workflow
            generator = barrier.generator

        while True:
            logging.debug('Transitioning workflow=%r, generator=%r, item=%r',
                          workflow, generator, item)
            try:
                try:
                    error = item is not None and item.error
                    if error:
                        next_item = generator.throw(*error)
                    elif isinstance(item, WorkflowItem) and item.done:
                        next_item = generator.send(item.result)
                    else:
                        next_item = generator.send(item)
                except StopIteration:
                    workflow.done = True
                except Return, e:
                    workflow.done = True
                    workflow.result = e.result
                except Exception, e:
                    logging.exception(
                        'error workflow=%r, generator=%r, item=%r',
                        workflow, generator, item)
                    workflow.done = True
                    workflow.error = sys.exc_info()
            finally:
                if workflow.done:
                    if workflow.root:
                        # Root workflow finished. This goes to the output
                        # queue so it can be received by the main thread.
                        return workflow
                    else:
                        # Sub-workflow finished. Reinject it into the
                        # workflow so a pending parent can catch it.
                        self.input_queue.put(workflow)
                        return

            barrier = Barrier(workflow, generator, next_item)
            self.enqueue_barrier(barrier)
            if barrier.outstanding:
                break

            # If a returned barrier has no oustanding parts, immediately
            # progress the workflow.
            item = barrier.get_item()

class PrintWorkflow(WorkflowItem):
    """Prints a message to stdout."""

    def run(self, message):
        yield []  # Make this into a generator
        print message

def get_coordinator():
    """Creates a coordinator and returns it."""
    workflow_queue = Queue.Queue()
    complete_queue = Queue.Queue()
    coordinator = WorkflowThread(workflow_queue, complete_queue)
    coordinator.register(WorkflowItem, workflow_queue)
    return coordinator
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import absolute_import, print_function

import inspect
import os
import subprocess
import time
from multiprocessing import Pool
from sys import stdout
from random import randint

import realhud

from tweepy.error import TweepError
import metadata_parser
import urllib.request
import mimetypes
import uuid

DRAGONFIRE_PATH = os.path.dirname(
    os.path.abspath(inspect.getfile(inspect.currentframe())))
FNULL = open(os.devnull, 'w')

songRunning = False

class TTA:
    def __init__(self, args):
        self.command = None
        self.message = None

        self.headless = args["headless"]
        self.silent = args["silent"]
        self.twitter = args["twitter"]
        if self.twitter:
            self.headless = True
            self.silent = True
        self.twitter_api = None
        self.twitter_user = None
        realhud.load_gif(DRAGONFIRE_PATH + "/realhud/animation/avatar.gif")

    def define(self, com="", msg="", sp="False"):
        self.command = com
        self.message = msg
        self.speak = sp

    def execute(self, duration):
        if self.twitter:
            return True
        try:
            subprocess.Popen(["notify-send", "Dragonfire", self.message])
        except BaseException:
            pass
        if self.command != "":
            time.sleep(duration)
            try:
                subprocess.Popen(self.command, stdout=FNULL, stderr=FNULL)
            except BaseException:
                pass
        # if self.speak == True:
        #   self.say(self.message)
        # else:

    def define_and_execute(self, com="", msg="", sp="False", duration=0):
        self.command = com
        self.message = msg
        self.speak = sp
        if self.twitter:
            return True
        try:
            subprocess.Popen(["notify-send", "Dragonfire", self.message])
        except BaseException:
            pass
        if self.command != "":
            time.sleep(duration)
            try:
                subprocess.Popen(self.command, stdout=FNULL, stderr=FNULL)
            except BaseException:
                pass

    def say(self, message, dynamic=False, end=False):
        if self.twitter:
            text = "@" + self.twitter_user + " " + message#.upper()
            text = (text[:280]) if len(text) > 280 else text
            if len(self.command) > 1:
                if self.command[0] == "sensible-browser":
                    text = text[:((len(self.command[1])+1)*(-1))] + " " + self.command[1]
                    page = metadata_parser.MetadataParser(url=self.command[1])
                    img_url = page.get_metadata_link('image')
                    if img_url:
                        response = urllib.request.urlopen(img_url)
                        img_data = response.read()
                        img_extension = mimetypes.guess_extension(response.headers['content-type'])
                        filename = "/tmp/tmp" + uuid.uuid4().hex + img_extension
                        with open(filename, 'wb') as f:
                            f.write(img_data)

                        try:
                            self.twitter_api.update_with_media(filename, text)
                            if randint(1,3) == 1:
                                self.twitter_api.create_friendship(self.twitter_user, follow=True)
                        except TweepError as e:
                            print("Warning: " + e.response.text)
                        finally:
                            os.remove(filename)
                        return True
            try:
                self.twitter_api.update_status(text)
                if randint(1,3) == 1:
                    self.twitter_api.create_friendship(self.twitter_user, follow=True)
            except TweepError as e:
                print("Warning: " + e.response.text)
            return True
        # if songRunning == True:
        #   subprocess.Popen(["rhythmbox-client","--pause"])
        if len(message) < 10000:
            if dynamic:
                if end:
                    print(message.upper())
                    print(
                        "_______________________________________________________________\n"
                    )
                else:
                    print("Dragonfire: " + message.upper(), end=' ')
                    stdout.flush()
            else:
                print("Dragonfire: " + message.upper())
                print(
                    "_______________________________________________________________\n"
                )
        if not self.silent:
            tts_proc = subprocess.Popen(
                "flite -voice slt -f /dev/stdin",
                stdin=subprocess.PIPE,
                stdout=FNULL,
                stderr=FNULL,
                shell=True)
            message = "".join([i if ord(i) < 128 else ' ' for i in message])
            tts_proc.stdin.write(message.encode())
            tts_proc.stdin.close()
            # print "TTS process started."

        pool = Pool(processes=1)
        if not self.headless:
            pool.apply_async(realhud.play_gif, [0.5, True])
            # print "Avatar process started."

        if not self.silent:
            tts_proc.wait()
        pool.terminate()
        # if songRunning == True:
        #   subprocess.Popen(["rhythmbox-client","--play"])

    def espeak(self, message):
        subprocess.Popen(["espeak", "-v", "en-uk-north", message])

if __name__ == "__main__":
    userin = TTA()
    userin.define([" "], " ")
    userin.say("Hello world!")
<commit_msg>Fix the character reduction method for Twitter character limit<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

from __future__ import absolute_import, print_function

import inspect
import os
import subprocess
import time
from multiprocessing import Pool
from sys import stdout
from random import randint

import realhud

from tweepy.error import TweepError
import metadata_parser
import urllib.request
import mimetypes
import uuid

DRAGONFIRE_PATH = os.path.dirname(
    os.path.abspath(inspect.getfile(inspect.currentframe())))
FNULL = open(os.devnull, 'w')
TWITTER_CHAR_LIMIT = 280

songRunning = False

class TTA:
    def __init__(self, args):
        self.command = None
        self.message = None

        self.headless = args["headless"]
        self.silent = args["silent"]
        self.twitter = args["twitter"]
        if self.twitter:
            self.headless = True
            self.silent = True
        self.twitter_api = None
        self.twitter_user = None
        realhud.load_gif(DRAGONFIRE_PATH + "/realhud/animation/avatar.gif")

    def define(self, com="", msg="", sp="False"):
        self.command = com
        self.message = msg
        self.speak = sp

    def execute(self, duration):
        if self.twitter:
            return True
        try:
            subprocess.Popen(["notify-send", "Dragonfire", self.message])
        except BaseException:
            pass
        if self.command != "":
            time.sleep(duration)
            try:
                subprocess.Popen(self.command, stdout=FNULL, stderr=FNULL)
            except BaseException:
                pass
        # if self.speak == True:
        #   self.say(self.message)
        # else:

    def define_and_execute(self, com="", msg="", sp="False", duration=0):
        self.command = com
        self.message = msg
        self.speak = sp
        if self.twitter:
            return True
        try:
            subprocess.Popen(["notify-send", "Dragonfire", self.message])
        except BaseException:
            pass
        if self.command != "":
            time.sleep(duration)
            try:
                subprocess.Popen(self.command, stdout=FNULL, stderr=FNULL)
            except BaseException:
                pass

    def say(self, message, dynamic=False, end=False):
        if self.twitter:
            text = "@" + self.twitter_user + " " + message#.upper()
            text = (text[:TWITTER_CHAR_LIMIT]) if len(text) > TWITTER_CHAR_LIMIT else text
            if len(self.command) > 1:
                if self.command[0] == "sensible-browser":
                    reduction = len(text + " " + self.command[1]) - TWITTER_CHAR_LIMIT
                    if reduction < 1:
                        reduction = None
                    text = text[:reduction] + " " + self.command[1]
                    page = metadata_parser.MetadataParser(url=self.command[1])
                    img_url = page.get_metadata_link('image')
                    if img_url:
                        response = urllib.request.urlopen(img_url)
                        img_data = response.read()
                        img_extension = mimetypes.guess_extension(response.headers['content-type'])
                        filename = "/tmp/tmp" + uuid.uuid4().hex + img_extension
                        with open(filename, 'wb') as f:
                            f.write(img_data)

                        try:
                            self.twitter_api.update_with_media(filename, text)
                            if randint(1,3) == 1:
                                self.twitter_api.create_friendship(self.twitter_user, follow=True)
                        except TweepError as e:
                            print("Warning: " + e.response.text)
                        finally:
                            os.remove(filename)
                        return True
            try:
                self.twitter_api.update_status(text)
                if randint(1,3) == 1:
                    self.twitter_api.create_friendship(self.twitter_user, follow=True)
            except TweepError as e:
                print("Warning: " + e.response.text)
            return True
        # if songRunning == True:
        #   subprocess.Popen(["rhythmbox-client","--pause"])
        if len(message) < 10000:
            if dynamic:
                if end:
                    print(message.upper())
                    print(
                        "_______________________________________________________________\n"
                    )
                else:
                    print("Dragonfire: " + message.upper(), end=' ')
                    stdout.flush()
            else:
                print("Dragonfire: " + message.upper())
                print(
                    "_______________________________________________________________\n"
                )
        if not self.silent:
            tts_proc = subprocess.Popen(
                "flite -voice slt -f /dev/stdin",
                stdin=subprocess.PIPE,
                stdout=FNULL,
                stderr=FNULL,
                shell=True)
            message = "".join([i if ord(i) < 128 else ' ' for i in message])
            tts_proc.stdin.write(message.encode())
            tts_proc.stdin.close()
            # print "TTS process started."

        pool = Pool(processes=1)
        if not self.headless:
            pool.apply_async(realhud.play_gif, [0.5, True])
            # print "Avatar process started."

        if not self.silent:
            tts_proc.wait()
        pool.terminate()
        # if songRunning == True:
        #   subprocess.Popen(["rhythmbox-client","--play"])

    def espeak(self, message):
        subprocess.Popen(["espeak", "-v", "en-uk-north", message])

if __name__ == "__main__":
    userin = TTA()
    userin.define([" "], " ")
    userin.say("Hello world!")
<|endoftext|>


<commit_before># coding=utf-8
import operator
from collections import OrderedDict

import coreschema
import uritemplate
from coreapi import Link, Document, Field
from coreapi.compat import force_text
from django.db import models
from django.utils.functional import Promise
from pkg_resources import parse_version
from rest_framework import serializers
from rest_framework.fields import IntegerField, URLField
from rest_framework.pagination import PageNumberPagination, LimitOffsetPagination, CursorPagination
from rest_framework.schemas import SchemaGenerator, insert_into, get_pk_description, field_to_schema

from drf_openapi.codec import _get_parameters

class VersionedSerializers:
    """Adapted from https://github.com/avanov/Rhetoric/ :)
    """
    OPERATORS = {
        '>': operator.gt,
        '<': operator.lt,
        '==': operator.eq,
        '>=': operator.ge,
        '<=': operator.le
    }

    """
    A map of version and serializer definition
    May be represented in the following form

    1. ``VERSION``
    2. ``==VERSION`` (the same as above)
    3. ``>VERSION``
    4. ``<VERSION``
    5. ``>=VERSION``
    6. ``<=Version``
    7. Comma-separated list of 1-7 evaluated as AND

    Must override in subclass, for example

    VERSION_MAP = (
        ('>1.3, <=1.6', MeSerializer16)
        ('>1.6', MeSerializer)
    )

    """
    VERSION_MAP = ()

    @classmethod
    def get(cls, request_version):
        for allowed_version, schema in cls.VERSION_MAP:
            distinct_versions = [version.strip() for version in allowed_version.split(',')]
            matched = True
            for distinct_version in distinct_versions:
                operation = cls.OPERATORS.get(distinct_version[:2])
                if operation:
                    # prepare cases #2, #5, #6
                    compare_with = distinct_version[2:]
                else:
                    operation = cls.OPERATORS.get(distinct_version[0])
                    if operation:
                        # prepare cases #3, #4
                        compare_with = distinct_version[1:]
                    else:
                        # prepare case #1
                        compare_with = distinct_version
                        operation = cls.OPERATORS['==']

                matched = operation(parse_version(request_version), parse_version(compare_with))
                if not matched:
                    matched = False

            if matched:
                return schema

        raise ValueError('Invalid request version {}'.format(request_version))
    get.__func__.__annotations__ = {'request_version': str}

class OpenApiSchemaGenerator(SchemaGenerator):
    def __init__(self, version, title=None, url=None, description=None, patterns=None, urlconf=None):
        self.version = version
        super(OpenApiSchemaGenerator, self).__init__(title, url, description, patterns, urlconf)

    def get_schema(self, request=None, public=False):
        if self.endpoints is None:
            inspector = self.endpoint_inspector_cls(self.patterns, self.urlconf)
            self.endpoints = inspector.get_api_endpoints()

        links = self.get_links(None if public else request)
        if not links:
            return None

        url = self.url
        if not url and request is not None:
            url = request.build_absolute_uri()

        return OpenApiDocument(
            version=self.version,
            title=self.title, description=self.description,
            url=url, content=links
        )

    def get_links(self, request=None):
        """
        Return a dictionary containing all the links that should be
        included in the API schema.
        """
        links = OrderedDict()

        # Generate (path, method, view) given (path, method, callback).
        paths = []
        view_endpoints = []
        for path, method, callback in self.endpoints:
            view = self.create_view(callback, method, request)
            if getattr(view, 'exclude_from_schema', False):
                continue
            path = self.coerce_path(path, method, view)
            paths.append(path)
            view_endpoints.append((path, method, view))

        # Only generate the path prefix for paths that will be included
        if not paths:
            return None
        prefix = self.determine_path_prefix(paths)

        for path, method, view in view_endpoints:
            if not self.has_view_permissions(path, method, view):
                continue
            link = self.get_link(path, method, view, version=request.version)
            subpath = path[len(prefix):]
            keys = self.get_keys(subpath, method, view)
            try:
                insert_into(links, keys, link)
            except Exception:
                continue
        return links

    def get_serializer_doc(self, serializer):
        if serializer.__doc__ is None:
            return ''

        doc = []
        for line in serializer.__doc__.splitlines():
            doc.append(line.strip())
        return '\n'.join(doc)

    def get_link(self, path, method, view, version=None):
        fields = self.get_path_fields(path, method, view)
        fields += self.get_serializer_fields(path, method, view, version=version)
        fields += self.get_pagination_fields(path, method, view)
        fields += self.get_filter_fields(path, method, view)

        if fields and any([field.location in ('form', 'body') for field in fields]):
            encoding = self.get_encoding(path, method, view)
        else:
            encoding = None

        description = self.get_description(path, method, view)

        method_name = getattr(view, 'action', method.lower())
        method_func = getattr(view, method_name, None)

        request_serializer_class = getattr(method_func, 'request_serializer', None)
        if request_serializer_class and issubclass(request_serializer_class, VersionedSerializers):
            request_doc = self.get_serializer_doc(request_serializer_class)
            if request_doc:
                description = description + '\n\n**Request Description:**\n' + request_doc

        response_serializer_class = getattr(method_func, 'response_serializer', None)
        if response_serializer_class and issubclass(response_serializer_class, VersionedSerializers):
            res_doc = self.get_serializer_doc(response_serializer_class)
            if res_doc:
                description = description + '\n\n**Response Description:**\n' + res_doc
            response_serializer_class = response_serializer_class.get(version)

        if not response_serializer_class and method_name in ('list', 'retrieve'):
            if hasattr(view, 'get_serializer_class'):
                response_serializer_class = view.get_serializer_class()
            elif hasattr(view, 'serializer_class'):
                response_serializer_class = view.serializer_class
            if response_serializer_class and method_name == 'list':
                response_serializer_class = self.get_paginator_serializer(
                    view, response_serializer_class)
        response_schema, error_status_codes = self.get_response_object(
            response_serializer_class, method_func.__doc__) if response_serializer_class else ({}, {})

        return OpenApiLink(
            response_schema=response_schema,
            error_status_codes=error_status_codes,
            url=path.replace('{version}', self.version),  # can't use format because there may be other param
            action=method.lower(),
            encoding=encoding,
            fields=fields,
            description=description
        )

    def get_paginator_serializer(self, view, child_serializer_class):
        class BaseFakeListSerializer(serializers.Serializer):
            results = child_serializer_class(many=True)

        class FakePrevNextListSerializer(BaseFakeListSerializer):
            next = URLField()
            previous = URLField()

        pager = view.pagination_class
        if hasattr(pager, 'default_pager'):
            # Must be a ProxyPagination
            pager = pager.default_pager

        if issubclass(pager, (PageNumberPagination, LimitOffsetPagination)):
            class FakeListSerializer(FakePrevNextListSerializer):
                count = IntegerField()
            return FakeListSerializer
        elif issubclass(pager, CursorPagination):
            return FakePrevNextListSerializer

        return BaseFakeListSerializer

    def get_path_fields(self, path, method, view):
        """
        Return a list of `coreapi.Field` instances corresponding to any
        templated path variables.
        """
        model = getattr(getattr(view, 'queryset', None), 'model', None)
        fields = []

        for variable in uritemplate.variables(path):

            if variable == 'version':
                continue

            title = ''
            description = ''
            schema_cls = coreschema.String
            kwargs = {}
            if model is not None:
                # Attempt to infer a field description if possible.
                try:
                    model_field = model._meta.get_field(variable)
                except:
                    model_field = None

                if model_field is not None and model_field.verbose_name:
                    title = force_text(model_field.verbose_name)

                if model_field is not None and model_field.help_text:
                    description = force_text(model_field.help_text)
                elif model_field is not None and model_field.primary_key:
                    description = get_pk_description(model, model_field)

                if hasattr(view, 'lookup_value_regex') and view.lookup_field == variable:
                    kwargs['pattern'] = view.lookup_value_regex
                elif isinstance(model_field, models.AutoField):
                    schema_cls = coreschema.Integer

            field = Field(
                name=variable,
                location='path',
                required=True,
                schema=schema_cls(title=title, description=description, **kwargs)
            )
            fields.append(field)

        return fields

    def get_serializer_fields(self, path, method, view, version=None):
        """
        Return a list of `coreapi.Field` instances corresponding to any
        request body input, as determined by the serializer class.
        """
        if method not in ('PUT', 'PATCH', 'POST'):
            return []

        if not hasattr(view, 'serializer_class') and not hasattr(view, 'get_serializer_class'):
            return []

        serializer_class = view.get_serializer_class() if hasattr(view, 'get_serializer_class') \
            else view.serializer_class
        serializer = serializer_class()

        if isinstance(serializer, serializers.ListSerializer):
            return [
                Field(
                    name='data',
                    location='body',
                    required=True,
                    schema=coreschema.Array()
                )
            ]

        if not isinstance(serializer, serializers.Serializer):
            return []

        fields = []
        for field in serializer.fields.values():
            if field.read_only or isinstance(field, serializers.HiddenField):
                continue

            required = field.required and method != 'PATCH'
            # if the attribute ('help_text') of this field is a lazy translation object, force it to generate a string
            description = str(field.help_text) if isinstance(field.help_text, Promise) else field.help_text
            field = Field(
                name=field.field_name,
                location='form',
                required=required,
                schema=field_to_schema(field),
                description=description,
            )
            fields.append(field)

        return fields

    def get_response_object(self, response_serializer_class, description):

        fields = []
        serializer = response_serializer_class()
        nested_obj = {}

        for field in serializer.fields.values():
            # If field is a serializer, attempt to get its schema.
            if isinstance(field, serializers.Serializer):
                subfield_schema = self.get_response_object(field.__class__, None)[0].get('schema')

                # If the schema exists, use it as the nested_obj
                if subfield_schema is not None:
                    nested_obj[field.field_name] = subfield_schema
                    nested_obj[field.field_name]['description'] = field.help_text
                    continue

            # Otherwise, carry-on and use the field's schema.
            fields.append(Field(
                name=field.field_name,
                location='form',
                required=field.required,
                schema=field_to_schema(field)
            ))

        res = _get_parameters(Link(fields=fields), None)

        if not res:
            if nested_obj:
                return {
                    'description': description,
                    'schema': {
                        'type': 'object',
                        'properties': nested_obj
                    }
                }, {}
            else:
                return {}, {}

        schema = res[0]['schema']
        schema['properties'].update(nested_obj)
        response_schema = {
            'description': description,
            'schema': schema
        }

        error_status_codes = {}

        response_meta = getattr(response_serializer_class, 'Meta', None)

        for status_code, description in getattr(response_meta, 'error_status_codes', {}).items():
            error_status_codes[status_code] = {'description': description}

        return response_schema, error_status_codes

class OpenApiDocument(Document):
    """OpenAPI-compliant document provides:
    - Versioning information
    """

    def __init__(self, version, url=None, title=None, description=None, media_type=None, content=None):
        super(OpenApiDocument, self).__init__(
            url=url,
            title=title,
            description=description,
            media_type=media_type,
            content=content
        )
        self._version = version

    @property
    def version(self):
        return self._version

class OpenApiLink(Link):
    """OpenAPI-compliant Link provides:
    - Schema to the response
    """

    def __init__(self, response_schema, error_status_codes,
                 url=None, action=None, encoding=None, transform=None, title=None,
                 description=None, fields=None):
        super(OpenApiLink, self).__init__(
            url=url,
            action=action,
            encoding=encoding,
            transform=transform,
            title=title,
            description=description,
            fields=fields
        )
        self._response_schema = response_schema
        self._error_status_codes = error_status_codes

    @property
    def response_schema(self):
        return self._response_schema

    @property
    def error_status_codes(self):
        return self._error_status_codes
<commit_msg>Fix imports for DRF 3.7 compat<commit_after># coding=utf-8
import operator
from collections import OrderedDict

import coreschema
import uritemplate
from coreapi import Link, Document, Field
from coreapi.compat import force_text
from django.db import models
from django.utils.functional import Promise
from pkg_resources import parse_version
from rest_framework import serializers
from rest_framework.fields import IntegerField, URLField
from rest_framework.pagination import PageNumberPagination, LimitOffsetPagination, CursorPagination
from rest_framework.schemas import SchemaGenerator
from rest_framework.schemas.generators import insert_into
from rest_framework.schemas.inspectors import get_pk_description, field_to_schema
from rest_framework.schemas.utils import is_list_view

from drf_openapi.codec import _get_parameters

class VersionedSerializers:
    """Adapted from https://github.com/avanov/Rhetoric/ :)
    """
    OPERATORS = {
        '>': operator.gt,
        '<': operator.lt,
        '==': operator.eq,
        '>=': operator.ge,
        '<=': operator.le
    }

    """
    A map of version and serializer definition
    May be represented in the following form

    1. ``VERSION``
    2. ``==VERSION`` (the same as above)
    3. ``>VERSION``
    4. ``<VERSION``
    5. ``>=VERSION``
    6. ``<=Version``
    7. Comma-separated list of 1-7 evaluated as AND

    Must override in subclass, for example

    VERSION_MAP = (
        ('>1.3, <=1.6', MeSerializer16)
        ('>1.6', MeSerializer)
    )

    """
    VERSION_MAP = ()

    @classmethod
    def get(cls, request_version):
        for allowed_version, schema in cls.VERSION_MAP:
            distinct_versions = [version.strip() for version in allowed_version.split(',')]
            matched = True
            for distinct_version in distinct_versions:
                operation = cls.OPERATORS.get(distinct_version[:2])
                if operation:
                    # prepare cases #2, #5, #6
                    compare_with = distinct_version[2:]
                else:
                    operation = cls.OPERATORS.get(distinct_version[0])
                    if operation:
                        # prepare cases #3, #4
                        compare_with = distinct_version[1:]
                    else:
                        # prepare case #1
                        compare_with = distinct_version
                        operation = cls.OPERATORS['==']

                matched = operation(parse_version(request_version), parse_version(compare_with))
                if not matched:
                    matched = False

            if matched:
                return schema

        raise ValueError('Invalid request version {}'.format(request_version))
    get.__func__.__annotations__ = {'request_version': str}

class OpenApiSchemaGenerator(SchemaGenerator):
    def __init__(self, version, title=None, url=None, description=None, patterns=None, urlconf=None):
        self.version = version
        super(OpenApiSchemaGenerator, self).__init__(title, url, description, patterns, urlconf)

    def get_schema(self, request=None, public=False):
        if self.endpoints is None:
            inspector = self.endpoint_inspector_cls(self.patterns, self.urlconf)
            self.endpoints = inspector.get_api_endpoints()

        links = self.get_links(None if public else request)
        if not links:
            return None

        url = self.url
        if not url and request is not None:
            url = request.build_absolute_uri()

        return OpenApiDocument(
            version=self.version,
            title=self.title, description=self.description,
            url=url, content=links
        )

    def get_links(self, request=None):
        """
        Return a dictionary containing all the links that should be
        included in the API schema.
        """
        links = OrderedDict()

        # Generate (path, method, view) given (path, method, callback).
        paths = []
        view_endpoints = []
        for path, method, callback in self.endpoints:
            view = self.create_view(callback, method, request)
            if getattr(view, 'exclude_from_schema', False):
                continue
            path = self.coerce_path(path, method, view)
            paths.append(path)
            view_endpoints.append((path, method, view))

        # Only generate the path prefix for paths that will be included
        if not paths:
            return None
        prefix = self.determine_path_prefix(paths)

        for path, method, view in view_endpoints:
            if not self.has_view_permissions(path, method, view):
                continue
            link = self.get_link(path, method, view, version=request.version)
            subpath = path[len(prefix):]
            keys = self.get_keys(subpath, method, view)
            try:
                insert_into(links, keys, link)
            except Exception:
                continue
        return links

    def get_serializer_doc(self, serializer):
        if serializer.__doc__ is None:
            return ''

        doc = []
        for line in serializer.__doc__.splitlines():
            doc.append(line.strip())
        return '\n'.join(doc)

    def get_link(self, path, method, view, version=None):
        fields = self.get_path_fields(path, method, view)
        fields += self.get_serializer_fields(path, method, view, version=version)
        fields += self.get_pagination_fields(path, method, view)
        fields += self.get_filter_fields(path, method, view)

        if fields and any([field.location in ('form', 'body') for field in fields]):
            encoding = self.get_encoding(path, method, view)
        else:
            encoding = None

        description = self.get_description(path, method, view)

        method_name = getattr(view, 'action', method.lower())
        method_func = getattr(view, method_name, None)

        request_serializer_class = getattr(method_func, 'request_serializer', None)
        if request_serializer_class and issubclass(request_serializer_class, VersionedSerializers):
            request_doc = self.get_serializer_doc(request_serializer_class)
            if request_doc:
                description = description + '\n\n**Request Description:**\n' + request_doc

        response_serializer_class = getattr(method_func, 'response_serializer', None)
        if response_serializer_class and issubclass(response_serializer_class, VersionedSerializers):
            res_doc = self.get_serializer_doc(response_serializer_class)
            if res_doc:
                description = description + '\n\n**Response Description:**\n' + res_doc
            response_serializer_class = response_serializer_class.get(version)

        if not response_serializer_class and method_name in ('list', 'retrieve'):
            if hasattr(view, 'get_serializer_class'):
                response_serializer_class = view.get_serializer_class()
            elif hasattr(view, 'serializer_class'):
                response_serializer_class = view.serializer_class
            if response_serializer_class and method_name == 'list':
                response_serializer_class = self.get_paginator_serializer(
                    view, response_serializer_class)
        response_schema, error_status_codes = self.get_response_object(
            response_serializer_class, method_func.__doc__) if response_serializer_class else ({}, {})

        return OpenApiLink(
            response_schema=response_schema,
            error_status_codes=error_status_codes,
            url=path.replace('{version}', self.version),  # can't use format because there may be other param
            action=method.lower(),
            encoding=encoding,
            fields=fields,
            description=description
        )

    def get_paginator_serializer(self, view, child_serializer_class):
        class BaseFakeListSerializer(serializers.Serializer):
            results = child_serializer_class(many=True)

        class FakePrevNextListSerializer(BaseFakeListSerializer):
            next = URLField()
            previous = URLField()

        pager = view.pagination_class
        if hasattr(pager, 'default_pager'):
            # Must be a ProxyPagination
            pager = pager.default_pager

        if issubclass(pager, (PageNumberPagination, LimitOffsetPagination)):
            class FakeListSerializer(FakePrevNextListSerializer):
                count = IntegerField()
            return FakeListSerializer
        elif issubclass(pager, CursorPagination):
            return FakePrevNextListSerializer

        return BaseFakeListSerializer

    def get_path_fields(self, path, method, view):
        """
        Return a list of `coreapi.Field` instances corresponding to any
        templated path variables.
        """
        model = getattr(getattr(view, 'queryset', None), 'model', None)
        fields = []

        for variable in uritemplate.variables(path):

            if variable == 'version':
                continue

            title = ''
            description = ''
            schema_cls = coreschema.String
            kwargs = {}
            if model is not None:
                # Attempt to infer a field description if possible.
                try:
                    model_field = model._meta.get_field(variable)
                except:
                    model_field = None

                if model_field is not None and model_field.verbose_name:
                    title = force_text(model_field.verbose_name)

                if model_field is not None and model_field.help_text:
                    description = force_text(model_field.help_text)
                elif model_field is not None and model_field.primary_key:
                    description = get_pk_description(model, model_field)

                if hasattr(view, 'lookup_value_regex') and view.lookup_field == variable:
                    kwargs['pattern'] = view.lookup_value_regex
                elif isinstance(model_field, models.AutoField):
                    schema_cls = coreschema.Integer

            field = Field(
                name=variable,
                location='path',
                required=True,
                schema=schema_cls(title=title, description=description, **kwargs)
            )
            fields.append(field)

        return fields

    def get_serializer_fields(self, path, method, view, version=None):
        """
        Return a list of `coreapi.Field` instances corresponding to any
        request body input, as determined by the serializer class.
        """
        if method not in ('PUT', 'PATCH', 'POST'):
            return []

        if not hasattr(view, 'serializer_class') and not hasattr(view, 'get_serializer_class'):
            return []

        serializer_class = view.get_serializer_class() if hasattr(view, 'get_serializer_class') \
            else view.serializer_class
        serializer = serializer_class()

        if isinstance(serializer, serializers.ListSerializer):
            return [
                Field(
                    name='data',
                    location='body',
                    required=True,
                    schema=coreschema.Array()
                )
            ]

        if not isinstance(serializer, serializers.Serializer):
            return []

        fields = []
        for field in serializer.fields.values():
            if field.read_only or isinstance(field, serializers.HiddenField):
                continue

            required = field.required and method != 'PATCH'
            # if the attribute ('help_text') of this field is a lazy translation object, force it to generate a string
            description = str(field.help_text) if isinstance(field.help_text, Promise) else field.help_text
            field = Field(
                name=field.field_name,
                location='form',
                required=required,
                schema=field_to_schema(field),
                description=description,
            )
            fields.append(field)

        return fields

    def get_response_object(self, response_serializer_class, description):

        fields = []
        serializer = response_serializer_class()
        nested_obj = {}

        for field in serializer.fields.values():
            # If field is a serializer, attempt to get its schema.
            if isinstance(field, serializers.Serializer):
                subfield_schema = self.get_response_object(field.__class__, None)[0].get('schema')

                # If the schema exists, use it as the nested_obj
                if subfield_schema is not None:
                    nested_obj[field.field_name] = subfield_schema
                    nested_obj[field.field_name]['description'] = field.help_text
                    continue

            # Otherwise, carry-on and use the field's schema.
            fields.append(Field(
                name=field.field_name,
                location='form',
                required=field.required,
                schema=field_to_schema(field)
            ))

        res = _get_parameters(Link(fields=fields), None)

        if not res:
            if nested_obj:
                return {
                    'description': description,
                    'schema': {
                        'type': 'object',
                        'properties': nested_obj
                    }
                }, {}
            else:
                return {}, {}

        schema = res[0]['schema']
        schema['properties'].update(nested_obj)
        response_schema = {
            'description': description,
            'schema': schema
        }

        error_status_codes = {}

        response_meta = getattr(response_serializer_class, 'Meta', None)

        for status_code, description in getattr(response_meta, 'error_status_codes', {}).items():
            error_status_codes[status_code] = {'description': description}

        return response_schema, error_status_codes

class OpenApiDocument(Document):
    """OpenAPI-compliant document provides:
    - Versioning information
    """

    def __init__(self, version, url=None, title=None, description=None, media_type=None, content=None):
        super(OpenApiDocument, self).__init__(
            url=url,
            title=title,
            description=description,
            media_type=media_type,
            content=content
        )
        self._version = version

    @property
    def version(self):
        return self._version

class OpenApiLink(Link):
    """OpenAPI-compliant Link provides:
    - Schema to the response
    """

    def __init__(self, response_schema, error_status_codes,
                 url=None, action=None, encoding=None, transform=None, title=None,
                 description=None, fields=None):
        super(OpenApiLink, self).__init__(
            url=url,
            action=action,
            encoding=encoding,
            transform=transform,
            title=title,
            description=description,
            fields=fields
        )
        self._response_schema = response_schema
        self._error_status_codes = error_status_codes

    @property
    def response_schema(self):
        return self._response_schema

    @property
    def error_status_codes(self):
        return self._error_status_codes
<|endoftext|>


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

#
# MIT License
#
# Copyright (c) 2016 Tomas Pavuk <433592@mail.muni.cz>, Institute of Computer Science, Masaryk University
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

"""
Detects external dns resolvers used in the specified local network.

Default output parameters:
    * Address and port of the broker: producer:9092
    * Kafka topic: results.output

Usage:
    dns_external_resolvers.py -iz <input-zookeeper-hostname>:<input-zookeeper-port> -it <input-topic>
    -oz <output-zookeeper-hostname>:<output-zookeeper-port> -ot <output-topic> -ln <local-network>/<subnet-mask>

To run this on the Stream4Flow, you need to receive flows by IPFIXCol and make them available via Kafka topic. Then you
can run the application
    $ ./run-application.sh ./detection/dns_external_resolvers/spark/dns_external_resolvers.py -iz producer:2181\
    -it ipfix.entry -oz producer:9092 -ot results.output -ln 10.10.0.0/16
"""

import argparse  # Arguments parser
import time  # Unix time to timestamp conversion

from netaddr import IPNetwork, IPAddress  # Checking if IP is in the network
from modules import kafkaIO  # IO operations with kafka topics
from modules import DNSResponseConverter  # Convert byte array to the IP address
from termcolor import cprint  # Colors in the console output

def get_output_json(key, value):
    """
    Create JSON with correct format.

    :param key: Source ip address
    :param value: Dictionary value for statistic
    :return: JSON string in desired format
    """

    # Convert Unix time to timestamp
    s, ms = divmod(value[0], 1000)
    timestamp = "%s.%03d" % (time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(s)), ms) + 'Z'

    return "{\"@type\": \"external_dns_resolver\", \"src_ip\": \"" + str(key[0]) + '\"' +\
           ", \"resolver_ip\": \"" + str(key[1]) + '\"' +\
           ", \"flows\": " + str(value[1]) + \
           ", \"timestamp\": \"" + str(timestamp) + "\"}\n"

def process_results(results, producer, s_output_topic):
    """
    Format and report detected records.

    :param results: Detected records
    :param producer: Kafka producer that sends the data to output topic
    :param s_output_topic: Name of the receiving kafka topic
    """

    output_json = ""
    # Transform given results into the JSON
    for key, value in results.iteritems():
        output_json += get_output_json(key, value)

    if output_json:
        # Print data to standard output
        cprint(output_json)

        # Send results to the specified kafka topic
        kafkaIO.send_data_to_kafka(output_json, producer, s_output_topic)

def get_external_dns_resolvers(dns_input_stream, all_data_stream, s_window_duration):
    """
    Gets used external dns resolvers from input stream

    :param dns_input_stream: Input flows
    :param s_window_duration: Length of the window in seconds
    :param all_data_stream: All incoming flows
    :return: Detected external resolvers
    """
    dns_resolved = dns_input_stream\
        .window(s_window_duration, s_window_duration) \
        .filter(lambda record: record["ipfix.DNSCrrType"] == 1) \
        .map(lambda record: ((record["ipfix.destinationIPv4Address"],
                              DNSResponseConverter.convert_dns_rdata(record["ipfix.DNSRData"], record["ipfix.DNSCrrType"])),
                             (record["ipfix.sourceIPv4Address"],
                              record["ipfix.flowStartMilliseconds"])))\

    detected_external = all_data_stream\
        .window(s_window_duration, s_window_duration) \
        .filter(lambda flow_json: flow_json["ipfix.protocolIdentifier"] == 6) \
        .map(lambda record: ((record["ipfix.sourceIPv4Address"], record["ipfix.destinationIPv4Address"]),
                              record["ipfix.flowStartMilliseconds"])) \
        .join(dns_resolved) \
        .filter(lambda record: abs(record[1][0] - record[1][1][1]) <= 5000) \
        .map(lambda record: ((record[0][0], record[1][1][0]), (record[1][1][1], 1))) \
        .reduceByKey(lambda actual, update: (actual[0],
                                             actual[1] + update[1]))

    return detected_external

def get_dns_stream(flows_stream):
    """
    Filter to get only flows containing DNS information.

    :param flows_stream: Input flows
    :return: Flows with DNS information
    """
    return flows_stream \
        .filter(lambda flow_json: ("ipfix.DNSName" in flow_json.keys()) and
                                  ("ipfix.sourceIPv4Address" in flow_json.keys()))

def get_flows_external_to_local(s_dns_stream, local_network):
    """
    Filter to contain flows going from the specified local network to the different network.

    :param s_dns_stream: Input flows
    :param local_network: Local network's address
    :return: Flows coming from local network to external networks
    """
    return s_dns_stream \
        .filter(lambda dns_json: (IPAddress(dns_json["ipfix.sourceIPv4Address"]) not in IPNetwork(local_network)) and
                                 (IPAddress(dns_json["ipfix.destinationIPv4Address"]) in IPNetwork(local_network)))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-iz", "--input_zookeeper", help="input zookeeper hostname:port", type=str, required=True)
    parser.add_argument("-it", "--input_topic", help="input kafka topic", type=str, required=True)
    parser.add_argument("-oz", "--output_zookeeper", help="output zookeeper hostname:port", type=str, required=True)
    parser.add_argument("-ot", "--output_topic", help="output kafka topic", type=str, required=True)
    parser.add_argument("-w", "--window_size", help="window size (in seconds)", type=int, required=False, default=60)
    parser.add_argument("-m", "--microbatch", help="microbatch (in seconds)", type=int, required=False, default=5)

    # Define Arguments for detection
    parser.add_argument("-ln", "--local_network", help="local network", type=str, required=True)

    # Parse arguments
    args = parser.parse_args()

    # Initialize input stream and parse it into JSON
    ssc, parsed_input_stream = kafkaIO\
        .initialize_and_parse_input_stream(args.input_zookeeper, args.input_topic, args.microbatch)

    # Prepare input stream
    dns_stream = get_dns_stream(parsed_input_stream)
    dns_external_to_local = get_flows_external_to_local(dns_stream, args.local_network)
    ipv4_stream = parsed_input_stream.filter(lambda flow_json: ("ipfix.sourceIPv4Address" in flow_json.keys()))

    # Initialize kafka producer
    kafka_producer = kafkaIO.initialize_kafka_producer(args.output_zookeeper)

    # Calculate and process DNS statistics
    get_external_dns_resolvers(dns_external_to_local, ipv4_stream, args.window_size) \
        .foreachRDD(lambda rdd: process_results(rdd.collectAsMap(), kafka_producer, args.output_topic))

    # Start Spark streaming context
    kafkaIO.spark_start(ssc)
<commit_msg>Changed window length for dns external resolvers<commit_after># -*- coding: utf-8 -*-

#
# MIT License
#
# Copyright (c) 2016 Tomas Pavuk <433592@mail.muni.cz>, Institute of Computer Science, Masaryk University
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#

"""
Detects external dns resolvers used in the specified local network.

Default output parameters:
    * Address and port of the broker: producer:9092
    * Kafka topic: results.output

Usage:
    dns_external_resolvers.py -iz <input-zookeeper-hostname>:<input-zookeeper-port> -it <input-topic>
    -oz <output-zookeeper-hostname>:<output-zookeeper-port> -ot <output-topic> -ln <local-network>/<subnet-mask>

To run this on the Stream4Flow, you need to receive flows by IPFIXCol and make them available via Kafka topic. Then you
can run the application
    $ ./run-application.sh ./detection/dns_external_resolvers/spark/dns_external_resolvers.py -iz producer:2181\
    -it ipfix.entry -oz producer:9092 -ot results.output -ln 10.10.0.0/16
"""

import argparse  # Arguments parser
import time  # Unix time to timestamp conversion

from netaddr import IPNetwork, IPAddress  # Checking if IP is in the network
from modules import kafkaIO  # IO operations with kafka topics
from modules import DNSResponseConverter  # Convert byte array to the IP address
from termcolor import cprint  # Colors in the console output

def get_output_json(key, value):
    """
    Create JSON with correct format.

    :param key: Source ip address
    :param value: Dictionary value for statistic
    :return: JSON string in desired format
    """

    # Convert Unix time to timestamp
    s, ms = divmod(value[0], 1000)
    timestamp = "%s.%03d" % (time.strftime("%Y-%m-%dT%H:%M:%S", time.gmtime(s)), ms) + 'Z'

    return "{\"@type\": \"external_dns_resolver\", \"src_ip\": \"" + str(key[0]) + '\"' +\
           ", \"resolver_ip\": \"" + str(key[1]) + '\"' +\
           ", \"flows\": " + str(value[1]) + \
           ", \"timestamp\": \"" + str(timestamp) + "\"}\n"

def process_results(results, producer, s_output_topic):
    """
    Format and report detected records.

    :param results: Detected records
    :param producer: Kafka producer that sends the data to output topic
    :param s_output_topic: Name of the receiving kafka topic
    """

    output_json = ""
    # Transform given results into the JSON
    for key, value in results.iteritems():
        output_json += get_output_json(key, value)

    if output_json:
        # Print data to standard output
        cprint(output_json)

        # Send results to the specified kafka topic
        kafkaIO.send_data_to_kafka(output_json, producer, s_output_topic)

def get_external_dns_resolvers(dns_input_stream, all_data_stream, s_window_duration):
    """
    Gets used external dns resolvers from input stream

    :param dns_input_stream: Input flows
    :param s_window_duration: Length of the window in seconds
    :param all_data_stream: All incoming flows
    :return: Detected external resolvers
    """
    dns_resolved = dns_input_stream\
        .window(s_window_duration, s_window_duration) \
        .filter(lambda record: record["ipfix.DNSCrrType"] == 1) \
        .map(lambda record: ((record["ipfix.destinationIPv4Address"],
                              DNSResponseConverter.convert_dns_rdata(record["ipfix.DNSRData"], record["ipfix.DNSCrrType"])),
                             (record["ipfix.sourceIPv4Address"],
                              record["ipfix.flowStartMilliseconds"])))\

    detected_external = all_data_stream\
        .window(s_window_duration, s_window_duration) \
        .filter(lambda flow_json: flow_json["ipfix.protocolIdentifier"] == 6) \
        .map(lambda record: ((record["ipfix.sourceIPv4Address"], record["ipfix.destinationIPv4Address"]),
                              record["ipfix.flowStartMilliseconds"])) \
        .join(dns_resolved) \
        .filter(lambda record: abs(record[1][0] - record[1][1][1]) <= 5000) \
        .map(lambda record: ((record[0][0], record[1][1][0]), (record[1][1][1], 1))) \
        .reduceByKey(lambda actual, update: (actual[0],
                                             actual[1] + update[1]))

    return detected_external

def get_dns_stream(flows_stream):
    """
    Filter to get only flows containing DNS information.

    :param flows_stream: Input flows
    :return: Flows with DNS information
    """
    return flows_stream \
        .filter(lambda flow_json: ("ipfix.DNSName" in flow_json.keys()) and
                                  ("ipfix.sourceIPv4Address" in flow_json.keys()))

def get_flows_external_to_local(s_dns_stream, local_network):
    """
    Filter to contain flows going from the specified local network to the different network.

    :param s_dns_stream: Input flows
    :param local_network: Local network's address
    :return: Flows coming from local network to external networks
    """
    return s_dns_stream \
        .filter(lambda dns_json: (IPAddress(dns_json["ipfix.sourceIPv4Address"]) not in IPNetwork(local_network)) and
                                 (IPAddress(dns_json["ipfix.destinationIPv4Address"]) in IPNetwork(local_network)))

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-iz", "--input_zookeeper", help="input zookeeper hostname:port", type=str, required=True)
    parser.add_argument("-it", "--input_topic", help="input kafka topic", type=str, required=True)
    parser.add_argument("-oz", "--output_zookeeper", help="output zookeeper hostname:port", type=str, required=True)
    parser.add_argument("-ot", "--output_topic", help="output kafka topic", type=str, required=True)
    parser.add_argument("-w", "--window_size", help="window size (in seconds)", type=int, required=False, default=360)
    parser.add_argument("-m", "--microbatch", help="microbatch (in seconds)", type=int, required=False, default=5)

    # Define Arguments for detection
    parser.add_argument("-ln", "--local_network", help="local network", type=str, required=True)

    # Parse arguments
    args = parser.parse_args()

    # Initialize input stream and parse it into JSON
    ssc, parsed_input_stream = kafkaIO\
        .initialize_and_parse_input_stream(args.input_zookeeper, args.input_topic, args.microbatch)

    # Prepare input stream
    dns_stream = get_dns_stream(parsed_input_stream)
    dns_external_to_local = get_flows_external_to_local(dns_stream, args.local_network)
    ipv4_stream = parsed_input_stream.filter(lambda flow_json: ("ipfix.sourceIPv4Address" in flow_json.keys()))

    # Initialize kafka producer
    kafka_producer = kafkaIO.initialize_kafka_producer(args.output_zookeeper)

    # Calculate and process DNS statistics
    get_external_dns_resolvers(dns_external_to_local, ipv4_stream, args.window_size) \
        .foreachRDD(lambda rdd: process_results(rdd.collectAsMap(), kafka_producer, args.output_topic))

    # Start Spark streaming context
    kafkaIO.spark_start(ssc)
<|endoftext|>