<commit_before>"""Module for running the worker process.

It's an infinite loop.
"""
from __future__ import absolute_import

import logging
import resource
import signal
import sys
import time

from kale import consumer
from kale import publisher
from kale import queue_info
from kale import settings
from kale import timeout
from kale import utils
from six.moves import range

logger = logging.getLogger(__name__)

SIGNALS_TO_HANDLE = (
    signal.SIGABRT,
    # This will catch Ctrl-C interrupts.
    signal.SIGINT,
    signal.SIGQUIT,
    signal.SIGTERM,
    # Handle Ctrl-Z/suspend.
    signal.SIGTSTP)

# Logging constants.
LOG_TASK_RESULT_DEFERRED = 'deferred'
LOG_TASK_RESULT_ERROR = 'error'
LOG_TASK_RESULT_SUCCESS = 'success'

class Worker(object):

    def __init__(self):
        """Initialize a worker instance."""

        self._consumer = consumer.Consumer()

        queue_class = utils.class_import_from_path(settings.QUEUE_CLASS)
        q_info = queue_info.QueueInfo(
            config_file=settings.QUEUE_CONFIG,
            sqs_talk=self._consumer,
            queue_cls=queue_class)

        queue_selector_class = utils.class_import_from_path(
            settings.QUEUE_SELECTOR)
        self._queue_selector = queue_selector_class(q_info)

        # The worker will publish permanently failed tasks to a
        # dead-letter-queue.
        self._publisher = publisher.Publisher()

        # Track total messages processed.
        self._total_messages_processed = 0

        # Intialize queue variables used by each batch.
        self._incomplete_messages = []
        self._successful_messages = []
        self._failed_messages = []
        self._permanent_failures = []
        self._batch_stop_time = time.time()
        self._batch_queue = None
        # Monitors whether the worker has been exposed to tasks and may
        # have bloated in memory.
        self._dirty = False

        # Setup signal handling for cleanup.
        for sig in SIGNALS_TO_HANDLE:
            signal.signal(sig, self._cleanup_worker)

        # Allow the client of this library to do any setup before
        # starting the worker.
        settings.ON_WORKER_STARTUP()

    def _on_pre_run_worker(self):
        """Callback function right at the beginning of starting the worker. """
        logger.info('Starting run loop for task worker.')

    def _on_exceeding_memory_limit(self, ru_maxrss):
        """Callback function when the process exceeds memory limit.

        :param int ru_maxrss: maximum resident set size used (in kilobytes).
        """
        logger.info('Memory usages exceeds max: %d. Exiting.' % ru_maxrss)

    def _on_sigtstp(self, num_completed, num_incomplete):
        """Callback function when SIGTSTP is triggered.

        :param int num_completed: the number of tasks completed in this batch.
        :param int num_incomplete: the number of tasks incomplete in the batch.
        """
        logger.info(
            ('Taskworker process suspended. Completed tasks: %d;'
             ' Incomplete: %d') % (
                num_completed, num_incomplete))

    def _on_shutdown(self, num_completed, num_incomplete):
        """Callback function when we shutdown the worker process.

        :param int num_completed: the number of tasks completed in this batch.
        :param int num_incomplete: the number of tasks incomplete in the batch.
        """
        logger.info(('Taskworker process shutting down. Completed tasks: %d;'
                     ' Incomplete: %d') % (num_completed, num_incomplete))

    def _on_pre_batch_run(self, messages):
        """Callback function before running a batch of tasks.

        :param list[KaleMessage] messages: a list of task messages to process.
        """
        logger.info('Start processing %d tasks in a batch in queue %s ...' % (
            len(messages), self._batch_queue.name))

    def _on_post_batch_run(self, num_completed, num_incomplete, messages):
        """Callback function after running a batch of tasks.

        :param int num_completed: the number of tasks completed in this batch.
        :param int num_incomplete: the number of tasks incomplete in the batch.
        :param list[KaleMessage] messages: a list of messages in this batch.
        """
        logger.info(('Finish processing message batch. Completed tasks: %d;'
                     ' Incomplete: %d') % (
            num_completed, num_incomplete))

    def _on_permanent_failure_batch(self):
        """Callback when there are permanently failed tasks in this batch."""
        logger.info(('Moving permamently %d failed tasks to the '
                     'dead-letter-queue %s.') % (
            len(self._permanent_failures), self._batch_queue.dlq_name))

    def _on_task_deferred(self, message, time_remaining_sec):
        """Callback function when a task is deferred.

        :param list[KaleMessage] message: a list of task messages.
        :param int time_remaining_sec: integer for seconds remained in the
            budget of running this batch of tasks.
        """
        task = message.task_inst
        logger.info(('Task deferred. Task id: %s; Queue: %s; '
                     'Time remaining: %d sec') % (
            task.task_id, self._batch_queue.name, time_remaining_sec))

    def _on_task_failed(self, message, time_remaining_sec, err,
                        permanent_failure):
        """Callback function when a task fails.

        :param KaleMessage message: an object of kale.message.KaleMessage.
        :param int time_remaining_sec: integer for seconds remained in the
            budget of running this batch of tasks.
        :param Exception err: object of Exception.
        :param bool permanent_failure: whether this task permanently fails.
        """
        task = message.task_inst
        logger.info(('Task failed. Task id: %s; Queue: %s; '
                     'Time remaining: %d sec') % (
            task.task_id, self._batch_queue.name, time_remaining_sec))

    def _on_task_succeeded(self, message, time_remaining_sec):
        """Callback function when a task succeeds.

        :param KaleMessage message: an object of kale.message.KaleMessage.
        :param int time_remaining_sec: integer for seconds remained in the
            budget of running this batch of tasks.
        """
        task = message.task_inst
        logger.info(('Task succeeded. Task id: %s; Queue: %s; '
                     'Time remaining: %d sec') % (
            task.task_id, self._batch_queue.name, time_remaining_sec))

    def run(self):
        """This method starts the task processing loop."""

        self._on_pre_run_worker()

        while self._check_process_resources():
            self._batch_queue = self._queue_selector.get_queue()
            for i in range(self._batch_queue.num_iterations):
                if not self._run_single_iteration():
                    # If the iteration didn't process any tasks break out
                    # of this loop and move on to another queue.
                    break

    def _check_process_resources(self):
        """Check if the process is still is abusing resources & should continue

        This will check the processes memory usage and gracefully exit if
        it exceeds the maximum value in settings.

        :return: True if still operational, otherwise it will exit the process.
        :rtype: bool
        """
        resource_data = resource.getrusage(resource.RUSAGE_SELF)

        if resource_data.ru_maxrss < (
                settings.DIE_ON_RESIDENT_SET_SIZE_MB * 1024):

            if self._dirty:
                # We only log when the worker has been infected by  tasks.
                logger.info('Worker process data.')
            return True

        # Allow the client of this library to do any setup before
        # shutting down the worker.
        settings.ON_WORKER_SHUTDOWN()

        self._on_exceeding_memory_limit(resource_data.ru_maxrss)

        # Use non-zero exit code.
        sys.exit(1)

    def _cleanup_worker(self, signum, frame):
        """Handle cleanup when the process is sent a signal.

        This will handle releasing tasks in flight and deleting tasks that have
        been completed.
        """
        logger.info('Process sent signal %d. Cleaning up tasks...' % signum)

        num_completed, num_incomplete = self._release_batch()

        # When the process is suspended we release tasks and then return to the
        # main loop.
        if signum == signal.SIGTSTP:
            self._on_sigtstp(num_completed, num_incomplete)
            return
        else:
            # Allow the client of this library to do any setup before
            # shutting down the worker.
            settings.ON_WORKER_SHUTDOWN()
            self._on_shutdown(num_completed, num_incomplete)
            sys.exit(0)

    def _run_single_iteration(self):
        """Run a single iteration of the task processing loop.

        :return: True if we were able to process a batch, False is there were
            no messages.
        """
        message_batch = self._consumer.fetch_batch(
            self._batch_queue.name,
            self._batch_queue.batch_size,
            self._batch_queue.visibility_timeout_sec,
            long_poll_time_sec=self._batch_queue.long_poll_time_sec)

        self._dirty = bool(message_batch)
        if not message_batch:
            # No any messages in this queue. Let's re-select queue again.
            return False

        self._on_pre_batch_run(message_batch)

        self._batch_stop_time = time.time() + \
            self._batch_queue.visibility_timeout_sec

        self._run_batch(message_batch)
        num_completed, num_incomplete = self._release_batch()

        self._on_post_batch_run(num_completed, num_incomplete, message_batch)

        return True

    def _release_batch(self):
        """Release the most recent batch back to SQS.
        This will delete tasks that succeeded and reset the
        visibility timeout for incomplete tasks.

        :return: A two-tuple of the count of tasks that were completed and
            the count of tasks that were incomplete.
        :rtype: tuple
        """
        # Delete from queues (failed messages are re-published as new tasks)
        messages_to_be_deleted = self._successful_messages + \
            self._failed_messages
        # Set timeout to 0 (if there is time left)
        # As an enhancement, we reset the timeout on tasks that didn't get
        # attempted if the remaining timeout is greater than some threshold.
        # and example of this being helpful is if a 5 minute task was given the
        # opportunity to run with 4 minutes left (and declines). This will
        # release the task 4 minutes before it previously would have.

        if (self._batch_stop_time - time.time()) > \
                settings.RESET_TIMEOUT_THRESHOLD:
            messages_to_be_released = self._incomplete_messages
        else:
            messages_to_be_released = []

        if messages_to_be_deleted:
            # Note: This includes failed tasks.
            self._consumer.delete_messages(messages_to_be_deleted,
                                           self._batch_queue.name)

        if messages_to_be_released:
            # This is only tasks that we didn't get the chance to attempt.
            self._consumer.release_messages(messages_to_be_released,
                                            self._batch_queue.name)

        # Send permanently failed tasks to the dead-letter-queue.
        if self._permanent_failures:
            self._publisher.publish_messages_to_dead_letter_queue(
                self._batch_queue.dlq_name, self._permanent_failures)
            self._on_permanent_failure_batch()

        # All messages start as incomplete.
        self._incomplete_messages = []
        self._successful_messages = []
        self._failed_messages = []
        self._permanent_failures = []

        return len(messages_to_be_deleted), len(messages_to_be_released)

    def _run_batch(self, message_batch):
        """Consume as many tasks as possible in visibility timeout.

        :param list[KaleMessage] message_batch: list of consumable messages.
        """

        # All messages start as incomplete.
        self._incomplete_messages = list(message_batch)
        self._successful_messages = []
        self._failed_messages = []

        # These messages will be sent to the dead-letter-queue.
        self._permanent_failures = []

        # Set visibility timeout start time.
        for message in message_batch:
            task_inst = message.task_inst
            time_remaining_sec = self._batch_stop_time - time.time()

            if task_inst.time_limit >= time_remaining_sec:
                # Greedily continue to process tasks, this task
                # is already in the incomplete_messages list
                self._on_task_deferred(message, time_remaining_sec)
                continue

            # Add cleanup method when tasks are timed out?
            try:
                with timeout.time_limit(task_inst.time_limit):
                    self.run_task(message)
            except Exception as err:
                # Re-publish failed tasks.
                # As an optimization we could run all of the failures from a
                # batch together.
                permanent_failure = not task_inst.__class__.handle_failure(
                    message, err)
                if permanent_failure and settings.USE_DEAD_LETTER_QUEUE:
                    self._permanent_failures.append(message)

                self._failed_messages.append(message)

                self._on_task_failed(message, time_remaining_sec, err,
                                     permanent_failure)
            else:
                self._successful_messages.append(message)
                self._on_task_succeeded(message, time_remaining_sec)
            finally:
                self.remove_message_or_exit(message)

            # Increment total messages counter.
            self._total_messages_processed += 1

    def remove_message_or_exit(self, message):
        try:
            self._incomplete_messages.remove(message)
        except ValueError:
            # Cleanup happened due to the signal handler - make sure we exit immediately.
            sys.exit(0)

    def run_task(self, message):
        """Run the task contained in the message.
        :param message: message.KaleMessage containing the task and arguments to run.
        """
        message.task_inst.run(*message.task_args, **message.task_kwargs)
<commit_msg>Reduce the noisiness of the logs.<commit_after>"""Module for running the worker process.

It's an infinite loop.
"""
from __future__ import absolute_import

import logging
import resource
import signal
import sys
import time

from kale import consumer
from kale import publisher
from kale import queue_info
from kale import settings
from kale import timeout
from kale import utils
from six.moves import range

logger = logging.getLogger(__name__)

SIGNALS_TO_HANDLE = (
    signal.SIGABRT,
    # This will catch Ctrl-C interrupts.
    signal.SIGINT,
    signal.SIGQUIT,
    signal.SIGTERM,
    # Handle Ctrl-Z/suspend.
    signal.SIGTSTP)

# Logging constants.
LOG_TASK_RESULT_DEFERRED = 'deferred'
LOG_TASK_RESULT_ERROR = 'error'
LOG_TASK_RESULT_SUCCESS = 'success'

class Worker(object):

    def __init__(self):
        """Initialize a worker instance."""

        self._consumer = consumer.Consumer()

        queue_class = utils.class_import_from_path(settings.QUEUE_CLASS)
        q_info = queue_info.QueueInfo(
            config_file=settings.QUEUE_CONFIG,
            sqs_talk=self._consumer,
            queue_cls=queue_class)

        queue_selector_class = utils.class_import_from_path(
            settings.QUEUE_SELECTOR)
        self._queue_selector = queue_selector_class(q_info)

        # The worker will publish permanently failed tasks to a
        # dead-letter-queue.
        self._publisher = publisher.Publisher()

        # Track total messages processed.
        self._total_messages_processed = 0

        # Intialize queue variables used by each batch.
        self._incomplete_messages = []
        self._successful_messages = []
        self._failed_messages = []
        self._permanent_failures = []
        self._batch_stop_time = time.time()
        self._batch_queue = None
        # Monitors whether the worker has been exposed to tasks and may
        # have bloated in memory.
        self._dirty = False

        # Setup signal handling for cleanup.
        for sig in SIGNALS_TO_HANDLE:
            signal.signal(sig, self._cleanup_worker)

        # Allow the client of this library to do any setup before
        # starting the worker.
        settings.ON_WORKER_STARTUP()

    def _on_pre_run_worker(self):
        """Callback function right at the beginning of starting the worker. """
        logger.info('Starting run loop for task worker.')

    def _on_exceeding_memory_limit(self, ru_maxrss):
        """Callback function when the process exceeds memory limit.

        :param int ru_maxrss: maximum resident set size used (in kilobytes).
        """
        logger.info('Memory usages exceeds max: %d. Exiting.' % ru_maxrss)

    def _on_sigtstp(self, num_completed, num_incomplete):
        """Callback function when SIGTSTP is triggered.

        :param int num_completed: the number of tasks completed in this batch.
        :param int num_incomplete: the number of tasks incomplete in the batch.
        """
        logger.info(
            ('Taskworker process suspended. Completed tasks: %d;'
             ' Incomplete: %d') % (
                num_completed, num_incomplete))

    def _on_shutdown(self, num_completed, num_incomplete):
        """Callback function when we shutdown the worker process.

        :param int num_completed: the number of tasks completed in this batch.
        :param int num_incomplete: the number of tasks incomplete in the batch.
        """
        logger.info(('Taskworker process shutting down. Completed tasks: %d;'
                     ' Incomplete: %d') % (num_completed, num_incomplete))

    def _on_pre_batch_run(self, messages):
        """Callback function before running a batch of tasks.

        :param list[KaleMessage] messages: a list of task messages to process.
        """
        logger.debug('Start processing %d tasks in a batch in queue %s ...' % (
            len(messages), self._batch_queue.name))

    def _on_post_batch_run(self, num_completed, num_incomplete, messages):
        """Callback function after running a batch of tasks.

        :param int num_completed: the number of tasks completed in this batch.
        :param int num_incomplete: the number of tasks incomplete in the batch.
        :param list[KaleMessage] messages: a list of messages in this batch.
        """
        logger.debug(('Finish processing message batch. Completed tasks: %d;'
                     ' Incomplete: %d') % (
            num_completed, num_incomplete))

    def _on_permanent_failure_batch(self):
        """Callback when there are permanently failed tasks in this batch."""
        logger.info(('Moving permamently %d failed tasks to the '
                     'dead-letter-queue %s.') % (
            len(self._permanent_failures), self._batch_queue.dlq_name))

    def _on_task_deferred(self, message, time_remaining_sec):
        """Callback function when a task is deferred.

        :param list[KaleMessage] message: a list of task messages.
        :param int time_remaining_sec: integer for seconds remained in the
            budget of running this batch of tasks.
        """
        task = message.task_inst
        logger.info(('Task deferred. Task id: %s; Queue: %s; '
                     'Time remaining: %d sec') % (
            task.task_id, self._batch_queue.name, time_remaining_sec))

    def _on_task_failed(self, message, time_remaining_sec, err,
                        permanent_failure):
        """Callback function when a task fails.

        :param KaleMessage message: an object of kale.message.KaleMessage.
        :param int time_remaining_sec: integer for seconds remained in the
            budget of running this batch of tasks.
        :param Exception err: object of Exception.
        :param bool permanent_failure: whether this task permanently fails.
        """
        task = message.task_inst
        logger.info(('Task failed. Task id: %s; Queue: %s; '
                     'Time remaining: %d sec') % (
            task.task_id, self._batch_queue.name, time_remaining_sec))

    def _on_task_succeeded(self, message, time_remaining_sec):
        """Callback function when a task succeeds.

        :param KaleMessage message: an object of kale.message.KaleMessage.
        :param int time_remaining_sec: integer for seconds remained in the
            budget of running this batch of tasks.
        """
        task = message.task_inst
        logger.debug(('Task succeeded. Task id: %s; Queue: %s; '
                     'Time remaining: %d sec') % (
            task.task_id, self._batch_queue.name, time_remaining_sec))

    def run(self):
        """This method starts the task processing loop."""

        self._on_pre_run_worker()

        while self._check_process_resources():
            self._batch_queue = self._queue_selector.get_queue()
            for i in range(self._batch_queue.num_iterations):
                if not self._run_single_iteration():
                    # If the iteration didn't process any tasks break out
                    # of this loop and move on to another queue.
                    break

    def _check_process_resources(self):
        """Check if the process is still is abusing resources & should continue

        This will check the processes memory usage and gracefully exit if
        it exceeds the maximum value in settings.

        :return: True if still operational, otherwise it will exit the process.
        :rtype: bool
        """
        resource_data = resource.getrusage(resource.RUSAGE_SELF)

        if resource_data.ru_maxrss < (
                settings.DIE_ON_RESIDENT_SET_SIZE_MB * 1024):

            if self._dirty:
                # We only log when the worker has been infected by  tasks.
                logger.info('Worker process data.')
            return True

        # Allow the client of this library to do any setup before
        # shutting down the worker.
        settings.ON_WORKER_SHUTDOWN()

        self._on_exceeding_memory_limit(resource_data.ru_maxrss)

        # Use non-zero exit code.
        sys.exit(1)

    def _cleanup_worker(self, signum, frame):
        """Handle cleanup when the process is sent a signal.

        This will handle releasing tasks in flight and deleting tasks that have
        been completed.
        """
        logger.info('Process sent signal %d. Cleaning up tasks...' % signum)

        num_completed, num_incomplete = self._release_batch()

        # When the process is suspended we release tasks and then return to the
        # main loop.
        if signum == signal.SIGTSTP:
            self._on_sigtstp(num_completed, num_incomplete)
            return
        else:
            # Allow the client of this library to do any setup before
            # shutting down the worker.
            settings.ON_WORKER_SHUTDOWN()
            self._on_shutdown(num_completed, num_incomplete)
            sys.exit(0)

    def _run_single_iteration(self):
        """Run a single iteration of the task processing loop.

        :return: True if we were able to process a batch, False is there were
            no messages.
        """
        message_batch = self._consumer.fetch_batch(
            self._batch_queue.name,
            self._batch_queue.batch_size,
            self._batch_queue.visibility_timeout_sec,
            long_poll_time_sec=self._batch_queue.long_poll_time_sec)

        self._dirty = bool(message_batch)
        if not message_batch:
            # No any messages in this queue. Let's re-select queue again.
            return False

        self._on_pre_batch_run(message_batch)

        self._batch_stop_time = time.time() + \
            self._batch_queue.visibility_timeout_sec

        self._run_batch(message_batch)
        num_completed, num_incomplete = self._release_batch()

        self._on_post_batch_run(num_completed, num_incomplete, message_batch)

        return True

    def _release_batch(self):
        """Release the most recent batch back to SQS.
        This will delete tasks that succeeded and reset the
        visibility timeout for incomplete tasks.

        :return: A two-tuple of the count of tasks that were completed and
            the count of tasks that were incomplete.
        :rtype: tuple
        """
        # Delete from queues (failed messages are re-published as new tasks)
        messages_to_be_deleted = self._successful_messages + \
            self._failed_messages
        # Set timeout to 0 (if there is time left)
        # As an enhancement, we reset the timeout on tasks that didn't get
        # attempted if the remaining timeout is greater than some threshold.
        # and example of this being helpful is if a 5 minute task was given the
        # opportunity to run with 4 minutes left (and declines). This will
        # release the task 4 minutes before it previously would have.

        if (self._batch_stop_time - time.time()) > \
                settings.RESET_TIMEOUT_THRESHOLD:
            messages_to_be_released = self._incomplete_messages
        else:
            messages_to_be_released = []

        if messages_to_be_deleted:
            # Note: This includes failed tasks.
            self._consumer.delete_messages(messages_to_be_deleted,
                                           self._batch_queue.name)

        if messages_to_be_released:
            # This is only tasks that we didn't get the chance to attempt.
            self._consumer.release_messages(messages_to_be_released,
                                            self._batch_queue.name)

        # Send permanently failed tasks to the dead-letter-queue.
        if self._permanent_failures:
            self._publisher.publish_messages_to_dead_letter_queue(
                self._batch_queue.dlq_name, self._permanent_failures)
            self._on_permanent_failure_batch()

        # All messages start as incomplete.
        self._incomplete_messages = []
        self._successful_messages = []
        self._failed_messages = []
        self._permanent_failures = []

        return len(messages_to_be_deleted), len(messages_to_be_released)

    def _run_batch(self, message_batch):
        """Consume as many tasks as possible in visibility timeout.

        :param list[KaleMessage] message_batch: list of consumable messages.
        """

        # All messages start as incomplete.
        self._incomplete_messages = list(message_batch)
        self._successful_messages = []
        self._failed_messages = []

        # These messages will be sent to the dead-letter-queue.
        self._permanent_failures = []

        # Set visibility timeout start time.
        for message in message_batch:
            task_inst = message.task_inst
            time_remaining_sec = self._batch_stop_time - time.time()

            if task_inst.time_limit >= time_remaining_sec:
                # Greedily continue to process tasks, this task
                # is already in the incomplete_messages list
                self._on_task_deferred(message, time_remaining_sec)
                continue

            # Add cleanup method when tasks are timed out?
            try:
                with timeout.time_limit(task_inst.time_limit):
                    self.run_task(message)
            except Exception as err:
                # Re-publish failed tasks.
                # As an optimization we could run all of the failures from a
                # batch together.
                permanent_failure = not task_inst.__class__.handle_failure(
                    message, err)
                if permanent_failure and settings.USE_DEAD_LETTER_QUEUE:
                    self._permanent_failures.append(message)

                self._failed_messages.append(message)

                self._on_task_failed(message, time_remaining_sec, err,
                                     permanent_failure)
            else:
                self._successful_messages.append(message)
                self._on_task_succeeded(message, time_remaining_sec)
            finally:
                self.remove_message_or_exit(message)

            # Increment total messages counter.
            self._total_messages_processed += 1

    def remove_message_or_exit(self, message):
        try:
            self._incomplete_messages.remove(message)
        except ValueError:
            # Cleanup happened due to the signal handler - make sure we exit immediately.
            sys.exit(0)

    def run_task(self, message):
        """Run the task contained in the message.
        :param message: message.KaleMessage containing the task and arguments to run.
        """
        message.task_inst.run(*message.task_args, **message.task_kwargs)
<|endoftext|>


<commit_before># Copyright 2011 James McCauley
#
# This file is part of POX.
#
# POX is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POX is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with POX.  If not, see <http://www.gnu.org/licenses/>.

"""
An L2 viro POX controller.
"""

from pox.core import core
import pox.openflow.libopenflow_01 as of
from pox.lib.util import dpid_to_str
from pox.lib.util import str_to_bool
from pox.lib.packet import *
from pox.lib.addresses import *
from pox.lib.util import *
import time
import sys
from viro_module import viroModule
from viro_constant import *
from viro_veil import *
from pox.lib.recoco import Timer

log = core.getLogger()
myViro = ""
Mydpid = 0
myvid = 0
round = 1

class ViroSwitch(object):

    def __init__(self, connection, transparent):

        self.connection = connection
        self.transparent = transparent

        # We want to hear PacketIn messages, so we listen
        connection.addListeners(self)

    def processViroPacket(self, packet, match=None, event=None):

        global myvid, Mydpid, myViro
        L = len(myvid)
        length = getdpidLength(Mydpid)

        opcode = getopcode(packet)

        if opcode == DISC_ECHO_REQ:  # Handles the echo neibghour discover message/packet
            # sends a disc_echo_reply packet

            packet_fields = printDiscoverPacket(packet, L, length)  # gets the fields from the packet
            nvid = packet_fields[1]  # direct neigbour VID

            print "Neighbour discover Request message received from: ", nvid
            r = createDISCOVER_ECHO_REPLY(myvid, Mydpid)
            mac = '00:14:4f:e2:b3:70'

            msg = self.createOPENFLOW_Message(of.OFPP_IN_PORT, mac, r, event.port)

            self.connection.send(msg)
            print "Neighbour discover Reply message sent"

        elif opcode == DISC_ECHO_REPLY:  # Handles the echo neibghour reply message/packet
            # process the packet

            packet_fields = printDiscoverPacket(packet, L, length)  # gets the fields from the packet
            nvid = packet_fields[1]  # direct neigbour VID
            nport = event.port  # direct neigbour port

            print "Neighbour discover Reply message received from: ", nvid
            myViro.updateRoutingTable(nvid, nport)

        else:

            # printPacket(packet, L)
            dst = getDest(packet, L)  # gets the packet dst
            src = getSrc(packet, L)  # gets the packet src

            # forward the packet if I am not the destination
            if dst != myvid:
                self.routeViroPacket(packet)
                return

            if opcode == RDV_QUERY:

                print "RDV_QUERY message received"
                if src == myvid:
                    print "I am the rdv point - processing the packet"
                    myViro.selfRVDQuery(packet)
                    return

                else:
                    rvdReplyPacket = myViro.rvdQuery(packet)

                    if (rvdReplyPacket == ''):
                        return

                    mac = '00:14:4f:e2:b3:70'  # A fake MAC address.
                    msg = self.createOPENFLOW_Message(of.OFPP_IN_PORT, mac, rvdReplyPacket, event.port)

                    self.connection.send(msg)
                    print "RDV_REPLY message sent"

            elif opcode == RDV_PUBLISH:
                myViro.rdvPublish(packet)

            elif opcode == RDV_REPLY:

                print "RDV_REPLY message received"
                myViro.rdvReply(packet)

            elif opcode == VIRO_DATA_OP:
                # The part where it handles VIRO data packet
                print "Received a VIRO Data Packet"

    def createOPENFLOW_Message(self, openflow_port, mac, packet, event_port=None):
        # encapsulating the VIRO packet into an ethernet frame
        e = ethernet(type=0x0802, src=EthAddr(mac))
        e.set_payload(packet)

        # composing openFlow message
        msg = of.ofp_packet_out()
        msg.data = e.pack()
        msg.actions.append(
            of.ofp_action_output(port=openflow_port))  # send the message to the same port as the openflow port

        if (event_port != None):
            msg.in_port = event_port
        return msg

    def _handle_PacketIn(self, event):
        """
        Handle packet in messages from the switch to implement above algorithm.
        """

        packet = event.parsed
        match = of.ofp_match.from_packet(packet)
        # matching the packet type
        try:
            if (match.dl_type == packet.VIRO_TYPE ):

                print  "VIRO packet received....."
                payload = packet.payload
                mypacket = payload
                [packetType] = struct.unpack("!H", mypacket[6:8])

                if (packetType == VIRO_CONTROL):
                    self.processViroPacket(mypacket, match, event)  # handling the VIRO REQUEST
                    return
        except:
            print "Error while processing packet"

    def runARound(self, round):
        global myViro
        routingTable = myViro.routingTable
        global Mydpid, myvid
        mydpid = Mydpid
        L = len(myvid)

        # start from round 2 since connectivity in round 1 is already learnt using the physical neighbors
        for i in range(2, round + 1):

            # see if routing entry for this round is already available in the routing table.
            if i in routingTable:
                if len(routingTable[i]) > 0:

                    # publish the information if it is already there
                    for t in routingTable[i]:
                        if t[1] == int(myvid, 2):
                            print "Sending rdv publish messages"
                            packet, dst = myViro.publish(t, i)
                            self.routeViroPacket(packet)

                else:
                    print "Sending rdv query messages"

                    packet, dst = myViro.query(i)
                    self.routeViroPacket(packet)
            else:

                print "Sending rdv query messages"
                packet, dst = myViro.query(i)
                self.routeViroPacket(packet)

    def startRound(self):
        global myvid, myViro, Mydpid, round
        L = len(myvid)

        print myvid, 'Starting Round : ', round
        self.runARound(round)
        round = round + 1

        if round > L:
            round = L

        print '\n\t----> Routing Table at :', myvid, '|', Mydpid, ' <----'
        for i in range(1, L + 1):
            if i in myViro.routingTable:
                for j in myViro.routingTable[i]:
                    print 'Bucket::', i, 'Nexthop:', bin2str(j[0], L), 'Port:', j[2], 'Gateway:', bin2str(j[1],
                                                                                L), 'Prefix:', \
                    j[3]
            else:
                print 'Bucket', i, '  --- E M P T Y --- '
        print 'RDV STORE: ', myViro.rdvStore
        print '\n --  --  --  --  -- --  --  --  --  -- --  --  --  --  -- \n'

    def routeViroPacket(self, packet):
        global myvid, myViro

        # Type of packet: rvds Query or Publish
        # k - bucket level

        L = len(myvid)
        dst = getDest(packet, L)

        # If it's me
        if (dst == myvid):
            print 'I am the destination!'
            self.processViroPacket(packet)
            return

            # get nextHop and port
        nextHop, port = myViro.getNextHop(packet)
        if ( nextHop != ''):

            hwrdst = '00:14:4f:e2:b3:70'

            msg = self.createOPENFLOW_Message(of.OFPP_IN_PORT, hwrdst, packet, int(port))

            self.connection.send(msg)
        else:
            print " Next hop is none "

class viro_controller(object):
    """
    Waits for OpenFlow switches to connect.
    """

    def __init__(self, transparent):
        core.openflow.addListeners(self)
        self.transparent = transparent
        self.myviroSwitch = ''

    def _handle_ConnectionUp(self, event):
        log.debug("Connection %s" % (event.connection,))
        self.myviroSwitch = ViroSwitch(event.connection, self.transparent)

        print "Starting Neighbor Discovery ...."
        global Mydpid, myViro, myvid

        Mydpid = dpidToStr(event.connection.dpid)  # gets the switch dpid identifier
        if (Mydpid == '00-00-00-00-00-01'):
            myvid = '000'
        elif (Mydpid == '00-00-00-00-00-02'):
            myvid = '001'
        elif (Mydpid == '00-00-00-00-00-03'):
            myvid = '010'
        elif (Mydpid == '00-00-00-00-00-04'):
            myvid = '011'
        elif (Mydpid == '00-00-00-00-00-05'):
            myvid = '100'
        elif (Mydpid == '00-00-00-00-00-06'):
            myvid = '101'
        elif (Mydpid == '00-00-00-00-00-07'):
            myvid = '110'
        elif (Mydpid == '00-00-00-00-00-08'):
            myvid = '111'

        myViro = viroModule(Mydpid, myvid)

        # Call neighbour discovery function after every DISCOVER_TIME seconds
        Timer(DISCOVER_TIME, self.neibghoorDiscover, args=[Mydpid, myvid, event], recurring=True)
        # Poulate routing table after every UPDATE_RT_TIME seconds
        Timer(UPDATE_RT_TIME, self.myviroSwitch.startRound, recurring=True)
        # Look for failures in the neigbours switches
        Timer(FAILURE_TIME, self.discoveryFailure, recurring=True)

    def neibghoorDiscover(self, mydip, myvid, event):

        try:
            dpid = mydip
            r = createDISCOVER_ECHO_REQ(myvid, dpid)
            mac = '00:14:4f:e2:b3:70'  # A fake MAC address.
            msg = self.myviroSwitch.createOPENFLOW_Message(of.OFPP_FLOOD, mac, r, None)
            event.connection.send(msg)
            print "Sending neighbour discover packets"

        except:
            print "Error .... not able to send discover packets"

    def discoveryFailure(self):
        # This is the function handling failure events
        # Random code, delete this when you are doing it.
        csci = 5221

def launch(transparent=False):
    """
    Starts an VIRO switch.
    """

    core.registerNew(viro_controller, str_to_bool(transparent))
<commit_msg>Refactored get VID from PID to be less hard-coded<commit_after># Copyright 2011 James McCauley
#
# This file is part of POX.
#
# POX is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# POX is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with POX.  If not, see <http://www.gnu.org/licenses/>.

"""
An L2 viro POX controller.
"""

from pox.core import core
import pox.openflow.libopenflow_01 as of
from pox.lib.util import dpid_to_str
from pox.lib.util import str_to_bool
from pox.lib.packet import *
from pox.lib.addresses import *
from pox.lib.util import *
import time
import sys
from viro_module import viroModule
from viro_constant import *
from viro_veil import *
from pox.lib.recoco import Timer

log = core.getLogger()
myViro = ""
mydpid = 0
myvid = 0
round = 1

class ViroSwitch(object):

    def __init__(self, connection, transparent):

        self.connection = connection
        self.transparent = transparent

        # We want to hear PacketIn messages, so we listen
        connection.addListeners(self)

    def processViroPacket(self, packet, match=None, event=None):

        global myvid, mydpid, myViro
        L = len(myvid)
        length = getdpidLength(mydpid)

        opcode = getopcode(packet)

        if opcode == DISC_ECHO_REQ:  # Handles the echo neibghour discover message/packet
            # sends a disc_echo_reply packet

            packet_fields = printDiscoverPacket(packet, L, length)  # gets the fields from the packet
            nvid = packet_fields[1]  # direct neigbour VID

            print "Neighbour discover Request message received from: ", nvid
            r = createDISCOVER_ECHO_REPLY(myvid, mydpid)
            mac = '00:14:4f:e2:b3:70'

            msg = self.createOPENFLOW_Message(of.OFPP_IN_PORT, mac, r, event.port)

            self.connection.send(msg)
            print "Neighbour discover Reply message sent"

        elif opcode == DISC_ECHO_REPLY:  # Handles the echo neibghour reply message/packet
            # process the packet

            packet_fields = printDiscoverPacket(packet, L, length)  # gets the fields from the packet
            nvid = packet_fields[1]  # direct neigbour VID
            nport = event.port  # direct neigbour port

            print "Neighbour discover Reply message received from: ", nvid
            myViro.updateRoutingTable(nvid, nport)

        else:

            # printPacket(packet, L)
            dst = getDest(packet, L)  # gets the packet dst
            src = getSrc(packet, L)  # gets the packet src

            # forward the packet if I am not the destination
            if dst != myvid:
                self.routeViroPacket(packet)
                return

            if opcode == RDV_QUERY:

                print "RDV_QUERY message received"
                if src == myvid:
                    print "I am the rdv point - processing the packet"
                    myViro.selfRVDQuery(packet)
                    return

                else:
                    rvdReplyPacket = myViro.rvdQuery(packet)

                    if (rvdReplyPacket == ''):
                        return

                    mac = '00:14:4f:e2:b3:70'  # A fake MAC address.
                    msg = self.createOPENFLOW_Message(of.OFPP_IN_PORT, mac, rvdReplyPacket, event.port)

                    self.connection.send(msg)
                    print "RDV_REPLY message sent"

            elif opcode == RDV_PUBLISH:
                myViro.rdvPublish(packet)

            elif opcode == RDV_REPLY:

                print "RDV_REPLY message received"
                myViro.rdvReply(packet)

            elif opcode == VIRO_DATA_OP:
                # The part where it handles VIRO data packet
                print "Received a VIRO Data Packet"

    def createOPENFLOW_Message(self, openflow_port, mac, packet, event_port=None):
        # encapsulating the VIRO packet into an ethernet frame
        e = ethernet(type=0x0802, src=EthAddr(mac))
        e.set_payload(packet)

        # composing openFlow message
        msg = of.ofp_packet_out()
        msg.data = e.pack()
        msg.actions.append(
            of.ofp_action_output(port=openflow_port))  # send the message to the same port as the openflow port

        if (event_port != None):
            msg.in_port = event_port
        return msg

    def _handle_PacketIn(self, event):
        """
        Handle packet in messages from the switch to implement above algorithm.
        """

        packet = event.parsed
        match = of.ofp_match.from_packet(packet)
        # matching the packet type
        try:
            if (match.dl_type == packet.VIRO_TYPE ):

                print  "VIRO packet received....."
                payload = packet.payload
                mypacket = payload
                [packetType] = struct.unpack("!H", mypacket[6:8])

                if (packetType == VIRO_CONTROL):
                    self.processViroPacket(mypacket, match, event)  # handling the VIRO REQUEST
                    return
        except:
            print "Error while processing packet"

    def runARound(self, round):
        global myViro
        routingTable = myViro.routingTable
        global mydpid, myvid
        mydpid = mydpid
        L = len(myvid)

        # start from round 2 since connectivity in round 1 is already learnt using the physical neighbors
        for i in range(2, round + 1):

            # see if routing entry for this round is already available in the routing table.
            if i in routingTable:
                if len(routingTable[i]) > 0:

                    # publish the information if it is already there
                    for t in routingTable[i]:
                        if t[1] == int(myvid, 2):
                            print "Sending rdv publish messages"
                            packet, dst = myViro.publish(t, i)
                            self.routeViroPacket(packet)

                else:
                    print "Sending rdv query messages"

                    packet, dst = myViro.query(i)
                    self.routeViroPacket(packet)
            else:

                print "Sending rdv query messages"
                packet, dst = myViro.query(i)
                self.routeViroPacket(packet)

    def startRound(self):
        global myvid, myViro, mydpid, round
        L = len(myvid)

        print myvid, 'Starting Round : ', round
        self.runARound(round)
        round = round + 1

        if round > L:
            round = L

        print '\n\t----> Routing Table at :', myvid, '|', mydpid, ' <----'
        for i in range(1, L + 1):
            if i in myViro.routingTable:
                for j in myViro.routingTable[i]:
                    print 'Bucket::', i, 'Nexthop:', bin2str(j[0], L), 'Port:', j[2], 'Gateway:', bin2str(j[1],
                                                                                L), 'Prefix:', \
                    j[3]
            else:
                print 'Bucket', i, '  --- E M P T Y --- '
        print 'RDV STORE: ', myViro.rdvStore
        print '\n --  --  --  --  -- --  --  --  --  -- --  --  --  --  -- \n'

    def routeViroPacket(self, packet):
        global myvid, myViro

        # Type of packet: rvds Query or Publish
        # k - bucket level

        L = len(myvid)
        dst = getDest(packet, L)

        # If it's me
        if (dst == myvid):
            print 'I am the destination!'
            self.processViroPacket(packet)
            return

            # get nextHop and port
        nextHop, port = myViro.getNextHop(packet)
        if ( nextHop != ''):

            hwrdst = '00:14:4f:e2:b3:70'

            msg = self.createOPENFLOW_Message(of.OFPP_IN_PORT, hwrdst, packet, int(port))

            self.connection.send(msg)
        else:
            print " Next hop is none "

class viro_controller(object):
    """
    Waits for OpenFlow switches to connect.
    """

    def __init__(self, transparent):
        core.openflow.addListeners(self)
        self.transparent = transparent
        self.myviroSwitch = ''

    def _handle_ConnectionUp(self, event):
        log.debug("Connection %s" % (event.connection,))
        self.myviroSwitch = ViroSwitch(event.connection, self.transparent)

        print "Starting Neighbor Discovery ...."
        global mydpid, myViro, myvid

        mydpid = dpidToStr(event.connection.dpid)   # gets the switch dpid identifier
        myvid = self.get_vid_from_pid(mydpid)
        myViro = viroModule(mydpid, myvid)

        # Call neighbour discovery function after every DISCOVER_TIME seconds
        Timer(DISCOVER_TIME, self.neibghoorDiscover, args=[mydpid, myvid, event], recurring=True)
        # Poulate routing table after every UPDATE_RT_TIME seconds
        Timer(UPDATE_RT_TIME, self.myviroSwitch.startRound, recurring=True)
        # Look for failures in the neigbours switches
        Timer(FAILURE_TIME, self.discoveryFailure, recurring=True)

    def get_vid_from_pid(self, pid):
        # To convert a pid string (assumed to be formatted like a MAC address: xx-xx-xx-xx-xx-xx)
        # starting at "00-00-00-00-00-01" to a vid string (of '1' and '0' characters)
        # starting at "000" we do the following:
        #   1. Remove "-" characters to make the string numeric
        #   2. Convert that string to an integer (assuming base 16)
        #   3. Subtract 1 so that "00-00-00-00-00-01" corresponds with "000"
        #   4. Convert the int back into a string using base 2
        #   5. Zero-pad the result the 3 bits to match the behavior of the original function
        return format(int(pid.replace('-', ''), 16) - 1, 'b').zfill(3)

    def neibghoorDiscover(self, mydip, myvid, event):

        try:
            dpid = mydip
            r = createDISCOVER_ECHO_REQ(myvid, dpid)
            mac = '00:14:4f:e2:b3:70'  # A fake MAC address.
            msg = self.myviroSwitch.createOPENFLOW_Message(of.OFPP_FLOOD, mac, r, None)
            event.connection.send(msg)
            print "Sending neighbour discover packets"

        except:
            print "Error .... not able to send discover packets"

    def discoveryFailure(self):
        # This is the function handling failure events
        # Random code, delete this when you are doing it.
        csci = 5221

def launch(transparent=False):
    """
    Starts an VIRO switch.
    """

    core.registerNew(viro_controller, str_to_bool(transparent))
<|endoftext|>


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

__all__ = ['bilibili_download']

from ..common import *

from .sina import sina_download_by_id
from .tudou import tudou_download_by_id
from .youku import youku_download_by_id

import re

def get_srt_xml(id):
    url = 'http://comment.bilibili.tv/%s.xml' % id
    return get_html(url)

def parse_srt_p(p):
    fields = p.split(',')
    assert len(fields) == 8, fields
    time, mode, font_size, font_color, pub_time, pool, user_id, history = fields
    time = float(time)

    mode = int(mode)
    assert 1 <= mode <= 8
    # mode 1~3: scrolling
    # mode 4: bottom
    # mode 5: top
    # mode 6: reverse?
    # mode 7: position
    # mode 8: advanced

    pool = int(pool)
    assert 0 <= pool <= 2
    # pool 0: normal
    # pool 1: srt
    # pool 2: special?

    font_size = int(font_size)

    font_color = '#%06x' % int(font_color)

    return pool, mode, font_size, font_color

def parse_srt_xml(xml):
    d = re.findall(r'<d p="([^"]+)">(.*)</d>', xml)
    for x, y in d:
        p = parse_srt_p(x)
    raise NotImplementedError()

def parse_cid_playurl(xml):
    from xml.dom.minidom import parseString
    doc = parseString(xml.encode('utf-8'))
    urls = [durl.getElementsByTagName('url')[0].firstChild.nodeValue for durl in doc.getElementsByTagName('durl')]
    return urls

def bilibili_download_by_cid(id, title, output_dir = '.', merge = True, info_only = False):
    url = 'http://interface.bilibili.tv/playurl?cid=' + id
    urls = parse_cid_playurl(get_html(url, 'utf-8'))

    if re.search(r'\.(flv|hlv)\b', urls[0]):
        type = 'flv'
    elif re.search(r'/mp4/', urls[0]):
        type = 'mp4'
    else:
        raise NotImplementedError(urls[0])

    size = 0
    for url in urls:
        _, _, temp = url_info(url)
        size += temp

    print_info(site_info, title, type, size)
    if not info_only:
        download_urls(urls, title, type, total_size = None, output_dir = output_dir, merge = merge)

def bilibili_download(url, output_dir = '.', merge = True, info_only = False):
    assert re.match(r'http://(www.bilibili.tv|bilibili.kankanews.com|bilibili.smgbb.cn)/video/av(\d+)', url)
    html = get_html(url)

    title = r1(r'
([^<>]+)
', html)
    title = unescape_html(title)
    title = escape_file_path(title)

    flashvars = r1_of([r'flashvars="([^"]+)"', r'"https://secure.bilibili.tv/secure,(cid=\d+)(?:&aid=\d+)?"'], html)
    assert flashvars
    t, id = flashvars.split('=', 1)
    id = id.split('&')[0]
    if t == 'cid':
        bilibili_download_by_cid(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    elif t == 'vid':
        sina_download_by_id(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    elif t == 'ykid':
        youku_download_by_id(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    elif t == 'uid':
        tudou_download_by_id(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    else:
        raise NotImplementedError(flashvars)

    if not info_only:
        print('Downloading %s ...' % (title + '.cmt.xml'))
        xml = get_srt_xml(id)
        with open(title + '.cmt.xml', 'w') as x:
            x.write(xml)

site_info = "bilibili.tv"
download = bilibili_download
download_playlist = playlist_not_supported('bilibili')
<commit_msg>fix #123 for bilibili<commit_after>#!/usr/bin/env python

__all__ = ['bilibili_download']

from ..common import *

from .sina import sina_download_by_id
from .tudou import tudou_download_by_id
from .youku import youku_download_by_id

import re

def get_srt_xml(id):
    url = 'http://comment.bilibili.tv/%s.xml' % id
    return get_html(url)

def parse_srt_p(p):
    fields = p.split(',')
    assert len(fields) == 8, fields
    time, mode, font_size, font_color, pub_time, pool, user_id, history = fields
    time = float(time)

    mode = int(mode)
    assert 1 <= mode <= 8
    # mode 1~3: scrolling
    # mode 4: bottom
    # mode 5: top
    # mode 6: reverse?
    # mode 7: position
    # mode 8: advanced

    pool = int(pool)
    assert 0 <= pool <= 2
    # pool 0: normal
    # pool 1: srt
    # pool 2: special?

    font_size = int(font_size)

    font_color = '#%06x' % int(font_color)

    return pool, mode, font_size, font_color

def parse_srt_xml(xml):
    d = re.findall(r'<d p="([^"]+)">(.*)</d>', xml)
    for x, y in d:
        p = parse_srt_p(x)
    raise NotImplementedError()

def parse_cid_playurl(xml):
    from xml.dom.minidom import parseString
    doc = parseString(xml.encode('utf-8'))
    urls = [durl.getElementsByTagName('url')[0].firstChild.nodeValue for durl in doc.getElementsByTagName('durl')]
    return urls

def bilibili_download_by_cid(id, title, output_dir = '.', merge = True, info_only = False):
    url = 'http://interface.bilibili.tv/playurl?cid=' + id
    urls = parse_cid_playurl(get_html(url, 'utf-8'))

    if re.search(r'\.(flv|hlv)\b', urls[0]):
        type = 'flv'
    elif re.search(r'/flv/', urls[0]):
        type = 'flv'
    elif re.search(r'/mp4/', urls[0]):
        type = 'mp4'
    else:
        raise NotImplementedError(urls[0])

    size = 0
    for url in urls:
        _, _, temp = url_info(url)
        size += temp

    print_info(site_info, title, type, size)
    if not info_only:
        download_urls(urls, title, type, total_size = None, output_dir = output_dir, merge = merge)

def bilibili_download(url, output_dir = '.', merge = True, info_only = False):
    assert re.match(r'http://(www.bilibili.tv|bilibili.kankanews.com|bilibili.smgbb.cn)/video/av(\d+)', url)
    html = get_html(url)

    title = r1(r'
([^<>]+)
', html)
    title = unescape_html(title)
    title = escape_file_path(title)

    flashvars = r1_of([r'flashvars="([^"]+)"', r'"https://secure.bilibili.tv/secure,(cid=\d+)(?:&aid=\d+)?"'], html)
    assert flashvars
    t, id = flashvars.split('=', 1)
    id = id.split('&')[0]
    if t == 'cid':
        bilibili_download_by_cid(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    elif t == 'vid':
        sina_download_by_id(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    elif t == 'ykid':
        youku_download_by_id(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    elif t == 'uid':
        tudou_download_by_id(id, title, output_dir = output_dir, merge = merge, info_only = info_only)
    else:
        raise NotImplementedError(flashvars)

    if not info_only:
        print('Downloading %s ...' % (title + '.cmt.xml'))
        xml = get_srt_xml(id)
        with open(title + '.cmt.xml', 'w') as x:
            x.write(xml)

site_info = "bilibili.tv"
download = bilibili_download
download_playlist = playlist_not_supported('bilibili')
<|endoftext|>


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

# Copyright 2015 SmartBear Software
#
# 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.

"""Swagger generic API client. This client handles the client-
server communication, and is invariant across implementations. Specifics of
the methods and models for each application are generated from the Swagger
templates."""

import __builtin__

import sys
import os
import re
import requests
import urllib
import urllib2
import httplib
import json
import datetime
import mimetypes
import random
import string

from magnum.common import utils

from oslo_utils import importutils

class ApiClient(object):
  """Generic API client for Swagger client library builds

  Attributes:
    host: The base path for the server to call
    headerName: a header to pass when making calls to the API
    headerValue: a header value to pass when making calls to the API
  """
  def __init__(self, host=None, headerName=None, headerValue=None):
    self.defaultHeaders = {}
    if (headerName is not None):
      self.defaultHeaders[headerName] = headerValue
    self.host = host
    self.cookie = None
    self.boundary = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(30))
    # Set default User-Agent.
    self.user_agent = 'Python-Swagger'

  @property
  def user_agent(self):
    return self.defaultHeaders['User-Agent']

  @user_agent.setter
  def user_agent(self, value):
    self.defaultHeaders['User-Agent'] = value

  def setDefaultHeader(self, headerName, headerValue):
    self.defaultHeaders[headerName] = headerValue

  def callAPI(self, resourcePath, method, queryParams, postData,
              ca_cert=None, cert=None, key=None, headerParams=None, files=None):

    url = self.host + resourcePath

    mergedHeaderParams = self.defaultHeaders.copy()
    if headerParams:
        mergedHeaderParams.update(headerParams)
    headers = {}
    if mergedHeaderParams:
      for param, value in mergedHeaderParams.iteritems():
        headers[param] = ApiClient.sanitizeForSerialization(value)

    if self.cookie:
      headers['Cookie'] = ApiClient.sanitizeForSerialization(self.cookie)

    data = None

    if queryParams:
      # Need to remove None values, these should not be sent
      sentQueryParams = {}
      for param, value in queryParams.items():
        if value is not None:
          sentQueryParams[param] = ApiClient.sanitizeForSerialization(value)
      url = url + '?' + urllib.urlencode(sentQueryParams)

    if method in ['GET']:
      #Options to add statements later on and for compatibility
      pass

    elif method in ['POST', 'PUT', 'DELETE']:
      if postData:
        postData = ApiClient.sanitizeForSerialization(postData)
        if 'Content-type' not in headers:
          headers['Content-type'] = 'application/json'
          data = json.dumps(postData)
        elif headers['Content-type'] == 'multipart/form-data':
          data = self.buildMultipartFormData(postData, files)
          headers['Content-type'] = 'multipart/form-data; boundary={0}'.format(self.boundary)
          headers['Content-length'] = str(len(data))
        else:
            data = urllib.urlencode(postData)

    else:
      raise Exception('Method ' + method + ' is not recognized.')

    utils.raise_exception_invalid_scheme(url)

    response = requests.request(method, url=url, headers=headers, data=data,
                                cert=(cert, key), verify=ca_cert)
    if 'Set-Cookie' in response.headers:
      self.cookie = response.headers['Set-Cookie']
    try:
      data = json.loads(response.content)
    except ValueError:  # PUT requests don't return anything
      data = None

    return data

  def toPathValue(self, obj):
    """Convert a string or object to a path-friendly value
    Args:
        obj -- object or string value
    Returns:
        string -- quoted value
    """
    if type(obj) == list:
      return ','.join(obj)
    else:
      return str(obj)

  @staticmethod
  def sanitizeForSerialization(obj):
    """
    Sanitize an object for Request.

    If obj is None, return None.
    If obj is str, int, long, float, bool, return directly.
    If obj is datetime.datetime, datetime.date convert to string in iso8601 format.
    If obj is list, santize each element in the list.
    If obj is dict, return the dict.
    If obj is swagger model, return the properties dict.
    """
    if isinstance(obj, type(None)):
      return None
    elif isinstance(obj, (unicode, str, int, long, float, bool, file)):
      return obj
    elif isinstance(obj, list):
      return [ApiClient.sanitizeForSerialization(subObj) for subObj in obj]
    elif isinstance(obj, (datetime.datetime, datetime.date)):
      return obj.isoformat()
    else:
      if isinstance(obj, dict):
        objDict = obj
      else:
        # Convert model obj to dict except attributes `swaggerTypes`, `attributeMap`
        # and attributes which value is not None.
        # Convert attribute name to json key in model definition for request.
        objDict = {obj.attributeMap[key]: val
                   for key, val in obj.__dict__.iteritems()
                   if key != 'swaggerTypes' and key != 'attributeMap' and val is not None}
      return {key: ApiClient.sanitizeForSerialization(val)
              for (key, val) in objDict.iteritems()}

  def buildMultipartFormData(self, postData, files):
    def escape_quotes(s):
      return s.replace('"', '\\"')

    lines = []

    for name, value in postData.items():
      lines.extend((
        '--{0}'.format(self.boundary),
        'Content-Disposition: form-data; name="{0}"'.format(escape_quotes(name)),
        '',
        str(value),
      ))

    for name, filepath in files.items():
      f = open(filepath, 'r')
      filename = filepath.split('/')[-1]
      mimetype = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
      lines.extend((
        '--{0}'.format(self.boundary),
        'Content-Disposition: form-data; name="{0}"; filename="{1}"'.format(escape_quotes(name), escape_quotes(filename)),
        'Content-Type: {0}'.format(mimetype),
        '',
        f.read()
      ))

    lines.extend((
      '--{0}--'.format(self.boundary),
      ''
    ))
    return '\r\n'.join(lines)

  def deserialize(self, obj, objClass):
    """Derialize a JSON string into an object.

    Args:
        obj -- string or object to be deserialized
        objClass -- class literal for deserialzied object, or string
            of class name
    Returns:
        object -- deserialized object"""

    # Have to accept objClass as string or actual type. Type could be a
    # native Python type, or one of the model classes.
    if type(objClass) == str:
      if 'list[' in objClass:
        match = re.match('list\[(.*)\]', objClass)
        subClass = match.group(1)
        return [self.deserialize(subObj, subClass) for subObj in obj]

      classname = objClass
      if classname in {'int', 'float', 'long', 'dict', 'list', 'str', 'bool'}:
        objClass = getattr(__builtin__, classname)
      elif classname == 'datetime':
        objClass = self.__parse_string_to_datetime
      else:  # not a native type, must be model class
        model = ('magnum.common.pythonk8sclient.client.models.%s.%s' %
                 (classname, classname))
        objClass = importutils.import_class(model)
    else:
      classname = None

    if ((getattr(objClass, '__module__', None) == __builtin__.__name__) or
        (classname == 'datetime')):
      return objClass(obj)

    instance = objClass()

    for attr, attrType in instance.swaggerTypes.iteritems():
        if obj is not None and instance.attributeMap[attr] in obj and type(obj) in [list, dict]:
          value = obj[instance.attributeMap[attr]]
          if attrType in {'str', 'int', 'long', 'float', 'bool'}:
            attrType = getattr(__builtin__, attrType)
            try:
              value = attrType(value)
            except UnicodeEncodeError:
              value = unicode(value)
            except TypeError:
              value = value
            setattr(instance, attr, value)
          elif (attrType == 'datetime'):
            setattr(instance, attr, self.__parse_string_to_datetime(value))
          elif 'list[' in attrType:
            match = re.match('list\[(.*)\]', attrType)
            subClass = match.group(1)
            subValues = []
            if not value:
              setattr(instance, attr, None)
            else:
              for subValue in value:
                subValues.append(self.deserialize(subValue, subClass))
            setattr(instance, attr, subValues)
          else:
            setattr(instance, attr, self.deserialize(value, attrType))
    return instance

  def __parse_string_to_datetime(self, string):
    """
    Parse datetime in string to datetime.

    The string should be in iso8601 datetime format.
    """
    try:
        from dateutil.parser import parse
        return parse(string)
    except ImportError:
        return string
<commit_msg>Fix typo error<commit_after>#!/usr/bin/env python

# Copyright 2015 SmartBear Software
#
# 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.

"""Swagger generic API client. This client handles the client-
server communication, and is invariant across implementations. Specifics of
the methods and models for each application are generated from the Swagger
templates."""

import __builtin__

import sys
import os
import re
import requests
import urllib
import urllib2
import httplib
import json
import datetime
import mimetypes
import random
import string

from magnum.common import utils

from oslo_utils import importutils

class ApiClient(object):
  """Generic API client for Swagger client library builds

  Attributes:
    host: The base path for the server to call
    headerName: a header to pass when making calls to the API
    headerValue: a header value to pass when making calls to the API
  """
  def __init__(self, host=None, headerName=None, headerValue=None):
    self.defaultHeaders = {}
    if (headerName is not None):
      self.defaultHeaders[headerName] = headerValue
    self.host = host
    self.cookie = None
    self.boundary = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(30))
    # Set default User-Agent.
    self.user_agent = 'Python-Swagger'

  @property
  def user_agent(self):
    return self.defaultHeaders['User-Agent']

  @user_agent.setter
  def user_agent(self, value):
    self.defaultHeaders['User-Agent'] = value

  def setDefaultHeader(self, headerName, headerValue):
    self.defaultHeaders[headerName] = headerValue

  def callAPI(self, resourcePath, method, queryParams, postData,
              ca_cert=None, cert=None, key=None, headerParams=None, files=None):

    url = self.host + resourcePath

    mergedHeaderParams = self.defaultHeaders.copy()
    if headerParams:
        mergedHeaderParams.update(headerParams)
    headers = {}
    if mergedHeaderParams:
      for param, value in mergedHeaderParams.iteritems():
        headers[param] = ApiClient.sanitizeForSerialization(value)

    if self.cookie:
      headers['Cookie'] = ApiClient.sanitizeForSerialization(self.cookie)

    data = None

    if queryParams:
      # Need to remove None values, these should not be sent
      sentQueryParams = {}
      for param, value in queryParams.items():
        if value is not None:
          sentQueryParams[param] = ApiClient.sanitizeForSerialization(value)
      url = url + '?' + urllib.urlencode(sentQueryParams)

    if method in ['GET']:
      #Options to add statements later on and for compatibility
      pass

    elif method in ['POST', 'PUT', 'DELETE']:
      if postData:
        postData = ApiClient.sanitizeForSerialization(postData)
        if 'Content-type' not in headers:
          headers['Content-type'] = 'application/json'
          data = json.dumps(postData)
        elif headers['Content-type'] == 'multipart/form-data':
          data = self.buildMultipartFormData(postData, files)
          headers['Content-type'] = 'multipart/form-data; boundary={0}'.format(self.boundary)
          headers['Content-length'] = str(len(data))
        else:
            data = urllib.urlencode(postData)

    else:
      raise Exception('Method ' + method + ' is not recognized.')

    utils.raise_exception_invalid_scheme(url)

    response = requests.request(method, url=url, headers=headers, data=data,
                                cert=(cert, key), verify=ca_cert)
    if 'Set-Cookie' in response.headers:
      self.cookie = response.headers['Set-Cookie']
    try:
      data = json.loads(response.content)
    except ValueError:  # PUT requests don't return anything
      data = None

    return data

  def toPathValue(self, obj):
    """Convert a string or object to a path-friendly value
    Args:
        obj -- object or string value
    Returns:
        string -- quoted value
    """
    if type(obj) == list:
      return ','.join(obj)
    else:
      return str(obj)

  @staticmethod
  def sanitizeForSerialization(obj):
    """
    Sanitize an object for Request.

    If obj is None, return None.
    If obj is str, int, long, float, bool, return directly.
    If obj is datetime.datetime, datetime.date convert to string in iso8601 format.
    If obj is list, santize each element in the list.
    If obj is dict, return the dict.
    If obj is swagger model, return the properties dict.
    """
    if isinstance(obj, type(None)):
      return None
    elif isinstance(obj, (unicode, str, int, long, float, bool, file)):
      return obj
    elif isinstance(obj, list):
      return [ApiClient.sanitizeForSerialization(subObj) for subObj in obj]
    elif isinstance(obj, (datetime.datetime, datetime.date)):
      return obj.isoformat()
    else:
      if isinstance(obj, dict):
        objDict = obj
      else:
        # Convert model obj to dict except attributes `swaggerTypes`, `attributeMap`
        # and attributes which value is not None.
        # Convert attribute name to json key in model definition for request.
        objDict = {obj.attributeMap[key]: val
                   for key, val in obj.__dict__.iteritems()
                   if key != 'swaggerTypes' and key != 'attributeMap' and val is not None}
      return {key: ApiClient.sanitizeForSerialization(val)
              for (key, val) in objDict.iteritems()}

  def buildMultipartFormData(self, postData, files):
    def escape_quotes(s):
      return s.replace('"', '\\"')

    lines = []

    for name, value in postData.items():
      lines.extend((
        '--{0}'.format(self.boundary),
        'Content-Disposition: form-data; name="{0}"'.format(escape_quotes(name)),
        '',
        str(value),
      ))

    for name, filepath in files.items():
      f = open(filepath, 'r')
      filename = filepath.split('/')[-1]
      mimetype = mimetypes.guess_type(filename)[0] or 'application/octet-stream'
      lines.extend((
        '--{0}'.format(self.boundary),
        'Content-Disposition: form-data; name="{0}"; filename="{1}"'.format(escape_quotes(name), escape_quotes(filename)),
        'Content-Type: {0}'.format(mimetype),
        '',
        f.read()
      ))

    lines.extend((
      '--{0}--'.format(self.boundary),
      ''
    ))
    return '\r\n'.join(lines)

  def deserialize(self, obj, objClass):
    """Deserialize a JSON string into an object.

    Args:
        obj -- string or object to be deserialized
        objClass -- class literal for deserialzied object, or string
            of class name
    Returns:
        object -- deserialized object"""

    # Have to accept objClass as string or actual type. Type could be a
    # native Python type, or one of the model classes.
    if type(objClass) == str:
      if 'list[' in objClass:
        match = re.match('list\[(.*)\]', objClass)
        subClass = match.group(1)
        return [self.deserialize(subObj, subClass) for subObj in obj]

      classname = objClass
      if classname in {'int', 'float', 'long', 'dict', 'list', 'str', 'bool'}:
        objClass = getattr(__builtin__, classname)
      elif classname == 'datetime':
        objClass = self.__parse_string_to_datetime
      else:  # not a native type, must be model class
        model = ('magnum.common.pythonk8sclient.client.models.%s.%s' %
                 (classname, classname))
        objClass = importutils.import_class(model)
    else:
      classname = None

    if ((getattr(objClass, '__module__', None) == __builtin__.__name__) or
        (classname == 'datetime')):
      return objClass(obj)

    instance = objClass()

    for attr, attrType in instance.swaggerTypes.iteritems():
        if obj is not None and instance.attributeMap[attr] in obj and type(obj) in [list, dict]:
          value = obj[instance.attributeMap[attr]]
          if attrType in {'str', 'int', 'long', 'float', 'bool'}:
            attrType = getattr(__builtin__, attrType)
            try:
              value = attrType(value)
            except UnicodeEncodeError:
              value = unicode(value)
            except TypeError:
              value = value
            setattr(instance, attr, value)
          elif (attrType == 'datetime'):
            setattr(instance, attr, self.__parse_string_to_datetime(value))
          elif 'list[' in attrType:
            match = re.match('list\[(.*)\]', attrType)
            subClass = match.group(1)
            subValues = []
            if not value:
              setattr(instance, attr, None)
            else:
              for subValue in value:
                subValues.append(self.deserialize(subValue, subClass))
            setattr(instance, attr, subValues)
          else:
            setattr(instance, attr, self.deserialize(value, attrType))
    return instance

  def __parse_string_to_datetime(self, string):
    """
    Parse datetime in string to datetime.

    The string should be in iso8601 datetime format.
    """
    try:
        from dateutil.parser import parse
        return parse(string)
    except ImportError:
        return string
<|endoftext|>