<commit_before>
from docutils import nodes, utils
from docutils import nodes
from docutils.parsers.rst import directives, states
from docutils.nodes import fully_normalize_name, whitespace_normalize_name
from docutils.parsers.rst.roles import set_classes

from sphinx.util.compat import Directive
from sphinx.errors import SphinxError

class PhixError(SphinxError):
    category = 'Phix error'

class argouml(nodes.General, nodes.Element):
    pass

# TODO: visit_argouml_node, depart_todo_node ?

class ArgoUmlDirective(Directive):
    # http://docutils.sourceforge.net/docs/howto/rst-directives.html

    align_h_values = ('left', 'center', 'right')
    align_v_values = ('top', 'middle', 'bottom')
    align_values = align_v_values + align_h_values

    def align(argument):
        # This is not callable as self.align.  We cannot make it a
        # staticmethod because we're saving an unbound method in
        # option_spec below.
        return directives.choice(argument, Image.align_values)

    has_content = False
    required_arguments = 2
    optional_arguments = 0
    final_argument_whitespace = True
    option_spec = {'alt': directives.unchanged,
                   'height': directives.length_or_unitless,
                   'width': directives.length_or_percentage_or_unitless,
                   'scale': directives.percentage,
                   'align': align,
                   'target': directives.unchanged_required,
                   'class': directives.class_option}

    def run(self):
        if 'align' in self.options:
            if isinstance(self.state, states.SubstitutionDef):
                # Check for align_v_values.
                if self.options['align'] not in self.align_v_values:
                    raise self.error(
                        'Error in "%s" directive: "%s" is not a valid value '
                        'for the "align" option within a substitution '
                        'definition.  Valid values for "align" are: "%s".'
                        % (self.name, self.options['align'],
                           '", "'.join(self.align_v_values)))
            elif self.options['align'] not in self.align_h_values:
                raise self.error(
                    'Error in "%s" directive: "%s" is not a valid value for '
                    'the "align" option.  Valid values for "align" are: "%s".'
                    % (self.name, self.options['align'],
                       '", "'.join(self.align_h_values)))
        messages = []
        reference = directives.uri(self.arguments[0])
        self.options['uri'] = reference
        reference_node = None
        if 'target' in self.options:
            block = states.escape2null(
                self.options['target']).splitlines()
            block = [line for line in block]
            target_type, data = self.state.parse_target(
                block, self.block_text, self.lineno)
            if target_type == 'refuri':
                reference_node = nodes.reference(refuri=data)
            elif target_type == 'refname':
                reference_node = nodes.reference(
                    refname=fully_normalize_name(data),
                    name=whitespace_normalize_name(data))
                reference_node.indirect_reference_name = data
                self.state.document.note_refname(reference_node)
            else:                           # malformed target
                messages.append(data)       # data is a system message
            del self.options['target']
        set_classes(self.options)
        image_node = nodes.image(self.block_text, **self.options)
        if reference_node:
            reference_node += image_node
            return messages + [reference_node]
        else:
            return messages + [image_node]

def setup(app):
    app.add_node(argouml)
    app.add_directive('argouml', ArgoUmlDirective)

<commit_msg>Print the argouml arguments for debugging purposes<commit_after>
from docutils import nodes, utils
from docutils import nodes
from docutils.parsers.rst import directives, states
from docutils.nodes import fully_normalize_name, whitespace_normalize_name
from docutils.parsers.rst.roles import set_classes

from sphinx.util.compat import Directive
from sphinx.errors import SphinxError

class PhixError(SphinxError):
    category = 'Phix error'

class argouml(nodes.General, nodes.Element):
    pass

# TODO: visit_argouml_node, depart_todo_node ?

class ArgoUmlDirective(Directive):
    # http://docutils.sourceforge.net/docs/howto/rst-directives.html

    align_h_values = ('left', 'center', 'right')
    align_v_values = ('top', 'middle', 'bottom')
    align_values = align_v_values + align_h_values

    def align(argument):
        # This is not callable as self.align.  We cannot make it a
        # staticmethod because we're saving an unbound method in
        # option_spec below.
        return directives.choice(argument, Image.align_values)

    has_content = False
    required_arguments = 2
    optional_arguments = 0
    final_argument_whitespace = True
    option_spec = {'alt': directives.unchanged,
                   'height': directives.length_or_unitless,
                   'width': directives.length_or_percentage_or_unitless,
                   'scale': directives.percentage,
                   'align': align,
                   'target': directives.unchanged_required,
                   'class': directives.class_option}

    def run(self):
        print "self.arguments[0] =", self.arguments[0]
        print "self.arguments[1] =", self.arguments[1]
        if 'align' in self.options:
            if isinstance(self.state, states.SubstitutionDef):
                # Check for align_v_values.
                if self.options['align'] not in self.align_v_values:
                    raise self.error(
                        'Error in "%s" directive: "%s" is not a valid value '
                        'for the "align" option within a substitution '
                        'definition.  Valid values for "align" are: "%s".'
                        % (self.name, self.options['align'],
                           '", "'.join(self.align_v_values)))
            elif self.options['align'] not in self.align_h_values:
                raise self.error(
                    'Error in "%s" directive: "%s" is not a valid value for '
                    'the "align" option.  Valid values for "align" are: "%s".'
                    % (self.name, self.options['align'],
                       '", "'.join(self.align_h_values)))
        messages = []
        reference = directives.uri(self.arguments[0])
        self.options['uri'] = reference
        reference_node = None
        if 'target' in self.options:
            block = states.escape2null(
                self.options['target']).splitlines()
            block = [line for line in block]
            target_type, data = self.state.parse_target(
                block, self.block_text, self.lineno)
            if target_type == 'refuri':
                reference_node = nodes.reference(refuri=data)
            elif target_type == 'refname':
                reference_node = nodes.reference(
                    refname=fully_normalize_name(data),
                    name=whitespace_normalize_name(data))
                reference_node.indirect_reference_name = data
                self.state.document.note_refname(reference_node)
            else:                           # malformed target
                messages.append(data)       # data is a system message
            del self.options['target']
        set_classes(self.options)
        image_node = nodes.image(self.block_text, **self.options)
        if reference_node:
            reference_node += image_node
            return messages + [reference_node]
        else:
            return messages + [image_node]

def setup(app):
    app.add_node(argouml)
    app.add_directive('argouml', ArgoUmlDirective)

<|endoftext|>


<commit_before># coding: utf-8
from __future__ import unicode_literals, division, absolute_import, print_function

import os
import sys
import subprocess
from getpass import getuser

from ._types import str_cls, type_name

if sys.platform == 'darwin':
    from ._osx.open_directory import get_user_login_shell
else:
    from ._linux.getent import get_user_login_shell

_envs = {'bytes': {}, 'unicode': {}}
_environ = {}
for key in ('HOME', 'LANG', 'USER', 'PATH'):
    if key in os.environ:
        _environ[key] = os.environ[key]

def get_shell_env(shell=None, for_subprocess=False):
    """
    Fetches the environmental variables that are set when a new shell is opened.

    :param shell:
        The shell to get the env from, if None, uses the current user's login
        shell

    :param for_subprocess:
        If True, and the code is being run in Sublime Text 2, the result will
        be byte strings instead of unicode strings

    :return:
        A 2-element tuple:

         - [0] unicode string shell path
         - [1] env dict with keys and values as unicode strings
    """

    if shell is not None and not isinstance(shell, str_cls):
        raise TypeError('shell must be a unicode string, not %s' % type_name(shell))

    if shell is None:
        shell = get_user_login_shell(get_user())
    _, shell_name = shell.rsplit('/', 1)

    output_type = 'bytes' if sys.version_info < (3,) and for_subprocess else 'unicode'

    if shell not in _envs[output_type]:
        if shell_name in ['tcsh', 'csh']:
            params = ['-l']
        else:
            params = ['-l', '-i']
        params.insert(0, shell)

        env_proc = subprocess.Popen(
            params,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            env=_environ
        )

        stdout, _ = env_proc.communicate(b'/usr/bin/env\n')

        _envs[output_type][shell] = {}

        entries = stdout.strip().split(b'\n(?=\\w+=)')
        for entry in entries:
            if entry == b'':
                continue
            parts = entry.split(b'=', 1)
            if len(parts) < 2:
                continue
            name = parts[0]
            value = parts[1]
            if output_type == 'unicode':
                name = name.decode('utf-8', 'replace')
                value = value.decode('utf-8', 'replace')
            _envs[output_type][shell][name] = value

    if output_type == 'bytes':
        shell = shell.encode('utf-8')

    return (shell, _envs[output_type][shell])

def get_user():
    """
    Returns the current username as a unicode string

    :return:
        A unicode string of the current user's username
    """

    output = getuser()
    if not isinstance(output, str_cls):
        output = output.decode('utf-8')
    return output
<commit_msg>Switch shell invocation to not be interactive<commit_after># coding: utf-8
from __future__ import unicode_literals, division, absolute_import, print_function

import os
import sys
import subprocess
from getpass import getuser

from ._types import str_cls, type_name

if sys.platform == 'darwin':
    from ._osx.open_directory import get_user_login_shell
else:
    from ._linux.getent import get_user_login_shell

_envs = {'bytes': {}, 'unicode': {}}
_environ = {}
for key in ('HOME', 'LANG', 'USER', 'PATH'):
    if key in os.environ:
        _environ[key] = os.environ[key]

def get_shell_env(shell=None, for_subprocess=False):
    """
    Fetches the environmental variables that are set when a new shell is opened.

    :param shell:
        The shell to get the env from, if None, uses the current user's login
        shell

    :param for_subprocess:
        If True, and the code is being run in Sublime Text 2, the result will
        be byte strings instead of unicode strings

    :return:
        A 2-element tuple:

         - [0] unicode string shell path
         - [1] env dict with keys and values as unicode strings
    """

    if shell is not None and not isinstance(shell, str_cls):
        raise TypeError('shell must be a unicode string, not %s' % type_name(shell))

    if shell is None:
        shell = get_user_login_shell(get_user())
    _, shell_name = shell.rsplit('/', 1)

    output_type = 'bytes' if sys.version_info < (3,) and for_subprocess else 'unicode'

    if shell not in _envs[output_type]:
        env_proc = subprocess.Popen(
            [shell, '-l'],
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            env=_environ
        )

        stdout, _ = env_proc.communicate(b'/usr/bin/env\n')

        _envs[output_type][shell] = {}

        entries = stdout.strip().split(b'\n(?=\\w+=)')
        for entry in entries:
            if entry == b'':
                continue
            parts = entry.split(b'=', 1)
            if len(parts) < 2:
                continue
            name = parts[0]
            value = parts[1]
            if output_type == 'unicode':
                name = name.decode('utf-8', 'replace')
                value = value.decode('utf-8', 'replace')
            _envs[output_type][shell][name] = value

    if output_type == 'bytes':
        shell = shell.encode('utf-8')

    return (shell, _envs[output_type][shell])

def get_user():
    """
    Returns the current username as a unicode string

    :return:
        A unicode string of the current user's username
    """

    output = getuser()
    if not isinstance(output, str_cls):
        output = output.decode('utf-8')
    return output
<|endoftext|>


<commit_before>from amqplib import client_0_8 as amqp
from carrot.backends import BaseMessage, BaseBackend
from carrot.serialize import serialize, deserialize

class Message(BaseMessage):
    """A message received by the broker.

    Usually you don't insantiate message objects yourself, but receive
    them using a :class:`carrot.messaging.Consumer`.

    :param backend: see :attr:`backend`.

    :param amqp_message: see :attr:`amqp_message`.

    :param channel: see :attr:`channel`.

    :param decoder: see :attr:`decoder`.

    .. attribute:: body
        The message body. This data is serialized,
        so you probably want to deserialize it using
        :meth:`carrot.backend.BaseMessage.decode()`.

    .. attribute:: delivery_tag
        The message delivery tag, uniquely identifying this message.

    .. attribute:: backend
        The message backend used.
        A subclass of :class:`carrot.backend.BaseBackend`.

    .. attribute:: amqp_message

        A :class:`amqplib.client_0_8.basic_message.Message` instance.

    .. attribute:: channel

        The AMQP channel. A :class:`amqplib.client_0_8.channel.Channel` instance.

    .. attribute:: decoder

        A function able to deserialize the serialized message data.

    """
    def __init__(self, backend, amqp_message, **kwargs):
        self.amqp_message = amqp_message
        self.backend = backend
        kwargs.update({
            "body": amqp_message.body,
            "delivery_tag": amqp_message.delivery_tag})

        super(BaseMessage, self).__init__(backend, **kwargs)

class Backend(BaseBackend):

    def __init__(self, connection, **kwargs):
        self.connection = connection
        self.channel = self.connection.connection.channel()
        self.encoder = kwargs.get("encoder", serialize)
        self.decoder = kwargs.get("decoder", deserialize)

    def queue_declare(self, queue, durable, exclusive, auto_delete):
        self.channel.queue_declare(queue=queue, durable=durable,
                                   exclusive=exclusive,
                                   auto_delete=auto_delete)

    def exchange_declare(self, exchange, type, durable, auto_delete):
        self.channel.exchange_declare(exchange=exchange, type=type,
                                      durable=durable,
                                      auto_delete=auto_delete)

    def queue_bind(self, queue, exchange, routing_key):
        self.channel.queue_bind(queue=queue, exchange=exchange,
                                routing_key=routing_key)

    def get(self, queue):
        message = self.channel.basic_get(queue)
        if not message:
            return None
        return Message(backend=self,
                       amqp_message=message,
                       decoder=self.decoder)

    def consume(self, queue, no_ack, callback, consumer_tag):
        self.channel.basic_consume(queue=queue, no_ack=no_ack,
                                   callback=callback,
                                   consumer_tag=consumer_tag)
        while True:
            self.channel.wait()

    def cancel(self, consumer_tag):
        self.channel.basic_cancel(consumer_tag)

    def close(self):
        if getattr(self, "channel") and self.channel.is_open:
            self.channel.close()

    def ack(self, delivery_tag):
        return self.channel.basic_ack(delivery_tag)

    def reject(self, delivery_tag):
        return self.channel.basic_reject(delivery_tag, requeue=False)

    def requeue(self, delivery_tag):
        return self.channel.basic_reject(delivery_tag, requeue=True)

    def prepare_message(self, message_data, delivery_mode):
        message = amqp.Message(message_data)
        message.properties["delivery_mode"] = delivery_mode
        return message

    def publish(self, message, exchange, routing_key):
        return self.channel.basic_publish(message, exchange=exchange,
                                          routing_key=routing_key)
<commit_msg>Rename document references to carrot.backend -> carrot.backends<commit_after>from amqplib import client_0_8 as amqp
from carrot.backends import BaseMessage, BaseBackend
from carrot.serialize import serialize, deserialize

class Message(BaseMessage):
    """A message received by the broker.

    Usually you don't insantiate message objects yourself, but receive
    them using a :class:`carrot.messaging.Consumer`.

    :param backend: see :attr:`backend`.

    :param amqp_message: see :attr:`amqp_message`.

    :param channel: see :attr:`channel`.

    :param decoder: see :attr:`decoder`.

    .. attribute:: body
        The message body. This data is serialized,
        so you probably want to deserialize it using
        :meth:`carrot.backends.BaseMessage.decode()`.

    .. attribute:: delivery_tag
        The message delivery tag, uniquely identifying this message.

    .. attribute:: backend
        The message backend used.
        A subclass of :class:`carrot.backends.BaseBackend`.

    .. attribute:: amqp_message

        A :class:`amqplib.client_0_8.basic_message.Message` instance.

    .. attribute:: channel

        The AMQP channel. A :class:`amqplib.client_0_8.channel.Channel` instance.

    .. attribute:: decoder

        A function able to deserialize the serialized message data.

    """
    def __init__(self, backend, amqp_message, **kwargs):
        self.amqp_message = amqp_message
        self.backend = backend
        kwargs.update({
            "body": amqp_message.body,
            "delivery_tag": amqp_message.delivery_tag})

        super(BaseMessage, self).__init__(backend, **kwargs)

class Backend(BaseBackend):

    def __init__(self, connection, **kwargs):
        self.connection = connection
        self.channel = self.connection.connection.channel()
        self.encoder = kwargs.get("encoder", serialize)
        self.decoder = kwargs.get("decoder", deserialize)

    def queue_declare(self, queue, durable, exclusive, auto_delete):
        self.channel.queue_declare(queue=queue, durable=durable,
                                   exclusive=exclusive,
                                   auto_delete=auto_delete)

    def exchange_declare(self, exchange, type, durable, auto_delete):
        self.channel.exchange_declare(exchange=exchange, type=type,
                                      durable=durable,
                                      auto_delete=auto_delete)

    def queue_bind(self, queue, exchange, routing_key):
        self.channel.queue_bind(queue=queue, exchange=exchange,
                                routing_key=routing_key)

    def get(self, queue):
        message = self.channel.basic_get(queue)
        if not message:
            return None
        return Message(backend=self,
                       amqp_message=message,
                       decoder=self.decoder)

    def consume(self, queue, no_ack, callback, consumer_tag):
        self.channel.basic_consume(queue=queue, no_ack=no_ack,
                                   callback=callback,
                                   consumer_tag=consumer_tag)
        while True:
            self.channel.wait()

    def cancel(self, consumer_tag):
        self.channel.basic_cancel(consumer_tag)

    def close(self):
        if getattr(self, "channel") and self.channel.is_open:
            self.channel.close()

    def ack(self, delivery_tag):
        return self.channel.basic_ack(delivery_tag)

    def reject(self, delivery_tag):
        return self.channel.basic_reject(delivery_tag, requeue=False)

    def requeue(self, delivery_tag):
        return self.channel.basic_reject(delivery_tag, requeue=True)

    def prepare_message(self, message_data, delivery_mode):
        message = amqp.Message(message_data)
        message.properties["delivery_mode"] = delivery_mode
        return message

    def publish(self, message, exchange, routing_key):
        return self.channel.basic_publish(message, exchange=exchange,
                                          routing_key=routing_key)
<|endoftext|>


<commit_before>#!/usr/bin/env python
from groupflow_shared import *
from mininet.net import *
from mininet.node import OVSSwitch, UserSwitch
from mininet.link import TCLink
from mininet.log import setLogLevel
from mininet.cli import CLI
from mininet.node import Node, RemoteController
from scipy.stats import truncnorm
from numpy.random import randint, uniform
from subprocess import *
import sys
import signal
from time import sleep, time
from datetime import datetime
from multiprocessing import Process
import numpy as np

def mcastTest(topo, interactive = False, hosts = [], log_file_name = 'test_log.log', util_link_weight = 10, link_weight_type = 'linear', replacement_mode='none'):
    membership_mean = 0.1
    membership_std_dev = 0.25
    membership_avg_bound = float(len(hosts)) / 8.0
    test_groups = []
    test_group_launch_times = []

    # Launch the external controller
    pox_arguments = ['pox.py', 'log', '--file=pox.log,w', 'openflow.discovery',
            'openflow.flow_tracker', '--query_interval=1', '--link_max_bw=27', '--link_cong_threshold=24', '--avg_smooth_factor=0.65', '--log_peak_usage=True',
            'misc.benchmark_terminator', 'openflow.igmp_manager', 'misc.groupflow_event_tracer',
            'openflow.groupflow', '--util_link_weight=' + str(util_link_weight), '--link_weight_type=' + link_weight_type, '--flow_replacement_mode=' + replacement_mode,
            'log.level', '--WARNING', '--openflow.flow_tracker=INFO']
    print 'Launching external controller: ' + str(pox_arguments[0])
    print 'Launch arguments:'
    print ' '.join(pox_arguments)

    with open(os.devnull, "w") as fnull:
        pox_process = Popen(pox_arguments, stdout=fnull, stderr=fnull, shell=False, close_fds=True)
        # Allow time for the log file to be generated
        sleep(1)

    # Determine the flow tracker log file
    pox_log_file = open('./pox.log', 'r')
    flow_log_path = None
    event_log_path = None
    got_flow_log_path = False
    got_event_log_path = False
    while (not got_flow_log_path) or (not got_event_log_path):
        pox_log = pox_log_file.readline()

        if 'Writing flow tracker info to file:' in pox_log:
            pox_log_split = pox_log.split()
            flow_log_path = pox_log_split[-1]
            got_flow_log_path = True

        if 'Writing event trace info to file:' in pox_log:
            pox_log_split = pox_log.split()
            event_log_path = pox_log_split[-1]
            got_event_log_path = True

    print 'Got flow tracker log file: ' + str(flow_log_path)
    print 'Got event trace log file: ' + str(event_log_path)
    print 'Controller initialized'
    pox_log_offset = pox_log_file.tell()
    pox_log_file.close()

    # External controller
    net = Mininet(topo, controller=RemoteController, switch=OVSSwitch, link=TCLink, build=False, autoSetMacs=True)
    # pox = RemoteController('pox', '127.0.0.1', 6633)
    net.addController('pox', RemoteController, ip = '127.0.0.1', port = 6633)
    net.start()
    for switch_name in topo.get_switch_list():
        #print switch_name + ' route add -host 127.0.0.1 dev lo'
        net.get(switch_name).controlIntf = net.get(switch_name).intf('lo')
        net.get(switch_name).cmd('route add -host 127.0.0.1 dev lo')
        #print 'pox' + ' route add -host ' + net.get(switch_name).IP() + ' dev lo'
        net.get('pox').cmd('route add -host ' + net.get(switch_name).IP() + ' dev lo')
        #print net.get(switch_name).cmd('ifconfig')

    topo.mcastConfig(net)

    #print 'Controller network configuration:'
    #print net.get('pox').cmd('ifconfig')
    #print net.get('pox').cmd('route')

    sleep_time = 8 + (float(len(hosts))/8)
    print 'Waiting ' + str(sleep_time) + ' seconds to allow for controller topology discovery'
    sleep(sleep_time)   # Allow time for the controller to detect the topology

    if interactive:
        CLI(net)
    else:
        mcast_group_last_octet = 1
        mcast_port = 5010
        rand_seed = int(time())
        print 'Using random seed: ' + str(rand_seed)
        np.random.seed(rand_seed)
        host_join_probabilities = generate_group_membership_probabilities(hosts, membership_mean, membership_std_dev, membership_avg_bound)
        print 'Host join probabilities: ' + ', '.join(str(p) for p in host_join_probabilities)
        host_join_sum = sum(p[1] for p in host_join_probabilities)
        print 'Measured mean join probability: ' + str(host_join_sum / len(host_join_probabilities))
        print 'Predicted average group size: ' + str(host_join_sum)
        i = 1
        while True:
            print 'Generating multicast group #' + str(i)
            # Choose a sending host using a uniform random distribution
            sender_index = randint(0,len(hosts))
            sender_host = hosts[sender_index]

            # Choose a random number of receivers by comparing a uniform random variable
            # against the previously generated group membership probabilities
            receivers = []
            for host_prob in host_join_probabilities:
                p = uniform(0, 1)
                if p <= host_prob[1]:
                    receivers.append(host_prob[0])

            # Initialize the group
            # Note - This method of group IP generation will need to be modified slightly to support more than
            # 255 groups
            mcast_ip = '224.1.1.{last_octet}'.format(last_octet = str(mcast_group_last_octet))
            test_groups.append(MulticastGroupDefinition(sender_host, receivers, mcast_ip, mcast_port, mcast_port + 1))
            launch_time = time()
            test_group_launch_times.append(launch_time)
            print 'Launching multicast group #' + str(i) + ' at time: ' + str(launch_time)
            test_groups[-1].launch_mcast_applications(net)
            mcast_group_last_octet = mcast_group_last_octet + 1
            mcast_port = mcast_port + 2
            i += 1
            sleep(5 + uniform(0, 5))

            # Read from the log file to determine if a link has become overloaded, and cease generating new groups if so
            print 'Check for congested link...'
            congested_link = False
            pox_log_file = open('./pox.log', 'r')
            pox_log_file.seek(pox_log_offset)
            for line in pox_log_file:
                if 'Network peak link throughout' in line:
                    line_split = line.split(' ')
                    print 'Peak Usage (Mbps): ' + line_split[-1],
                if 'Network avg link throughout:' in line:
                    line_split = line.split(' ')
                    print 'Mean Usage (Mbps): ' + line_split[-1],
                if 'Link Fully Utilized!' in line:
                    congested_link = True
                    break
            pox_log_offset = pox_log_file.tell()
            pox_log_file.close()
            if congested_link:
                print 'Detected fully utilized link, terminating simulation.'
                break
            else:
                print 'No congestion detected.'

    print 'Terminating network applications'
    for group in test_groups:
        group.terminate_mcast_applications()
    print 'Terminating controller'
    pox_process.send_signal(signal.SIGINT)
    sleep(1)
    print 'Waiting for network application termination...'
    for group in test_groups:
        group.wait_for_application_termination()
    print 'Network applications terminated'
    print 'Waiting for controller termination...'
    pox_process.send_signal(signal.SIGKILL)
    pox_process.wait()
    print 'Controller terminated'
    pox_process = None
    net.stop()

    if not interactive:
        write_final_stats_log(log_file_name, flow_log_path, event_log_path, membership_mean, membership_std_dev, membership_avg_bound, test_groups, test_group_launch_times, topo)

topos = { 'mcast_test': ( lambda: MulticastTestTopo() ) }

def print_usage_text():
    print 'GroupFlow Multicast Testing with Mininet'
    print 'Usage:'
    print '1) No arguments:'
    print '> mininet_multicast_pox'
    print 'If no arguments are provided, the script will launch a hard-coded test topology with Mininet in interactive mode.'
    print ''
    print '2) Custom topology:'
    print '> mininet_multicast_pox <topology_path>'
    print 'topology_path: If a single argument is given, the argument will be interpreted as a path to a BRITE topology. Otherwise, this functions identically to the no argument mode.'
    print ''
    print '3) Automated benchmarking:'
    print '> mininet_multicast_pox <topology_path> <iterations_to_run> <log_file_prefix> <index_of_first_log_file> <parameter_sets (number is variable and unlimited)>'
    print 'Parameter sets have the form: flow_replacement_mode,link_weight_type,util_link_weight'

if __name__ == '__main__':
    setLogLevel( 'info' )
    if len(sys.argv) >= 2:
        if '-h' in str(sys.argv[1]) or 'help' in str(sys.argv[1]):
            print_usage_text()
            sys.exit()

    if len(sys.argv) >= 6:
        # Automated simulations - Differing link usage weights in Groupflow Module
        log_prefix = sys.argv[3]
        num_iterations = int(sys.argv[2])
        first_index = int(sys.argv[4])
        util_params = []
        for param_index in range(5, len(sys.argv)):
            param_split = sys.argv[param_index].split(',')
            util_params.append((param_split[0], param_split[1], float(param_split[2])))
        topo = BriteTopo(sys.argv[1])
        hosts = topo.get_host_list()
        start_time = time()
        print 'Simulations started at: ' + str(datetime.now())
        for i in range(0,num_iterations):
            for util_param in util_params:
                p = Process(target=mcastTest, args=(topo, False, hosts, log_prefix + '_' + ','.join([util_param[0], util_param[1], str(util_param[2])]) + '_' + str(i + first_index) + '.log', util_param[2], util_param[1], util_param[0]))
                sim_start_time = time()
                p.start()
                p.join()
                sim_end_time = time()

                # Make extra sure the network terminated cleanly
                call(['python', 'kill_running_test.py'])

                print 'Simulation ' + str(i+1) + '_' + ','.join([util_param[0], util_param[1], str(util_param[2])]) + ' completed at: ' + str(datetime.now()) + ' (runtime: ' + str(sim_end_time - sim_start_time) + ' seconds)'

        end_time = time()
        print ' '
        print 'Simulations completed at: ' + str(datetime.now())
        print 'Total runtime: ' + str(end_time - start_time) + ' seconds'
        print 'Average runtime per sim: ' + str((end_time - start_time) / (num_iterations * len(util_params))) + ' seconds'

    elif len(sys.argv) >= 2:
        # Interactive mode - configures POX and multicast routes, but no automatic traffic generation
        print 'Launching BRITE defined multicast test topology'
        topo = BriteTopo(sys.argv[1])
        hosts = topo.get_host_list()
        mcastTest(topo, True, hosts)

    else:
        # Interactive mode with barebones topology
        print 'Launching default multicast test topology'
        topo = MulticastTestTopo()
        hosts = topo.get_host_list()
        mcastTest(topo, True, hosts)
<commit_msg>Added option to fix group size to a constant.<commit_after>#!/usr/bin/env python
from groupflow_shared import *
from mininet.net import *
from mininet.node import OVSSwitch, UserSwitch
from mininet.link import TCLink
from mininet.log import setLogLevel
from mininet.cli import CLI
from mininet.node import Node, RemoteController
from scipy.stats import truncnorm
from numpy.random import randint, uniform
from subprocess import *
import sys
import signal
from time import sleep, time
from datetime import datetime
from multiprocessing import Process
import numpy as np

ENABLE_FIXED_GROUP_SIZE = True
FIXED_GROUP_SIZE = 4

def mcastTest(topo, interactive = False, hosts = [], log_file_name = 'test_log.log', util_link_weight = 10, link_weight_type = 'linear', replacement_mode='none'):
    membership_mean = 0.1
    membership_std_dev = 0.25
    membership_avg_bound = float(len(hosts)) / 8.0
    test_groups = []
    test_group_launch_times = []

    # Launch the external controller
    pox_arguments = ['pox.py', 'log', '--file=pox.log,w', 'openflow.discovery',
            'openflow.flow_tracker', '--query_interval=1', '--link_max_bw=27', '--link_cong_threshold=24', '--avg_smooth_factor=0.65', '--log_peak_usage=True',
            'misc.benchmark_terminator', 'openflow.igmp_manager', 'misc.groupflow_event_tracer',
            'openflow.groupflow', '--util_link_weight=' + str(util_link_weight), '--link_weight_type=' + link_weight_type, '--flow_replacement_mode=' + replacement_mode,
            'log.level', '--WARNING', '--openflow.flow_tracker=INFO']
    print 'Launching external controller: ' + str(pox_arguments[0])
    print 'Launch arguments:'
    print ' '.join(pox_arguments)

    with open(os.devnull, "w") as fnull:
        pox_process = Popen(pox_arguments, stdout=fnull, stderr=fnull, shell=False, close_fds=True)
        # Allow time for the log file to be generated
        sleep(1)

    # Determine the flow tracker log file
    pox_log_file = open('./pox.log', 'r')
    flow_log_path = None
    event_log_path = None
    got_flow_log_path = False
    got_event_log_path = False
    while (not got_flow_log_path) or (not got_event_log_path):
        pox_log = pox_log_file.readline()

        if 'Writing flow tracker info to file:' in pox_log:
            pox_log_split = pox_log.split()
            flow_log_path = pox_log_split[-1]
            got_flow_log_path = True

        if 'Writing event trace info to file:' in pox_log:
            pox_log_split = pox_log.split()
            event_log_path = pox_log_split[-1]
            got_event_log_path = True

    print 'Got flow tracker log file: ' + str(flow_log_path)
    print 'Got event trace log file: ' + str(event_log_path)
    print 'Controller initialized'
    pox_log_offset = pox_log_file.tell()
    pox_log_file.close()

    # External controller
    net = Mininet(topo, controller=RemoteController, switch=OVSSwitch, link=TCLink, build=False, autoSetMacs=True)
    # pox = RemoteController('pox', '127.0.0.1', 6633)
    net.addController('pox', RemoteController, ip = '127.0.0.1', port = 6633)
    net.start()
    for switch_name in topo.get_switch_list():
        #print switch_name + ' route add -host 127.0.0.1 dev lo'
        net.get(switch_name).controlIntf = net.get(switch_name).intf('lo')
        net.get(switch_name).cmd('route add -host 127.0.0.1 dev lo')
        #print 'pox' + ' route add -host ' + net.get(switch_name).IP() + ' dev lo'
        net.get('pox').cmd('route add -host ' + net.get(switch_name).IP() + ' dev lo')
        #print net.get(switch_name).cmd('ifconfig')

    topo.mcastConfig(net)

    #print 'Controller network configuration:'
    #print net.get('pox').cmd('ifconfig')
    #print net.get('pox').cmd('route')

    sleep_time = 8 + (float(len(hosts))/8)
    print 'Waiting ' + str(sleep_time) + ' seconds to allow for controller topology discovery'
    sleep(sleep_time)   # Allow time for the controller to detect the topology

    if interactive:
        CLI(net)
    else:
        mcast_group_last_octet = 1
        mcast_port = 5010
        rand_seed = int(time())
        print 'Using random seed: ' + str(rand_seed)
        np.random.seed(rand_seed)
        host_join_probabilities = generate_group_membership_probabilities(hosts, membership_mean, membership_std_dev, membership_avg_bound)
        print 'Host join probabilities: ' + ', '.join(str(p) for p in host_join_probabilities)
        host_join_sum = sum(p[1] for p in host_join_probabilities)
        print 'Measured mean join probability: ' + str(host_join_sum / len(host_join_probabilities))
        print 'Predicted average group size: ' + str(host_join_sum)
        i = 1
        while True:
            print 'Generating multicast group #' + str(i)
            # Choose a sending host using a uniform random distribution
            sender_index = randint(0,len(hosts))
            sender_host = hosts[sender_index]

            receivers = []
            if ENABLE_FIXED_GROUP_SIZE:
                while len(receivers) < FIXED_GROUP_SIZE:
                    receivers.append(hosts[randint(0,len(hosts))])
                    receivers = list(set(receivers))
            else:
                # Choose a random number of receivers by comparing a uniform random variable
                # against the previously generated group membership probabilities
                for host_prob in host_join_probabilities:
                    p = uniform(0, 1)
                    if p <= host_prob[1]:
                        receivers.append(host_prob[0])

            # Initialize the group
            # Note - This method of group IP generation will need to be modified slightly to support more than
            # 255 groups
            mcast_ip = '224.1.1.{last_octet}'.format(last_octet = str(mcast_group_last_octet))
            test_groups.append(MulticastGroupDefinition(sender_host, receivers, mcast_ip, mcast_port, mcast_port + 1))
            launch_time = time()
            test_group_launch_times.append(launch_time)
            print 'Launching multicast group #' + str(i) + ' at time: ' + str(launch_time)
            test_groups[-1].launch_mcast_applications(net)
            mcast_group_last_octet = mcast_group_last_octet + 1
            mcast_port = mcast_port + 2
            i += 1
            sleep(5 + uniform(0, 5))

            # Read from the log file to determine if a link has become overloaded, and cease generating new groups if so
            print 'Check for congested link...'
            congested_link = False
            pox_log_file = open('./pox.log', 'r')
            pox_log_file.seek(pox_log_offset)
            for line in pox_log_file:
                if 'Network peak link throughout' in line:
                    line_split = line.split(' ')
                    print 'Peak Usage (Mbps): ' + line_split[-1],
                if 'Network avg link throughout:' in line:
                    line_split = line.split(' ')
                    print 'Mean Usage (Mbps): ' + line_split[-1],
                if 'Link Fully Utilized!' in line:
                    congested_link = True
                    break
            pox_log_offset = pox_log_file.tell()
            pox_log_file.close()
            if congested_link:
                print 'Detected fully utilized link, terminating simulation.'
                break
            else:
                print 'No congestion detected.'

    print 'Terminating network applications'
    for group in test_groups:
        group.terminate_mcast_applications()
    print 'Terminating controller'
    pox_process.send_signal(signal.SIGINT)
    sleep(1)
    print 'Waiting for network application termination...'
    for group in test_groups:
        group.wait_for_application_termination()
    print 'Network applications terminated'
    print 'Waiting for controller termination...'
    pox_process.send_signal(signal.SIGKILL)
    pox_process.wait()
    print 'Controller terminated'
    pox_process = None
    net.stop()

    if not interactive:
        write_final_stats_log(log_file_name, flow_log_path, event_log_path, membership_mean, membership_std_dev, membership_avg_bound, test_groups, test_group_launch_times, topo)

topos = { 'mcast_test': ( lambda: MulticastTestTopo() ) }

def print_usage_text():
    print 'GroupFlow Multicast Testing with Mininet'
    print 'Usage:'
    print '1) No arguments:'
    print '> mininet_multicast_pox'
    print 'If no arguments are provided, the script will launch a hard-coded test topology with Mininet in interactive mode.'
    print ''
    print '2) Custom topology:'
    print '> mininet_multicast_pox <topology_path>'
    print 'topology_path: If a single argument is given, the argument will be interpreted as a path to a BRITE topology. Otherwise, this functions identically to the no argument mode.'
    print ''
    print '3) Automated benchmarking:'
    print '> mininet_multicast_pox <topology_path> <iterations_to_run> <log_file_prefix> <index_of_first_log_file> <parameter_sets (number is variable and unlimited)>'
    print 'Parameter sets have the form: flow_replacement_mode,link_weight_type,util_link_weight'

if __name__ == '__main__':
    setLogLevel( 'info' )
    if len(sys.argv) >= 2:
        if '-h' in str(sys.argv[1]) or 'help' in str(sys.argv[1]):
            print_usage_text()
            sys.exit()

    if len(sys.argv) >= 6:
        # Automated simulations - Differing link usage weights in Groupflow Module
        log_prefix = sys.argv[3]
        num_iterations = int(sys.argv[2])
        first_index = int(sys.argv[4])
        util_params = []
        for param_index in range(5, len(sys.argv)):
            param_split = sys.argv[param_index].split(',')
            util_params.append((param_split[0], param_split[1], float(param_split[2])))
        topo = BriteTopo(sys.argv[1])
        hosts = topo.get_host_list()
        start_time = time()
        print 'Simulations started at: ' + str(datetime.now())
        for i in range(0,num_iterations):
            for util_param in util_params:
                p = Process(target=mcastTest, args=(topo, False, hosts, log_prefix + '_' + ','.join([util_param[0], util_param[1], str(util_param[2])]) + '_' + str(i + first_index) + '.log', util_param[2], util_param[1], util_param[0]))
                sim_start_time = time()
                p.start()
                p.join()
                sim_end_time = time()

                # Make extra sure the network terminated cleanly
                call(['python', 'kill_running_test.py'])

                print 'Simulation ' + str(i+1) + '_' + ','.join([util_param[0], util_param[1], str(util_param[2])]) + ' completed at: ' + str(datetime.now()) + ' (runtime: ' + str(sim_end_time - sim_start_time) + ' seconds)'

        end_time = time()
        print ' '
        print 'Simulations completed at: ' + str(datetime.now())
        print 'Total runtime: ' + str(end_time - start_time) + ' seconds'
        print 'Average runtime per sim: ' + str((end_time - start_time) / (num_iterations * len(util_params))) + ' seconds'

    elif len(sys.argv) >= 2:
        # Interactive mode - configures POX and multicast routes, but no automatic traffic generation
        print 'Launching BRITE defined multicast test topology'
        topo = BriteTopo(sys.argv[1])
        hosts = topo.get_host_list()
        mcastTest(topo, True, hosts)

    else:
        # Interactive mode with barebones topology
        print 'Launching default multicast test topology'
        topo = MulticastTestTopo()
        hosts = topo.get_host_list()
        mcastTest(topo, True, hosts)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""Sample Sheet generation for BCL2FASTQ pipeline
"""
# --- standard library imports
#
import sys
import os
import logging
import argparse
from collections import namedtuple

#--- third-party imports
#
import yaml
import requests
import xml.etree.ElementTree as ET

#--- project specific imports
#
from pipelines import get_machine_run_flowcell_id
#from bcl2fastq import MuxUnit
# WARNING changes here, must be reflected in bcl2fastq.py as well
MuxUnit = namedtuple('MuxUnit', ['run_id', 'flowcell_id', 'mux_id', 'lane_ids',
                                 'mux_dir', 'barcode_mismatches'])

__author__ = "Lavanya Veeravalli"
__email__ = "veeravallil@gis.a-star.edu.sg"
__copyright__ = "2016 Genome Institute of Singapore"
__license__ = "The MIT License (MIT)"

# global logger
logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter(
    '[{asctime}] {levelname:8s} {filename} {message}', style='{'))
logger.addHandler(handler)

SAMPLESHEET_CSV = "samplesheet.csv"
USEBASES_CFG = "usebases.yaml"
MUXINFO_CFG = "muxinfo.yaml"
DEFAULT_BARCODE_MISMATCHES = None
SAMPLESHEET_HEADER = '[Data]'+'\n'+ 'Lane,Sample_ID,Sample_Name,Sample_Plate,Sample_Well,I7_Index_ID,index,I5_Index_ID,index2,Sample_Project,Description'

def getdirs(args):
    """gets directories from args and checks existance
    """
    rundir = args.rundir
    if not os.path.exists(rundir):
        logger.fatal("rundir '%s' does not exist under Run directory.\n" % (rundir))
        sys.exit(1)

    runinfo = os.path.join(rundir + '/RunInfo.xml')
    if not os.path.exists(runinfo):
        logger.fatal("RunInfo '%s' does not exist under Run directory.\n" % (runinfo))
        sys.exit(1)

    outdir = args.outdir
    if not os.path.exists(outdir):
        logger.fatal("output directory '%s' does not exist.\n" % (outdir))
        sys.exit(1)

    return(rundir, outdir, runinfo)

def generate_usebases(barcode_lens, runinfo):
    """FIXME:add-doc
    """
    tree = ET.parse(runinfo)
    root = tree.getroot()
    ub_list = []

    # for each lane and its barcode lengths
    for k, v in sorted(barcode_lens.items()):
        # v is list of barcode_len tuples
        assert len(set(v)) == 1, ("Different barcode length in lane {}".format(k))
        bc1, bc2 = v[0]# since all v's are the same

        ub = ""
        #if test:
        for read in root.iter('Read'):
            numcyc = int(read.attrib['NumCycles'])-1
            if read.attrib['IsIndexedRead'] == 'N':
                ub += 'Y' + str(numcyc) + 'n*,'
            elif read.attrib['IsIndexedRead'] == 'Y':
                if read.attrib['Number'] == '2':    ### BC1
                    if bc1 > 0:
                        ub += 'I'+str(bc1)+'n*,'
                    else:
                        ub += 'n*'+','
                if read.attrib['Number'] == '3':    ### BC2
                    if bc2 > 0:
                        ub += 'I'+str(bc2)+'n*,'
                    else:
                        ub += 'n*'+','
        ub = ub[:-1]
        ub_list.append(str(k + ':' + ub))
    return ub_list

def main():
    """
    The main function
    """

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--force-overwrite",
                        action="store_true",
                        help="Force overwriting of output files")
    parser.add_argument("-r", "--rundir",
                        dest="rundir",
                        required=True,
                        help="rundir, e.g. /mnt/seq/userrig/HS004/HS004-PE-R00139_BC6A7HANXX")
    parser.add_argument("-o", "--outdir",
                        required=True,
                        dest="outdir",
                        help="Output directory")
    parser.add_argument('-v', '--verbose', action='count', default=0,
                        help="Increase verbosity")
    parser.add_argument('-q', '--quiet', action='count', default=0,
                        help="Decrease verbosity")
    args = parser.parse_args()

    # Repeateable -v and -q for setting logging level.
    # See https://www.reddit.com/r/Python/comments/3nctlm/what_python_tools_should_i_be_using_on_every/
    # and https://gist.github.com/andreas-wilm/b6031a84a33e652680d4
    # script -vv -> DEBUG
    # script -v -> INFO
    # script -> WARNING
    # script -q -> ERROR
    # script -qq -> CRITICAL
    # script -qqq -> no logging at all
    logger.setLevel(logging.WARN + 10*args.quiet - 10*args.verbose)

    (rundir, outdir, runinfo) = getdirs(args)
    samplesheet_csv = os.path.join(outdir, SAMPLESHEET_CSV)
    usebases_cfg = os.path.join(outdir, USEBASES_CFG)
    muxinfo_cfg = os.path.join(outdir, MUXINFO_CFG)
    for f in [samplesheet_csv, usebases_cfg, muxinfo_cfg]:
        if not args.force_overwrite and os.path.exists(f):
            logger.fatal("Refusing to overwrite existing file {}".format(f))
            sys.exit(1)

    _, run_num, flowcellid = get_machine_run_flowcell_id(rundir)

    logger.info("Querying ELM for {}".format(run_num))
    #DEVELOPMENT URL
    #rest_url = 'http://dlap51v:8080/elm/rest/seqrun/illumina/' + run_num + '/detailanalysis/json'
    #PRODUCTION url
    rest_url = 'http://qldb01.gis.a-star.edu.sg:8080/rest/seqrun/illumina/' + run_num + '/detailanalysis/json'
    #NEW DEVELOPMENT url
    #rest_url = 'http://qldb01:8180/rest/seqrun/illumina/'+ run_num + '/detailanalysis/json'
    response = requests.get(rest_url)
    if response.status_code != requests.codes.ok:
        response.raise_for_status()
    rest_data = response.json()

    run_id = rest_data['runId']
    #counter = 0
    if rest_data['runPass'] != 'Pass':
        logger.warning("Skipping non-passed run")
        # NOTE: exit 0 and missing output files is the upstream signal for a failed run
        sys.exit(0)

    # this is the master samplesheet
    logger.info("Writing to {}".format(samplesheet_csv))
    # keys: lanes, values are barcode lens in lane (always two tuples, -1 if not present)
    barcode_lens = {}
    mux_units = dict()

    with open(samplesheet_csv, 'w') as fh_out:
        fh_out.write(SAMPLESHEET_HEADER + '\n')
        for rows in rest_data['lanes']:
            if rows['lanePass'] != 'Pass':
                continue
            BCL_Mismatch = []
            if "MUX" in rows['libraryId']:
                # multiplexed
                #counter = 0
                for child in rows['Children']:
                    #print(child)
                    #counter += 1
                    #id = 'S' + str(counter)
                    if 'BCL_Mismatch' in child:
                        BCL_Mismatch.append(child['BCL_Mismatch'])
                        # older samples have no values and that's okay

                    if "-" in child['barcode']:
                        # dual index
                        index = child['barcode'].split('-')
                        sample = rows['laneId']+',Sample_'+child['libraryId']+','+child['libraryId']+'-'+child['barcode']+',,,,'+ index[0] +',,'+ index[1] + ',' +'Project_'+rows['libraryId']+','+child['libtech']
                        index_lens = (len((index[0])), len((index[1])))
                    else:
                        sample = rows['laneId']+',Sample_'+child['libraryId']+','+child['libraryId']+'-'+child['barcode']+',,,,'+child['barcode']+',,,'+'Project_'+rows['libraryId']+','+child['libtech']
                        index_lens = (len(child['barcode']), -1)
                        print(sample)
                    barcode_lens.setdefault(rows['laneId'], []).append(index_lens)
                    fh_out.write(sample+ '\n')

            else:# non-multiplexed
                sample = rows['laneId']+',Sample_'+rows['libraryId']+','+rows['libraryId']+'-NoIndex'+',,,,,,,'+'Project_'+rows['libraryId']+','+rows['libtech']
                index_lens = (-1, -1)
                barcode_lens.setdefault(rows['laneId'], []).append(index_lens)
                fh_out.write(sample + '\n')

            #Barcode mismatch has to be the same for all the libraries in one MUX. Otherwise default mismatch value to be used
            if len(set(BCL_Mismatch)) == 1:
                barcode_mismatches = BCL_Mismatch[0]
            else:
                barcode_mismatches = DEFAULT_BARCODE_MISMATCHES# FIXME get from ELM
            mu = MuxUnit._make([run_id, flowcellid, rows['libraryId'], [rows['laneId']], 'Project_' + rows['libraryId'], barcode_mismatches])
            # merge lane into existing mux if needed
            if mu.mux_id in mux_units:
                mu_orig = mux_units[mu.mux_id]
                assert mu.barcode_mismatches == mu_orig.barcode_mismatches
                assert len(mu.lane_ids) == 1# is a list by design but just one element. otherwise below fails
                lane_ids = mu_orig.lane_ids.extend(mu.lane_ids)
                mu_orig = mu_orig._replace(lane_ids=lane_ids)
            else:
                mux_units[mu.mux_id] = mu

    logger.info("Writing to {}".format(usebases_cfg))
    usebases = generate_usebases(barcode_lens, runinfo)
    with open(usebases_cfg, 'w') as fh:
        fh.write(yaml.dump(dict(usebases=usebases), default_flow_style=True))

    logger.info("Writing to {}".format(muxinfo_cfg))
    with open(muxinfo_cfg, 'w') as fh:
        fh.write(yaml.dump([dict(mu._asdict()) for mu in mux_units.values()], default_flow_style=True))

if __name__ == "__main__":
    main()
    logger.info("Successful program exit")
<commit_msg>Removed some print details about sample<commit_after>#!/usr/bin/env python3
"""Sample Sheet generation for BCL2FASTQ pipeline
"""
# --- standard library imports
#
import sys
import os
import logging
import argparse
from collections import namedtuple

#--- third-party imports
#
import yaml
import requests
import xml.etree.ElementTree as ET

#--- project specific imports
#
from pipelines import get_machine_run_flowcell_id
#from bcl2fastq import MuxUnit
# WARNING changes here, must be reflected in bcl2fastq.py as well
MuxUnit = namedtuple('MuxUnit', ['run_id', 'flowcell_id', 'mux_id', 'lane_ids',
                                 'mux_dir', 'barcode_mismatches'])

__author__ = "Lavanya Veeravalli"
__email__ = "veeravallil@gis.a-star.edu.sg"
__copyright__ = "2016 Genome Institute of Singapore"
__license__ = "The MIT License (MIT)"

# global logger
logger = logging.getLogger(__name__)
handler = logging.StreamHandler()
handler.setFormatter(logging.Formatter(
    '[{asctime}] {levelname:8s} {filename} {message}', style='{'))
logger.addHandler(handler)

SAMPLESHEET_CSV = "samplesheet.csv"
USEBASES_CFG = "usebases.yaml"
MUXINFO_CFG = "muxinfo.yaml"
DEFAULT_BARCODE_MISMATCHES = None
SAMPLESHEET_HEADER = '[Data]'+'\n'+ 'Lane,Sample_ID,Sample_Name,Sample_Plate,Sample_Well,I7_Index_ID,index,I5_Index_ID,index2,Sample_Project,Description'

def getdirs(args):
    """gets directories from args and checks existance
    """
    rundir = args.rundir
    if not os.path.exists(rundir):
        logger.fatal("rundir '%s' does not exist under Run directory.\n" % (rundir))
        sys.exit(1)

    runinfo = os.path.join(rundir + '/RunInfo.xml')
    if not os.path.exists(runinfo):
        logger.fatal("RunInfo '%s' does not exist under Run directory.\n" % (runinfo))
        sys.exit(1)

    outdir = args.outdir
    if not os.path.exists(outdir):
        logger.fatal("output directory '%s' does not exist.\n" % (outdir))
        sys.exit(1)

    return(rundir, outdir, runinfo)

def generate_usebases(barcode_lens, runinfo):
    """FIXME:add-doc
    """
    tree = ET.parse(runinfo)
    root = tree.getroot()
    ub_list = []

    # for each lane and its barcode lengths
    for k, v in sorted(barcode_lens.items()):
        # v is list of barcode_len tuples
        assert len(set(v)) == 1, ("Different barcode length in lane {}".format(k))
        bc1, bc2 = v[0]# since all v's are the same

        ub = ""
        #if test:
        for read in root.iter('Read'):
            numcyc = int(read.attrib['NumCycles'])-1
            if read.attrib['IsIndexedRead'] == 'N':
                ub += 'Y' + str(numcyc) + 'n*,'
            elif read.attrib['IsIndexedRead'] == 'Y':
                if read.attrib['Number'] == '2':    ### BC1
                    if bc1 > 0:
                        ub += 'I'+str(bc1)+'n*,'
                    else:
                        ub += 'n*'+','
                if read.attrib['Number'] == '3':    ### BC2
                    if bc2 > 0:
                        ub += 'I'+str(bc2)+'n*,'
                    else:
                        ub += 'n*'+','
        ub = ub[:-1]
        ub_list.append(str(k + ':' + ub))
    return ub_list

def main():
    """
    The main function
    """

    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--force-overwrite",
                        action="store_true",
                        help="Force overwriting of output files")
    parser.add_argument("-r", "--rundir",
                        dest="rundir",
                        required=True,
                        help="rundir, e.g. /mnt/seq/userrig/HS004/HS004-PE-R00139_BC6A7HANXX")
    parser.add_argument("-o", "--outdir",
                        required=True,
                        dest="outdir",
                        help="Output directory")
    parser.add_argument('-v', '--verbose', action='count', default=0,
                        help="Increase verbosity")
    parser.add_argument('-q', '--quiet', action='count', default=0,
                        help="Decrease verbosity")
    args = parser.parse_args()

    # Repeateable -v and -q for setting logging level.
    # See https://www.reddit.com/r/Python/comments/3nctlm/what_python_tools_should_i_be_using_on_every/
    # and https://gist.github.com/andreas-wilm/b6031a84a33e652680d4
    # script -vv -> DEBUG
    # script -v -> INFO
    # script -> WARNING
    # script -q -> ERROR
    # script -qq -> CRITICAL
    # script -qqq -> no logging at all
    logger.setLevel(logging.WARN + 10*args.quiet - 10*args.verbose)

    (rundir, outdir, runinfo) = getdirs(args)
    samplesheet_csv = os.path.join(outdir, SAMPLESHEET_CSV)
    usebases_cfg = os.path.join(outdir, USEBASES_CFG)
    muxinfo_cfg = os.path.join(outdir, MUXINFO_CFG)
    for f in [samplesheet_csv, usebases_cfg, muxinfo_cfg]:
        if not args.force_overwrite and os.path.exists(f):
            logger.fatal("Refusing to overwrite existing file {}".format(f))
            sys.exit(1)

    _, run_num, flowcellid = get_machine_run_flowcell_id(rundir)

    logger.info("Querying ELM for {}".format(run_num))
    #DEVELOPMENT URL
    #rest_url = 'http://dlap51v:8080/elm/rest/seqrun/illumina/' + run_num + '/detailanalysis/json'
    #PRODUCTION url
    rest_url = 'http://qldb01.gis.a-star.edu.sg:8080/rest/seqrun/illumina/' + run_num + '/detailanalysis/json'
    #NEW DEVELOPMENT url
    #rest_url = 'http://qldb01:8180/rest/seqrun/illumina/'+ run_num + '/detailanalysis/json'
    response = requests.get(rest_url)
    if response.status_code != requests.codes.ok:
        response.raise_for_status()
    rest_data = response.json()

    run_id = rest_data['runId']
    #counter = 0
    if rest_data['runPass'] != 'Pass':
        logger.warning("Skipping non-passed run")
        # NOTE: exit 0 and missing output files is the upstream signal for a failed run
        sys.exit(0)

    # this is the master samplesheet
    logger.info("Writing to {}".format(samplesheet_csv))
    # keys: lanes, values are barcode lens in lane (always two tuples, -1 if not present)
    barcode_lens = {}
    mux_units = dict()

    with open(samplesheet_csv, 'w') as fh_out:
        fh_out.write(SAMPLESHEET_HEADER + '\n')
        for rows in rest_data['lanes']:
            if rows['lanePass'] != 'Pass':
                continue
            BCL_Mismatch = []
            if "MUX" in rows['libraryId']:
                # multiplexed
                #counter = 0
                for child in rows['Children']:
                    #counter += 1
                    #id = 'S' + str(counter)
                    if 'BCL_Mismatch' in child:
                        BCL_Mismatch.append(child['BCL_Mismatch'])
                        # older samples have no values and that's okay

                    if "-" in child['barcode']:
                        # dual index
                        index = child['barcode'].split('-')
                        sample = rows['laneId']+',Sample_'+child['libraryId']+','+child['libraryId']+'-'+child['barcode']+',,,,'+ index[0] +',,'+ index[1] + ',' +'Project_'+rows['libraryId']+','+child['libtech']
                        index_lens = (len((index[0])), len((index[1])))
                    else:
                        sample = rows['laneId']+',Sample_'+child['libraryId']+','+child['libraryId']+'-'+child['barcode']+',,,,'+child['barcode']+',,,'+'Project_'+rows['libraryId']+','+child['libtech']
                        index_lens = (len(child['barcode']), -1)
                    barcode_lens.setdefault(rows['laneId'], []).append(index_lens)
                    fh_out.write(sample+ '\n')

            else:# non-multiplexed
                sample = rows['laneId']+',Sample_'+rows['libraryId']+','+rows['libraryId']+'-NoIndex'+',,,,,,,'+'Project_'+rows['libraryId']+','+rows['libtech']
                index_lens = (-1, -1)
                barcode_lens.setdefault(rows['laneId'], []).append(index_lens)
                fh_out.write(sample + '\n')

            #Barcode mismatch has to be the same for all the libraries in one MUX. Otherwise default mismatch value to be used
            if len(set(BCL_Mismatch)) == 1:
                barcode_mismatches = BCL_Mismatch[0]
            else:
                barcode_mismatches = DEFAULT_BARCODE_MISMATCHES# FIXME get from ELM
            mu = MuxUnit._make([run_id, flowcellid, rows['libraryId'], [rows['laneId']], 'Project_' + rows['libraryId'], barcode_mismatches])
            # merge lane into existing mux if needed
            if mu.mux_id in mux_units:
                mu_orig = mux_units[mu.mux_id]
                assert mu.barcode_mismatches == mu_orig.barcode_mismatches
                assert len(mu.lane_ids) == 1# is a list by design but just one element. otherwise below fails
                lane_ids = mu_orig.lane_ids.extend(mu.lane_ids)
                mu_orig = mu_orig._replace(lane_ids=lane_ids)
            else:
                mux_units[mu.mux_id] = mu

    logger.info("Writing to {}".format(usebases_cfg))
    usebases = generate_usebases(barcode_lens, runinfo)
    with open(usebases_cfg, 'w') as fh:
        fh.write(yaml.dump(dict(usebases=usebases), default_flow_style=True))

    logger.info("Writing to {}".format(muxinfo_cfg))
    with open(muxinfo_cfg, 'w') as fh:
        fh.write(yaml.dump([dict(mu._asdict()) for mu in mux_units.values()], default_flow_style=True))

if __name__ == "__main__":
    main()
    logger.info("Successful program exit")
<|endoftext|>


<commit_before># Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class PyNumba(PythonPackage):
    """NumPy aware dynamic Python compiler using LLVM"""

    homepage = "https://numba.pydata.org/"
    pypi = "numba/numba-0.35.0.tar.gz"

    version('0.50.1', sha256='89e81b51b880f9b18c82b7095beaccc6856fcf84ba29c4f0ced42e4e5748a3a7')
    version('0.48.0', sha256='9d21bc77e67006b5723052840c88cc59248e079a907cc68f1a1a264e1eaba017')
    version('0.40.1', sha256='52d046c13bcf0de79dbfb936874b7228f141b9b8e3447cc35855e9ad3e12aa33')
    version('0.35.0', sha256='11564937757605bee590c5758c73cfe9fd6d569726b56d970316a6228971ecc3')

    depends_on('python@3.6:', type=('build', 'run'), when='@0.48.0:')
    depends_on('python@3.3:3.7.9999', type=('build', 'run'), when='@0.40.1:0.47.999')
    depends_on('python@3.3:3.6.9999', type=('build', 'run'), when='@:0.35.0')
    depends_on('py-numpy@1.15:', type=('build', 'run'), when='@0.48.0:')
    depends_on('py-numpy@1.10:1.99', type=('build', 'run'))
    depends_on('py-setuptools', type=('build', 'run'))

    # Note: As of 1 Nov 2018, 0.25 was the latest version of py-llvmlite.
    # That's why it was chosen as an upper bound in the following depends_on
    # calls.  If newer versions maintain backwards compatibility, the calls
    # can be updated accordingly.
    depends_on('py-llvmlite@0.33:0.34', type=('build', 'run'), when='@0.50.1')
    depends_on('py-llvmlite@0.31.0:0.32.0', type=('build', 'run'), when='@0.48.0')
    depends_on('py-llvmlite@0.25:', type=('build', 'run'), when='@0.40.1:')
    depends_on('py-llvmlite@0.20:0.25', type=('build', 'run'), when='@0.35.1')

    depends_on('py-argparse', type=('build', 'run'), when='^python@:2.6')
    depends_on('py-funcsigs', type=('build', 'run'), when='^python@:3.3.99')
    depends_on('py-singledispatch', type=('build', 'run'), when='^python@:3.3.99')

    # Version 6.0.0 of llvm had a hidden symbol which breaks numba at runtime.
    # See https://reviews.llvm.org/D44140
    conflicts('^llvm@6.0.0')
<commit_msg>[py-numba] added version 0.51.1 (#21182)<commit_after># Copyright 2013-2021 Lawrence Livermore National Security, LLC and other
# Spack Project Developers. See the top-level COPYRIGHT file for details.
#
# SPDX-License-Identifier: (Apache-2.0 OR MIT)

from spack import *

class PyNumba(PythonPackage):
    """NumPy aware dynamic Python compiler using LLVM"""

    homepage = "https://numba.pydata.org/"
    pypi = "numba/numba-0.35.0.tar.gz"

    version('0.51.1', sha256='1e765b1a41535684bf3b0465c1d0a24dcbbff6af325270c8f4dad924c0940160')
    version('0.50.1', sha256='89e81b51b880f9b18c82b7095beaccc6856fcf84ba29c4f0ced42e4e5748a3a7')
    version('0.48.0', sha256='9d21bc77e67006b5723052840c88cc59248e079a907cc68f1a1a264e1eaba017')
    version('0.40.1', sha256='52d046c13bcf0de79dbfb936874b7228f141b9b8e3447cc35855e9ad3e12aa33')
    version('0.35.0', sha256='11564937757605bee590c5758c73cfe9fd6d569726b56d970316a6228971ecc3')

    depends_on('python@3.6:', type=('build', 'run'), when='@0.48.0:')
    depends_on('python@3.3:3.7.9999', type=('build', 'run'), when='@0.40.1:0.47.999')
    depends_on('python@3.3:3.6.9999', type=('build', 'run'), when='@:0.35.0')
    depends_on('py-numpy@1.15:', type=('build', 'run'), when='@0.48.0:')
    depends_on('py-numpy@1.10:1.99', type=('build', 'run'))
    depends_on('py-setuptools', type=('build', 'run'))

    # Note: As of 1 Nov 2018, 0.25 was the latest version of py-llvmlite.
    # That's why it was chosen as an upper bound in the following depends_on
    # calls.  If newer versions maintain backwards compatibility, the calls
    # can be updated accordingly.
    depends_on('py-llvmlite@0.34.0:0.34.999', type=('build', 'run'), when='@0.51.1')
    depends_on('py-llvmlite@0.33.0:0.33.999', type=('build', 'run'), when='@0.50.1')
    depends_on('py-llvmlite@0.31.0:0.31.999', type=('build', 'run'), when='@0.48.0')
    depends_on('py-llvmlite@0.25.0:',         type=('build', 'run'), when='@0.40.1')
    depends_on('py-llvmlite@0.20:0.25',       type=('build', 'run'), when='@0.35.1')

    depends_on('py-argparse', type=('build', 'run'), when='^python@:2.6')
    depends_on('py-funcsigs', type=('build', 'run'), when='^python@:3.3.99')
    depends_on('py-singledispatch', type=('build', 'run'), when='^python@:3.3.99')

    # Version 6.0.0 of llvm had a hidden symbol which breaks numba at runtime.
    # See https://reviews.llvm.org/D44140
    conflicts('^llvm@6.0.0')
<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# Copyright 2018 The Bazel Authors. All rights reserved.
#
# 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.

import argparse
import os
import queue
import sys
import threading
import urllib.error
import urllib.request
import yaml

import gcloud

CONFIG_URL = "https://raw.githubusercontent.com/bazelbuild/continuous-integration/master/buildkite/instances.yml"
LOCAL_CONFIG_FILE_NAME = "instances.yml"

WORK_QUEUE = queue.Queue()

def worker():
    while True:
        item = WORK_QUEUE.get()
        if not item:
            break
        try:
            # We take a few keys out of the config item. The rest is passed
            # as-is to create_instance_template() and thus to the gcloud
            # command line tool.
            count = item.pop("count")
            instance_group_name = item.pop("name")
            project = item.pop("project")
            zone = item.pop("zone", None)
            region = item.pop("region", None)

            if not project:
                raise Exception("Invalid instance config, no project name set")

            if not zone and not region:
                raise Exception(
                    "Invalid instance config, either zone or region must be specified"
                )

            template_name = instance_group_name + "-template"

            if zone is not None:
                if (gcloud.delete_instance_group(
                        instance_group_name, project=project,
                        zone=zone).returncode == 0):
                    print("Deleted existing instance group: {}".format(
                        instance_group_name))
            elif region is not None:
                if (gcloud.delete_instance_group(
                        instance_group_name, project=project,
                        region=region).returncode == 0):
                    print("Deleted existing instance group: {}".format(
                        instance_group_name))

            if gcloud.delete_instance_template(
                    template_name, project=project).returncode == 0:
                print("Deleted existing VM template: {}".format(template_name))

            gcloud.create_instance_template(
                template_name, project=project, **item)

            kwargs = {
                "project": project,
                "base_instance_name": instance_group_name,
                "size": count,
                "template": template_name,
            }
            if zone is not None:
                kwargs["zone"] = zone
            elif region is not None:
                kwargs["region"] = region
            gcloud.create_instance_group(instance_group_name, **kwargs)
        finally:
            WORK_QUEUE.task_done()

def read_config_file(use_local_config):
    content = None
    if use_local_config:
        path = os.path.join(os.getcwd(), LOCAL_CONFIG_FILE_NAME)
        with open(path, "rb") as fd:
            content = fd.read().decode("utf-8")
    else:
        with urllib.request.urlopen(CONFIG_URL) as resp:
            content = resp.read().decode("utf-8")
    return yaml.safe_load(content)

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="Bazel CI Instance Creation")
    parser.add_argument(
        "names",
        type=str,
        nargs="*",
        help="List of instance (group) names that should be created. "
        'These values must correspond to "name" entries in the '
        'Yaml configuration, e.g. "bk-pipeline-ubuntu1804-java8".',
    )
    parser.add_argument(
        "--local_config",
        action="store_true",
        help="Whether to read the configuration from CWD/%s" %
        LOCAL_CONFIG_FILE_NAME,
    )

    args = parser.parse_args(argv)
    config = read_config_file(args.local_config)

    # Verify names passed on the command-line.
    valid_names = [item["name"] for item in config["instance_groups"]]
    for name in args.names:
        if name not in valid_names:
            print("Unknown instance name: {}!".format(name))
            print("\nValid instance names are: {}".format(
                " ".join(valid_names)))
            return 1
    if not args.names:
        parser.print_help()
        print("\nValid instance names are: {}".format(" ".join(valid_names)))
        return 1

    # Put VM creation instructions into the work queue.
    for instance in config["instance_groups"]:
        if instance["name"] not in args.names:
            continue
        WORK_QUEUE.put({**config["default_vm"], **instance})

    # Spawn worker threads that will create the VMs.
    threads = []
    for _ in range(WORK_QUEUE.qsize()):
        t = threading.Thread(target=worker)
        t.start()
        threads.append(t)

    # Wait for all VMs to be created.
    WORK_QUEUE.join()

    # Signal worker threads to exit.
    for _ in range(len(threads)):
        WORK_QUEUE.put(None)

    # Wait for worker threads to exit.
    for t in threads:
        t.join()

    return 0

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Print information about successful GCE operations. (#605)<commit_after>#!/usr/bin/env python3
#
# Copyright 2018 The Bazel Authors. All rights reserved.
#
# 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.

import argparse
import os
import queue
import sys
import threading
import urllib.error
import urllib.request
import yaml

import gcloud

CONFIG_URL = "https://raw.githubusercontent.com/bazelbuild/continuous-integration/master/buildkite/instances.yml"
LOCAL_CONFIG_FILE_NAME = "instances.yml"

WORK_QUEUE = queue.Queue()

def worker():
    while True:
        item = WORK_QUEUE.get()
        if not item:
            break
        try:
            # We take a few keys out of the config item. The rest is passed
            # as-is to create_instance_template() and thus to the gcloud
            # command line tool.
            count = item.pop("count")
            instance_group_name = item.pop("name")
            project = item.pop("project")
            zone = item.pop("zone", None)
            region = item.pop("region", None)

            if not project:
                raise Exception("Invalid instance config, no project name set")

            if not zone and not region:
                raise Exception(
                    "Invalid instance config, either zone or region must be specified"
                )

            template_name = instance_group_name + "-template"

            if zone is not None:
                if (gcloud.delete_instance_group(
                        instance_group_name, project=project,
                        zone=zone).returncode == 0):
                    print("Deleted existing instance group: {}".format(
                        instance_group_name))
            elif region is not None:
                if (gcloud.delete_instance_group(
                        instance_group_name, project=project,
                        region=region).returncode == 0):
                    print("Deleted existing instance group: {}".format(
                        instance_group_name))

            if gcloud.delete_instance_template(
                    template_name, project=project).returncode == 0:
                print("Deleted existing VM template: {}".format(template_name))

            print("Creating a new template and then a new instance group for project {}".format(project))

            gcloud.create_instance_template(
                template_name, project=project, **item)
            print("Created instance template {}}".format(template_name))

            kwargs = {
                "project": project,
                "base_instance_name": instance_group_name,
                "size": count,
                "template": template_name,
            }
            if zone is not None:
                kwargs["zone"] = zone
            elif region is not None:
                kwargs["region"] = region
            gcloud.create_instance_group(instance_group_name, **kwargs)
            print("Created instance group {}}".format(instance_group_name))
        finally:
            WORK_QUEUE.task_done()

def read_config_file(use_local_config):
    content = None
    if use_local_config:
        path = os.path.join(os.getcwd(), LOCAL_CONFIG_FILE_NAME)
        with open(path, "rb") as fd:
            content = fd.read().decode("utf-8")
    else:
        with urllib.request.urlopen(CONFIG_URL) as resp:
            content = resp.read().decode("utf-8")
    return yaml.safe_load(content)

def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]

    parser = argparse.ArgumentParser(description="Bazel CI Instance Creation")
    parser.add_argument(
        "names",
        type=str,
        nargs="*",
        help="List of instance (group) names that should be created. "
        'These values must correspond to "name" entries in the '
        'Yaml configuration, e.g. "bk-pipeline-ubuntu1804-java8".',
    )
    parser.add_argument(
        "--local_config",
        action="store_true",
        help="Whether to read the configuration from CWD/%s" %
        LOCAL_CONFIG_FILE_NAME,
    )

    args = parser.parse_args(argv)
    config = read_config_file(args.local_config)

    # Verify names passed on the command-line.
    valid_names = [item["name"] for item in config["instance_groups"]]
    for name in args.names:
        if name not in valid_names:
            print("Unknown instance name: {}!".format(name))
            print("\nValid instance names are: {}".format(
                " ".join(valid_names)))
            return 1
    if not args.names:
        parser.print_help()
        print("\nValid instance names are: {}".format(" ".join(valid_names)))
        return 1

    # Put VM creation instructions into the work queue.
    for instance in config["instance_groups"]:
        if instance["name"] not in args.names:
            continue
        WORK_QUEUE.put({**config["default_vm"], **instance})

    # Spawn worker threads that will create the VMs.
    threads = []
    for _ in range(WORK_QUEUE.qsize()):
        t = threading.Thread(target=worker)
        t.start()
        threads.append(t)

    # Wait for all VMs to be created.
    WORK_QUEUE.join()

    # Signal worker threads to exit.
    for _ in range(len(threads)):
        WORK_QUEUE.put(None)

    # Wait for worker threads to exit.
    for t in threads:
        t.join()

    return 0

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


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

# Hive Netius System
# Copyright (C) 2008-2012 Hive Solutions Lda.
#
# This file is part of Hive Netius System.
#
# Hive Netius System 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.
#
# Hive Netius System 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 Hive Netius System. If not, see <http://www.gnu.org/licenses/>.

__author__ = "João Magalhães joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2012 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "GNU General Public License (GPL), Version 3"
""" The license for the module """

import struct

import netius.common
import netius.clients

GRANTED = 0x5a
REJECTED = 0x5b
FAILED_CLIENT = 0x5c
FAILED_AUTH = 0x5d

GRANTED_EXTRA = 0x00

class SOCKSConnection(netius.Connection):

    def __init__(self, owner, socket, address, ssl = False):
        netius.Connection.__init__(self, owner, socket, address, ssl = ssl)
        self.parser = netius.common.SOCKSParser(self)

        self.parser.bind("on_data", self.on_data)
        self.parser.bind("on_auth", self.on_auth)

    def send_response(self, status = GRANTED):
        data = struct.pack("!BBHI", 0, status, 0, 0)
        self.send(data)

    def send_response_extra(self, status = GRANTED_EXTRA):
        version = self.parser.version
        type = self.parser.type
        port = self.parser.port
        address = self.parser.get_address()
        format = "!BBBB%dsH" % len(address)
        data = struct.pack(format, version, status, 0, type, address, port)
        self.send(data)

    def send_auth(self, version = None, method = 0x00):
        version = version or self.parser.version
        data = struct.pack("!BB", version, method)
        self.send(data)

    def get_version(self):
        return self.parser.version

    def parse(self, data):
        return self.parser.parse(data)

    def on_data(self):
        self.owner.on_data_socks(self, self.parser)

    def on_auth(self):
        self.owner.on_auth_socks(self, self.parser)

class SOCKSServer(netius.StreamServer):

    def __init__(self, rules = {}, name = None, handler = None, *args, **kwargs):
        netius.StreamServer.__init__(
            self,
            name = name,
            handler = handler,
            *args,
            **kwargs
        )
        self.rules = rules
        self.conn_map = {}

        self.raw_client = netius.clients.RawClient(
            thread = False,
            *args,
            **kwargs
        )
        self.raw_client.bind("connect", self._on_raw_connect)
        self.raw_client.bind("data", self._on_raw_data)
        self.raw_client.bind("close", self._on_raw_close)

        self.container = netius.Container(*args, **kwargs)
        self.container.add_base(self)
        self.container.add_base(self.raw_client)

    def start(self):
        self.raw_client.load()
        self.container.start()

    def stop(self):
        self.container.stop()

    def on_data(self, connection, data):
        netius.StreamServer.on_data(self, connection, data)

        if hasattr(connection, "tunnel_c"): connection.tunnel_c.send(data)
        else: connection.parse(data)

    def on_data_socks(self, connection, parser):
        host = parser.get_host()
        port = parser.port

        _connection = self.raw_client.connect(host, port)
        connection.tunnel_c = _connection
        self.conn_map[_connection] = connection

    def on_auth_socks(self, connection, parser):
        auth_methods = parser.auth_methods

        if not 0 in auth_methods:
            raise netius.ParserError("Authentication is not supported")

        connection.send_auth(method = 0)

    def on_connection_d(self, connection):
        netius.StreamServer.on_connection_d(self, connection)

        if hasattr(connection, "tunnel_c"): connection.tunnel_c.close()

    def new_connection(self, socket, address, ssl = False):
        return SOCKSConnection(self, socket, address, ssl = ssl)

    def _on_raw_connect(self, client, _connection):
        connection = self.conn_map[_connection]
        version = connection.get_version()
        if version == 0x04: connection.send_response(status = GRANTED)
        elif version == 0x05: connection.send_response_extra(status = GRANTED_EXTRA)

    def _on_raw_data(self, client, _connection, data):
        connection = self.conn_map[_connection]
        connection.send(data)

    def _on_raw_close(self, client, _connection):
        connection = self.conn_map[_connection]
        connection.close(flush = True)
        del self.conn_map[_connection]

if __name__ == "__main__":
    import logging
    server = SOCKSServer(level = logging.INFO)
    server.serve(env = True)
<commit_msg>new socks server initializer<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# Hive Netius System
# Copyright (C) 2008-2012 Hive Solutions Lda.
#
# This file is part of Hive Netius System.
#
# Hive Netius System 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.
#
# Hive Netius System 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 Hive Netius System. If not, see <http://www.gnu.org/licenses/>.

__author__ = "João Magalhães joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2012 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "GNU General Public License (GPL), Version 3"
""" The license for the module """

import struct

import netius.common
import netius.clients

GRANTED = 0x5a
REJECTED = 0x5b
FAILED_CLIENT = 0x5c
FAILED_AUTH = 0x5d

GRANTED_EXTRA = 0x00

class SOCKSConnection(netius.Connection):

    def __init__(self, owner, socket, address, ssl = False):
        netius.Connection.__init__(self, owner, socket, address, ssl = ssl)
        self.parser = netius.common.SOCKSParser(self)

        self.parser.bind("on_data", self.on_data)
        self.parser.bind("on_auth", self.on_auth)

    def send_response(self, status = GRANTED):
        data = struct.pack("!BBHI", 0, status, 0, 0)
        self.send(data)

    def send_response_extra(self, status = GRANTED_EXTRA):
        version = self.parser.version
        type = self.parser.type
        port = self.parser.port
        address = self.parser.get_address()
        format = "!BBBB%dsH" % len(address)
        data = struct.pack(format, version, status, 0, type, address, port)
        self.send(data)

    def send_auth(self, version = None, method = 0x00):
        version = version or self.parser.version
        data = struct.pack("!BB", version, method)
        self.send(data)

    def get_version(self):
        return self.parser.version

    def parse(self, data):
        return self.parser.parse(data)

    def on_data(self):
        self.owner.on_data_socks(self, self.parser)

    def on_auth(self):
        self.owner.on_auth_socks(self, self.parser)

class SOCKSServer(netius.StreamServer):

    def __init__(self, rules = {}, *args, **kwargs):
        netius.StreamServer.__init__(self, *args, **kwargs)
        self.rules = rules
        self.conn_map = {}

        self.raw_client = netius.clients.RawClient(
            thread = False,
            *args,
            **kwargs
        )
        self.raw_client.bind("connect", self._on_raw_connect)
        self.raw_client.bind("data", self._on_raw_data)
        self.raw_client.bind("close", self._on_raw_close)

        self.container = netius.Container(*args, **kwargs)
        self.container.add_base(self)
        self.container.add_base(self.raw_client)

    def start(self):
        self.raw_client.load()
        self.container.start()

    def stop(self):
        self.container.stop()

    def on_data(self, connection, data):
        netius.StreamServer.on_data(self, connection, data)

        if hasattr(connection, "tunnel_c"): connection.tunnel_c.send(data)
        else: connection.parse(data)

    def on_data_socks(self, connection, parser):
        host = parser.get_host()
        port = parser.port

        _connection = self.raw_client.connect(host, port)
        connection.tunnel_c = _connection
        self.conn_map[_connection] = connection

    def on_auth_socks(self, connection, parser):
        auth_methods = parser.auth_methods

        if not 0 in auth_methods:
            raise netius.ParserError("Authentication is not supported")

        connection.send_auth(method = 0)

    def on_connection_d(self, connection):
        netius.StreamServer.on_connection_d(self, connection)

        if hasattr(connection, "tunnel_c"): connection.tunnel_c.close()

    def new_connection(self, socket, address, ssl = False):
        return SOCKSConnection(self, socket, address, ssl = ssl)

    def _on_raw_connect(self, client, _connection):
        connection = self.conn_map[_connection]
        version = connection.get_version()
        if version == 0x04: connection.send_response(status = GRANTED)
        elif version == 0x05: connection.send_response_extra(status = GRANTED_EXTRA)

    def _on_raw_data(self, client, _connection, data):
        connection = self.conn_map[_connection]
        connection.send(data)

    def _on_raw_close(self, client, _connection):
        connection = self.conn_map[_connection]
        connection.close(flush = True)
        del self.conn_map[_connection]

if __name__ == "__main__":
    import logging
    server = SOCKSServer(level = logging.INFO)
    server.serve(env = True)
<|endoftext|>


<commit_before>import timeit

def autotimeit(stmt, setup="pass", repeat=3, mintime=0.2):
    timer = timeit.Timer(stmt, setup)
    number, time1 = autoscaler(timer, mintime)
    time2 = timer.repeat(repeat=repeat - 1, number=number)
    return min(time2 + [time1]) / number

def autoscaler(timer, mintime):
    number = 1
    for i in range(12):
        time = timer.timeit(number)
        if time > mintime:
            return number, time
        number *= 10
    raise RuntimeError("function is too fast to test")
<commit_msg>CLN: add typing to autotimeit.py<commit_after>import timeit
from typing import Tuple

def autotimeit(
    stmt: str, setup: str = "pass", repeat: int = 3, mintime: float = 0.2
) -> float:
    timer = timeit.Timer(stmt, setup)
    number, time1 = autoscaler(timer, mintime)
    time2 = timer.repeat(repeat=repeat - 1, number=number)
    return min(time2 + [time1]) / number

def autoscaler(timer: timeit.Timer, mintime: float) -> Tuple[int, float]:
    number = 1
    for _ in range(12):
        time = timer.timeit(number)
        if time > mintime:
            return number, time
        number *= 10
    raise RuntimeError("function is too fast to test")
<|endoftext|>