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

'''
The MIT License (MIT)

Copyright (c) 2015 Benjamin Xiao

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
'''

from __future__ import print_function

import argparse
import datetime
import os
import shlex
import signal
import subprocess
import time
import traceback

class NetemInstance:
    '''Wrapper around netem module and tc command.'''

    def __init__(self, nic, inbound, include, exclude):
        self.inbound = inbound
        self.include = include
        self.exclude = exclude

        if self.inbound:
            # Create virtual ifb device to do inbound impairment on
            self.real_nic = nic
            self.nic = 'ifb1'
        else:
            # Do normal outbound impairment so no virtual device necessary
            self.nic = nic

    @staticmethod
    def _call(command):
        '''Run command.'''
        subprocess.call(shlex.split(command))

    @staticmethod
    def _check_call(command):
        '''Run command, raising CalledProcessError if it fails.'''
        subprocess.check_call(shlex.split(command))

    @staticmethod
    def _generate_filters(filter_list):
        filter_strings = []
        filter_strings_ipv6 = []
        for tcfilter in filter_list:
            filter_tokens = tcfilter.split(',')
            try:
                filter_string = ''
                filter_string_ipv6 = ''
                for token in filter_tokens:
                    token_split = token.split('=')
                    key = token_split[0]
                    value = token_split[1]
                    # Check for ipv6 addresses and add them to the appropriate
                    # filter string
                    if key == 'src' or key == 'dst':
                        if '::' in value:
                            filter_string_ipv6 += 'match ip6 {0} {1} '.format(
                                key, value)
                        else:
                            filter_string += 'match ip {0} {1} '.format(
                                key, value)
                    else:
                        filter_string += 'match ip {0} {1} '.format(key, value)
                        filter_string_ipv6 += 'match ip6 {0} {1} '.format(
                            key, value)
                    if key == 'sport' or key == 'dport':
                        filter_string += '0xffff '
                        filter_string_ipv6 += '0xffff '
            except IndexError:
                print('Invalid filter parameters')

            if filter_string:
                filter_strings.append(filter_string)
            if filter_string_ipv6:
                filter_strings_ipv6.append(filter_string_ipv6)

        return filter_strings, filter_strings_ipv6

    def initialize(self):
        '''Set up traffic control.'''
        if self.inbound:
            self._check_call('modprobe ifb')
            self._check_call('ip link set dev {0} up'.format(self.nic))
            # Delete ingress device before trying to add
            self._call('tc qdisc del dev {0} ingress'.format(self.real_nic))
            # Add ingress device
            self._check_call(
                'tc qdisc replace dev {0} ingress'.format(self.real_nic))
            # Add filter to redirect ingress to virtual ifb device
            self._check_call(
                'tc filter replace dev {0} parent ffff: protocol ip prio 1 '
                'u32 match u32 0 0 flowid 1:1 action mirred egress redirect '
                'dev {1}'.format(self.real_nic, self.nic))

        # Delete network impairments from any previous runs of this script
        self._call('tc qdisc del root dev {0}'.format(self.nic))

        # Create prio qdisc so we can redirect some traffic to be unimpaired
        self._check_call(
            'tc qdisc add dev {0} root handle 1: prio'.format(self.nic))

        # Apply selective impairment based on include and exclude parameters
        # Work around broken default append behavior. Add default 'src=0/0' if
        # include list is empty
        if len(self.include) == 0:
            self.include.append('src=0/0')
            self.include.append('src=::/0')

        print('Including the following for network impairment:')
        include_filters, include_filters_ipv6 = self._generate_filters(
            self.include)
        for filter_string in include_filters:
            include_filter = 'tc filter add dev {0} protocol ip parent 1:0 ' \
                'prio 3 u32 {1}flowid 1:3'.format(self.nic, filter_string)
            print(include_filter)
            self._check_call(include_filter)

        for filter_string_ipv6 in include_filters_ipv6:
            include_filter_ipv6 = 'tc filter add dev {0} protocol ipv6 ' \
                'parent 1:0 prio 4 u32 {1}flowid 1:3'.format(
                    self.nic, filter_string_ipv6)
            print(include_filter_ipv6)
            self._check_call(include_filter_ipv6)

        print('Excluding the following from network impairment:')
        exclude_filters, exclude_filters_ipv6 = self._generate_filters(
            self.exclude)
        for filter_string in exclude_filters:
            exclude_filter = 'tc filter add dev {0} protocol ip parent 1:0 ' \
                'prio 1 u32 {1}flowid 1:2'.format(self.nic, filter_string)
            print(exclude_filter)
            self._check_call(exclude_filter)

        for filter_string_ipv6 in exclude_filters_ipv6:
            exclude_filter_ipv6 = 'tc filter add dev {0} protocol ipv6 ' \
                'parent 1:0 prio 2 u32 {1}flowid 1:2'.format(
                    self.nic, filter_string_ipv6)
            print(exclude_filter_ipv6)
            self._check_call(exclude_filter_ipv6)

    def netem(
            self,
            loss_ratio=0,
            loss_corr=0,
            dup_ratio=0,
            delay=0,
            jitter=0,
            delay_jitter_corr=0,
            reorder_ratio=0,
            reorder_corr=0,
            toggle=[1000000]):
        '''Enable packet loss.'''
        self._check_call(
            'tc qdisc add dev {0} parent 1:3 handle 30: netem'.format(
                self.nic))
        while len(toggle) != 0:
            impair_cmd = 'tc qdisc change dev {0} parent 1:3 handle 30: ' \
                'netem loss {1}% {2}% duplicate {3}% delay {4}ms {5}ms {6}% ' \
                'reorder {7}% {8}%'.format(
                    self.nic, loss_ratio, loss_corr, dup_ratio, delay, jitter,
                    delay_jitter_corr, reorder_ratio, reorder_corr)
            print('Setting network impairment:')
            print(impair_cmd)
            # Set network impairment
            self._check_call(impair_cmd)
            print(
                'Impairment timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))
            if len(toggle) == 0:
                return
            self._check_call(
                'tc qdisc change dev {0} parent 1:3 handle 30: netem'.format(
                    self.nic))
            print(
                'Impairment stopped timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))

    def rate(self, limit, buffer, latency, toggle):
        '''Enable packet reorder.'''
        self._check_call(
            'tc qdisc add dev {0} parent 1:3 handle 30: tbf rate 1000mbit '
            'buffer {1} latency {2}ms'.format(self.nic, buffer, latency))
        while len(toggle) != 0:
            impair_cmd = 'tc qdisc change dev {0} parent 1:3 handle 30: tbf ' \
                'rate {1}kbit buffer {2} latency {3}ms'.format(
                    self.nic, limit, buffer, latency)
            print('Setting network impairment:')
            print(impair_cmd)
            # Set network impairment
            self._check_call(impair_cmd)
            print(
                'Impairment timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))
            if len(toggle) == 0:
                return
            self._check_call(
                'tc qdisc change dev {0} parent 1:3 handle 30: tbf rate '
                '1000mbit buffer {1} latency {2}ms'.format(
                    self.nic, buffer, latency))
            print(
                'Impairment stopped timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))

    def teardown(self):
        '''Reset traffic control rules.'''
        if self.inbound:
            self._call(
                'tc filter del dev {0} parent ffff: protocol ip prio 1'.format(
                    self.real_nic))
            self._call('tc qdisc del dev {0} ingress'.format(self.real_nic))
            self._call('ip link set dev ifb0 down')
        self._call('tc qdisc del root dev {0}'.format(self.nic))
        print('Network impairment teardown complete.')

def init_signals(netem):
    '''Catch signals in order to stop network impairment before exiting.'''

    def signal_action(signum, frame):
        '''To be executed upon exit signal.'''
        print()
        netem.teardown()
        # Print blank line before quitting to deal with some crappy
        # terminal behavior
        print()
        exit(5)

    # Catch SIGINT and SIGTERM so that we can clean up
    for sig in [signal.SIGINT, signal.SIGTERM]:
        signal.signal(sig, signal_action)

    print('Network impairment starting.',
          'Press Ctrl-C to restore normal behavior and quit.')

def main():
    '''Start application.'''
    args = parse_args()

    if os.geteuid() != 0:
        print('You need root permissions to enable impairment!',
              'Please run with sudo or as root.')
        exit(1)

    try:
        netem = NetemInstance(
            args.nic, args.inbound, args.include, args.exclude)
        netem.initialize()

        init_signals(netem)

        # Do impairment
        if args.subparser_name == 'netem':
            netem.netem(
                args.loss_ratio,
                args.loss_corr,
                args.dup_ratio,
                args.delay,
                args.jitter,
                args.delay_jitter_corr,
                args.reorder_ratio,
                args.reorder_corr,
                args.toggle)
        elif args.subparser_name == 'rate':
            netem.rate(args.limit, args.buffer, args.latency, args.toggle)

        # Shutdown cleanly
        netem.teardown()
    except subprocess.CalledProcessError:
        traceback.print_exc()
        netem.teardown()
        exit(5)

def parse_args():
    '''Parse command-line arguments.'''
    argparser = argparse.ArgumentParser(
        description='Network Impairment Test Tool')

    argparser.add_argument(
        '-n', '--nic',
        choices=os.listdir('/sys/class/net'),
        metavar='INTERFACE',
        required=True,
        help='name of the network interface to be impaired')
    argparser.add_argument(
        '--inbound',
        action='store_true',
        help='do inbound impairment on the interface instead of outbound')
    argparser.add_argument(
        '--include',
        action='append',
        default=[],
        help='ip addresses and/or ports to include in network '
        'impairment (example: --include src=ip,sport=portnum '
        '--include dst=ip,dport=portnum)')
    argparser.add_argument(
        '--exclude',
        action='append',
        default=['dport=22', 'sport=22'],
        help='ip addresses and/or ports to exclude from network '
        'impairment (example: --exclude src=ip,sport=portnum '
        '--exclude dst=ip,dport=portnum)')

    subparsers = argparser.add_subparsers(
        title='impairments',
        dest='subparser_name',
        description='specify which impairment to enable',
        help='valid impairments')

    # loss args
    netem_args = subparsers.add_parser('netem', help='enable packet loss')
    netem_args.add_argument(
        '--loss_ratio',
        type=int,
        default=0,
        help='specify percentage of packets that will be lost')
    netem_args.add_argument(
        '--loss_corr',
        type=int,
        default=0,
        help='specify a correlation factor for the random packet loss')
    # dup args
    netem_args.add_argument(
        '--dup_ratio',
        type=int,
        default=0,
        help='specify percentage of packets that will be duplicated')
    # delay/jitter args
    netem_args.add_argument(
        '--delay',
        type=int,
        default=0,
        help='specify an overall delay for each packet')
    netem_args.add_argument(
        '--jitter',
        type=int,
        default=0,
        help='specify amount of jitter in milliseconds')
    netem_args.add_argument(
        '--delay_jitter_corr',
        type=int,
        default=0,
        help='specify a correlation factor for the random jitter')
    # reorder args
    netem_args.add_argument(
        '--reorder_ratio',
        type=int,
        default=0,
        help='specify percentage of packets that will be reordered')
    netem_args.add_argument(
        '--reorder_corr',
        type=int,
        default=0,
        help='specify a correlation factor for the random reordering')
    # toggle parameter
    netem_args.add_argument(
        '--toggle',
        nargs='+',
        type=int,
        default=[1000000],
        help='toggles impairment on and off on specific intervals (example: '
        '--toggle 6 3 5 1 will enable impairment for 6 seconds, turn it off '
        'for 3, turn it on for 5, and turn it off for 1')

    # rate limit args
    rate_args = subparsers.add_parser('rate', help='enable packet reorder')
    rate_args.add_argument(
        '--limit',
        type=int,
        default=0,
        help='specify rate limit in kb')
    rate_args.add_argument(
        '--buffer',
        type=int,
        default=2000,
        help='specify how many tokens in terms of bytes should be available')
    rate_args.add_argument(
        '--latency',
        type=int,
        default=20,
        help='specify the maximum time packets can stay in the '
        'queue before getting dropped')
    rate_args.add_argument(
        '--toggle',
        nargs='+',
        type=int,
        default=[1000000],
        help='toggles impairment on and off on specific intervals (example: '
        '--toggle 6 3 5 1 will enable impairment for 6 seconds, turn it off '
        'for 3, turn it on for 5, and turn it off for 1')

    return argparser.parse_args()

if __name__ == '__main__':
    main()
<commit_msg>Don't use a mutable object (in this case [1000000]) as a default argument<commit_after>#!/usr/bin/env python

'''
The MIT License (MIT)

Copyright (c) 2015 Benjamin Xiao

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
'''

from __future__ import print_function

import argparse
import datetime
import os
import shlex
import signal
import subprocess
import time
import traceback

class NetemInstance:
    '''Wrapper around netem module and tc command.'''

    def __init__(self, nic, inbound, include, exclude):
        self.inbound = inbound
        self.include = include
        self.exclude = exclude

        if self.inbound:
            # Create virtual ifb device to do inbound impairment on
            self.real_nic = nic
            self.nic = 'ifb1'
        else:
            # Do normal outbound impairment so no virtual device necessary
            self.nic = nic

    @staticmethod
    def _call(command):
        '''Run command.'''
        subprocess.call(shlex.split(command))

    @staticmethod
    def _check_call(command):
        '''Run command, raising CalledProcessError if it fails.'''
        subprocess.check_call(shlex.split(command))

    @staticmethod
    def _generate_filters(filter_list):
        filter_strings = []
        filter_strings_ipv6 = []
        for tcfilter in filter_list:
            filter_tokens = tcfilter.split(',')
            try:
                filter_string = ''
                filter_string_ipv6 = ''
                for token in filter_tokens:
                    token_split = token.split('=')
                    key = token_split[0]
                    value = token_split[1]
                    # Check for ipv6 addresses and add them to the appropriate
                    # filter string
                    if key == 'src' or key == 'dst':
                        if '::' in value:
                            filter_string_ipv6 += 'match ip6 {0} {1} '.format(
                                key, value)
                        else:
                            filter_string += 'match ip {0} {1} '.format(
                                key, value)
                    else:
                        filter_string += 'match ip {0} {1} '.format(key, value)
                        filter_string_ipv6 += 'match ip6 {0} {1} '.format(
                            key, value)
                    if key == 'sport' or key == 'dport':
                        filter_string += '0xffff '
                        filter_string_ipv6 += '0xffff '
            except IndexError:
                print('Invalid filter parameters')

            if filter_string:
                filter_strings.append(filter_string)
            if filter_string_ipv6:
                filter_strings_ipv6.append(filter_string_ipv6)

        return filter_strings, filter_strings_ipv6

    def initialize(self):
        '''Set up traffic control.'''
        if self.inbound:
            self._check_call('modprobe ifb')
            self._check_call('ip link set dev {0} up'.format(self.nic))
            # Delete ingress device before trying to add
            self._call('tc qdisc del dev {0} ingress'.format(self.real_nic))
            # Add ingress device
            self._check_call(
                'tc qdisc replace dev {0} ingress'.format(self.real_nic))
            # Add filter to redirect ingress to virtual ifb device
            self._check_call(
                'tc filter replace dev {0} parent ffff: protocol ip prio 1 '
                'u32 match u32 0 0 flowid 1:1 action mirred egress redirect '
                'dev {1}'.format(self.real_nic, self.nic))

        # Delete network impairments from any previous runs of this script
        self._call('tc qdisc del root dev {0}'.format(self.nic))

        # Create prio qdisc so we can redirect some traffic to be unimpaired
        self._check_call(
            'tc qdisc add dev {0} root handle 1: prio'.format(self.nic))

        # Apply selective impairment based on include and exclude parameters
        # Work around broken default append behavior. Add default 'src=0/0' if
        # include list is empty
        if len(self.include) == 0:
            self.include.append('src=0/0')
            self.include.append('src=::/0')

        print('Including the following for network impairment:')
        include_filters, include_filters_ipv6 = self._generate_filters(
            self.include)
        for filter_string in include_filters:
            include_filter = 'tc filter add dev {0} protocol ip parent 1:0 ' \
                'prio 3 u32 {1}flowid 1:3'.format(self.nic, filter_string)
            print(include_filter)
            self._check_call(include_filter)

        for filter_string_ipv6 in include_filters_ipv6:
            include_filter_ipv6 = 'tc filter add dev {0} protocol ipv6 ' \
                'parent 1:0 prio 4 u32 {1}flowid 1:3'.format(
                    self.nic, filter_string_ipv6)
            print(include_filter_ipv6)
            self._check_call(include_filter_ipv6)

        print('Excluding the following from network impairment:')
        exclude_filters, exclude_filters_ipv6 = self._generate_filters(
            self.exclude)
        for filter_string in exclude_filters:
            exclude_filter = 'tc filter add dev {0} protocol ip parent 1:0 ' \
                'prio 1 u32 {1}flowid 1:2'.format(self.nic, filter_string)
            print(exclude_filter)
            self._check_call(exclude_filter)

        for filter_string_ipv6 in exclude_filters_ipv6:
            exclude_filter_ipv6 = 'tc filter add dev {0} protocol ipv6 ' \
                'parent 1:0 prio 2 u32 {1}flowid 1:2'.format(
                    self.nic, filter_string_ipv6)
            print(exclude_filter_ipv6)
            self._check_call(exclude_filter_ipv6)

    def netem(
            self,
            loss_ratio=0,
            loss_corr=0,
            dup_ratio=0,
            delay=0,
            jitter=0,
            delay_jitter_corr=0,
            reorder_ratio=0,
            reorder_corr=0,
            toggle=None):
        '''Enable packet loss.'''
        if toggle is None:
            toggle = [1000000]
        self._check_call(
            'tc qdisc add dev {0} parent 1:3 handle 30: netem'.format(
                self.nic))
        while len(toggle) != 0:
            impair_cmd = 'tc qdisc change dev {0} parent 1:3 handle 30: ' \
                'netem loss {1}% {2}% duplicate {3}% delay {4}ms {5}ms {6}% ' \
                'reorder {7}% {8}%'.format(
                    self.nic, loss_ratio, loss_corr, dup_ratio, delay, jitter,
                    delay_jitter_corr, reorder_ratio, reorder_corr)
            print('Setting network impairment:')
            print(impair_cmd)
            # Set network impairment
            self._check_call(impair_cmd)
            print(
                'Impairment timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))
            if len(toggle) == 0:
                return
            self._check_call(
                'tc qdisc change dev {0} parent 1:3 handle 30: netem'.format(
                    self.nic))
            print(
                'Impairment stopped timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))

    def rate(self, limit, buffer, latency, toggle):
        '''Enable packet reorder.'''
        self._check_call(
            'tc qdisc add dev {0} parent 1:3 handle 30: tbf rate 1000mbit '
            'buffer {1} latency {2}ms'.format(self.nic, buffer, latency))
        while len(toggle) != 0:
            impair_cmd = 'tc qdisc change dev {0} parent 1:3 handle 30: tbf ' \
                'rate {1}kbit buffer {2} latency {3}ms'.format(
                    self.nic, limit, buffer, latency)
            print('Setting network impairment:')
            print(impair_cmd)
            # Set network impairment
            self._check_call(impair_cmd)
            print(
                'Impairment timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))
            if len(toggle) == 0:
                return
            self._check_call(
                'tc qdisc change dev {0} parent 1:3 handle 30: tbf rate '
                '1000mbit buffer {1} latency {2}ms'.format(
                    self.nic, buffer, latency))
            print(
                'Impairment stopped timestamp: {0}'.format(
                    datetime.datetime.today()))
            time.sleep(toggle.pop(0))

    def teardown(self):
        '''Reset traffic control rules.'''
        if self.inbound:
            self._call(
                'tc filter del dev {0} parent ffff: protocol ip prio 1'.format(
                    self.real_nic))
            self._call('tc qdisc del dev {0} ingress'.format(self.real_nic))
            self._call('ip link set dev ifb0 down')
        self._call('tc qdisc del root dev {0}'.format(self.nic))
        print('Network impairment teardown complete.')

def init_signals(netem):
    '''Catch signals in order to stop network impairment before exiting.'''

    def signal_action(signum, frame):
        '''To be executed upon exit signal.'''
        print()
        netem.teardown()
        # Print blank line before quitting to deal with some crappy
        # terminal behavior
        print()
        exit(5)

    # Catch SIGINT and SIGTERM so that we can clean up
    for sig in [signal.SIGINT, signal.SIGTERM]:
        signal.signal(sig, signal_action)

    print('Network impairment starting.',
          'Press Ctrl-C to restore normal behavior and quit.')

def main():
    '''Start application.'''
    args = parse_args()

    if os.geteuid() != 0:
        print('You need root permissions to enable impairment!',
              'Please run with sudo or as root.')
        exit(1)

    try:
        netem = NetemInstance(
            args.nic, args.inbound, args.include, args.exclude)
        netem.initialize()

        init_signals(netem)

        # Do impairment
        if args.subparser_name == 'netem':
            netem.netem(
                args.loss_ratio,
                args.loss_corr,
                args.dup_ratio,
                args.delay,
                args.jitter,
                args.delay_jitter_corr,
                args.reorder_ratio,
                args.reorder_corr,
                args.toggle)
        elif args.subparser_name == 'rate':
            netem.rate(args.limit, args.buffer, args.latency, args.toggle)

        # Shutdown cleanly
        netem.teardown()
    except subprocess.CalledProcessError:
        traceback.print_exc()
        netem.teardown()
        exit(5)

def parse_args():
    '''Parse command-line arguments.'''
    argparser = argparse.ArgumentParser(
        description='Network Impairment Test Tool')

    argparser.add_argument(
        '-n', '--nic',
        choices=os.listdir('/sys/class/net'),
        metavar='INTERFACE',
        required=True,
        help='name of the network interface to be impaired')
    argparser.add_argument(
        '--inbound',
        action='store_true',
        help='do inbound impairment on the interface instead of outbound')
    argparser.add_argument(
        '--include',
        action='append',
        default=[],
        help='ip addresses and/or ports to include in network '
        'impairment (example: --include src=ip,sport=portnum '
        '--include dst=ip,dport=portnum)')
    argparser.add_argument(
        '--exclude',
        action='append',
        default=['dport=22', 'sport=22'],
        help='ip addresses and/or ports to exclude from network '
        'impairment (example: --exclude src=ip,sport=portnum '
        '--exclude dst=ip,dport=portnum)')

    subparsers = argparser.add_subparsers(
        title='impairments',
        dest='subparser_name',
        description='specify which impairment to enable',
        help='valid impairments')

    # loss args
    netem_args = subparsers.add_parser('netem', help='enable packet loss')
    netem_args.add_argument(
        '--loss_ratio',
        type=int,
        default=0,
        help='specify percentage of packets that will be lost')
    netem_args.add_argument(
        '--loss_corr',
        type=int,
        default=0,
        help='specify a correlation factor for the random packet loss')
    # dup args
    netem_args.add_argument(
        '--dup_ratio',
        type=int,
        default=0,
        help='specify percentage of packets that will be duplicated')
    # delay/jitter args
    netem_args.add_argument(
        '--delay',
        type=int,
        default=0,
        help='specify an overall delay for each packet')
    netem_args.add_argument(
        '--jitter',
        type=int,
        default=0,
        help='specify amount of jitter in milliseconds')
    netem_args.add_argument(
        '--delay_jitter_corr',
        type=int,
        default=0,
        help='specify a correlation factor for the random jitter')
    # reorder args
    netem_args.add_argument(
        '--reorder_ratio',
        type=int,
        default=0,
        help='specify percentage of packets that will be reordered')
    netem_args.add_argument(
        '--reorder_corr',
        type=int,
        default=0,
        help='specify a correlation factor for the random reordering')
    # toggle parameter
    netem_args.add_argument(
        '--toggle',
        nargs='+',
        type=int,
        default=[1000000],
        help='toggles impairment on and off on specific intervals (example: '
        '--toggle 6 3 5 1 will enable impairment for 6 seconds, turn it off '
        'for 3, turn it on for 5, and turn it off for 1')

    # rate limit args
    rate_args = subparsers.add_parser('rate', help='enable packet reorder')
    rate_args.add_argument(
        '--limit',
        type=int,
        default=0,
        help='specify rate limit in kb')
    rate_args.add_argument(
        '--buffer',
        type=int,
        default=2000,
        help='specify how many tokens in terms of bytes should be available')
    rate_args.add_argument(
        '--latency',
        type=int,
        default=20,
        help='specify the maximum time packets can stay in the '
        'queue before getting dropped')
    rate_args.add_argument(
        '--toggle',
        nargs='+',
        type=int,
        default=[1000000],
        help='toggles impairment on and off on specific intervals (example: '
        '--toggle 6 3 5 1 will enable impairment for 6 seconds, turn it off '
        'for 3, turn it on for 5, and turn it off for 1')

    return argparser.parse_args()

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># Copyright 2019 Google Inc.
#
# 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.

# This file is autogenerated by copybara, please do not edit.

buildtools_version = "e8d502966b089e9a4d1cc960aca8a982e97ce601"
buildtools_sha256 = "faa51477cfd8315d5df037fe692921adc533d793903551009942400ba36dbdf3"
<commit_msg>Update dependency bazelbuild/buildtools to latest version<commit_after># Copyright 2019 Google Inc.
#
# 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.

# This file is autogenerated by copybara, please do not edit.

buildtools_version = "b08a19329f40f76175ee7353b1d22782af80526f"
buildtools_sha256 = "5d27bf4972295bb15eb544ea249a83cdbd7b0c1af086b5d34e7ae4a12381895d"
<|endoftext|>


<commit_before><commit_msg>Branch merge<commit_after><|endoftext|>


<commit_before># -*- coding: utf8 -*-

# 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 logging
import requests
import rdflib
import SPARQLWrapper

def print_debug_info():
    logging.debug("requests=" % requests.__version__)
    logging.debug("rdflib=" % rdflib.__version__)
    logging.debug("SPARQLWrapper=" % SPARQLWrapper.__version__)

print_debug_info()

<commit_msg>fixed<commit_after># -*- coding: utf8 -*-

# 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 logging
import requests
import rdflib
import SPARQLWrapper

def print_debug_info():
    logging.debug("requests=%s" % requests.__version__)
    logging.debug("rdflib=%s" % rdflib.__version__)
    logging.debug("SPARQLWrapper=%s" % SPARQLWrapper.__version__)

print_debug_info()

<|endoftext|>


<commit_before>from datetime import datetime
from requests.exceptions import ConnectionError, ReadTimeout, SSLError
from requests.packages.urllib3.exceptions import ReadTimeoutError, ProtocolError
from .exceptions import *

import requests
import hashlib
import socket
import ssl
import time

BASE_URL = 'http://api.stats.com/v1'
REST_TIMEOUT = None
API_SECRET = None

class StatsAPI(object):

    """Access REST API or Streaming API resources.
    :param key: Stats application key
    :param secret: Stats application secret
    :param accept: Stats accept type (json, xml)
    """

    api_key = None
    api_secret = None
    api_accept = None

    def __init__(
            self,
            key=None,
            secret=None,
            accept='json'):
        """Initialize with your application credentials"""

        if not all([key, secret]):
            raise Exception('Missing authentication parameter')

        self.api_key = key
        self.api_secret = secret
        self.api_accept = accept

    def _prepare_url(self, domain, path, params=None):
        auth = self.generate_auth()
        url = '{}{}?api_key={}&sig={}'.format(domain, path, auth['api_key'], auth['sig'])
        return url

    def _get_endpoint(self, resource, params=None):
        """Substitute any parameters in the resource path with :PARAM."""
        if ':' in resource:
            if params is not None:
                for k in params:
                    try:
                        url_param = ':%s' % k
                        resource = resource.replace(url_param, str(params[k]))
                    except:
                        pass
                return resource, resource

            else:
                return resource, resource

        else:
            return resource, resource

    def request(self, resource, rest_params=None, params=None, files=None):
        """Request a Stats REST API .
        :param resource: A valid Stats endpoint
        :param params: Dictionary with endpoint parameters or None (default)
        :returns: StatsResponse
        :raises: StatsConnectionError
        """
        resource, endpoint = self._get_endpoint(resource, rest_params)
        # if endpoint not in ENDPOINTS:
        #    raise Exception('Endpoint "%s" unsupported' % endpoint)

        # payload = {'key1': 'value1', 'key2': 'value2'}
        # r = requests.get("http://httpbin.org/get", params=payload)
        session = requests.Session()
        # session.headers = {'User-Agent': USER_AGENT}
        method = 'GET'
        subdomain = BASE_URL
        url = self._prepare_url(subdomain, resource, params)
        session.stream = False
        timeout = REST_TIMEOUT
        if method == 'POST':
            data = params
            params = None
        else:
            data = None
        try:
            r = session.request(
                method,
                url,
                data=data,
                params=params,
                timeout=timeout,
                files=files,
                # proxies=self.proxies
            )
        except (ConnectionError, ProtocolError, ReadTimeout, ReadTimeoutError,
                SSLError, ssl.SSLError, socket.error) as e:
            raise StatsConnectionError(e)
        return StatsResponse(r)

    def generate_auth(self):
        """Returns the dictionary for the authentication get params
        """
        return {'accept': self.api_accept, 'api_key': self.api_key, 'sig': self.generate_sig()}

    def generate_sig(self):
        """Returns the SHA256 signature for making requests
        """
        timestamp = int(time.time())
        hash_it = '{}{}{}'.format(self.api_key, self.api_secret, timestamp)
        return hashlib.sha256(hash_it).hexdigest()

class StatsResponse(object):

    """Response from either a REST API call.
    :param response: The requests.Response object returned by the API call
    """

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

    @property
    def headers(self):
        """:returns: Dictionary of API response header contents."""
        return self.response.headers

    @property
    def status_code(self):
        """:returns: HTTP response status code."""
        return self.response.status_code

    @property
    def text(self):
        """:returns: Raw API response text."""
        return self.response.text

    def json(self):
        """:returns: response as JSON object."""
        return self.response.json()<commit_msg>return url in response<commit_after>from requests.exceptions import ConnectionError, ReadTimeout, SSLError
from requests.packages.urllib3.exceptions import ReadTimeoutError, ProtocolError
from .exceptions import *

import requests
import hashlib
import socket
import ssl
import time

BASE_URL = 'http://api.stats.com/v1'
REST_TIMEOUT = None
API_SECRET = None

class StatsAPI(object):

    """Access REST API or Streaming API resources.
    :param key: Stats application key
    :param secret: Stats application secret
    :param accept: Stats accept type (json, xml)
    """

    api_key = None
    api_secret = None
    api_accept = None

    def __init__(
            self,
            key=None,
            secret=None,
            accept='json'):
        """Initialize with your application credentials"""

        if not all([key, secret]):
            raise Exception('Missing authentication parameter')

        self.api_key = key
        self.api_secret = secret
        self.api_accept = accept

    def _prepare_url(self, domain, path, params=None):
        auth = self.generate_auth()
        url = '{}{}?api_key={}&sig={}'.format(domain, path, auth['api_key'], auth['sig'])
        return url

    def _get_endpoint(self, resource, params=None):
        """Substitute any parameters in the resource path with :PARAM."""
        if ':' in resource:
            if params is not None:
                for k in params:
                    try:
                        url_param = ':%s' % k
                        resource = resource.replace(url_param, str(params[k]))
                    except:
                        pass
                return resource, resource

            else:
                return resource, resource

        else:
            return resource, resource

    def request(self, resource, rest_params=None, params=None, files=None):
        """Request a Stats REST API .
        :param resource: A valid Stats endpoint
        :param params: Dictionary with endpoint parameters or None (default)
        :returns: StatsResponse
        :raises: StatsConnectionError
        """
        resource, endpoint = self._get_endpoint(resource, rest_params)
        # if endpoint not in ENDPOINTS:
        #    raise Exception('Endpoint "%s" unsupported' % endpoint)

        # payload = {'key1': 'value1', 'key2': 'value2'}
        # r = requests.get("http://httpbin.org/get", params=payload)
        session = requests.Session()
        # session.headers = {'User-Agent': USER_AGENT}
        method = 'GET'
        subdomain = BASE_URL
        url = self._prepare_url(subdomain, resource, params)
        session.stream = False
        timeout = REST_TIMEOUT
        if method == 'POST':
            data = params
            params = None
        else:
            data = None
        try:
            r = session.request(
                method,
                url,
                data=data,
                params=params,
                timeout=timeout,
                files=files,
                # proxies=self.proxies
            )
        except (ConnectionError, ProtocolError, ReadTimeout, ReadTimeoutError,
                SSLError, ssl.SSLError, socket.error) as e:
            raise StatsConnectionError(e)
        return StatsResponse(r)

    def generate_auth(self):
        """Returns the dictionary for the authentication get params
        """
        return {'accept': self.api_accept, 'api_key': self.api_key, 'sig': self.generate_sig()}

    def generate_sig(self):
        """Returns the SHA256 signature for making requests
        """
        timestamp = int(time.time())
        hash_it = '{}{}{}'.format(self.api_key, self.api_secret, timestamp)
        return hashlib.sha256(hash_it).hexdigest()

class StatsResponse(object):

    """Response from either a REST API call.
    :param response: The requests.Response object returned by the API call
    """

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

    @property
    def headers(self):
        """:returns: Dictionary of API response header contents."""
        return self.response.headers

    @property
    def status_code(self):
        """:returns: HTTP response status code."""
        return self.response.status_code

    @property
    def text(self):
        """:returns: Raw API response text."""
        return self.response.text

    @property
    def url(self):
        """:returns: Raw API response url."""
        return self.response.url

    def json(self):
        """:returns: response as JSON object."""
        return self.response.json()<|endoftext|>


<commit_before>import re
import datetime
import logbook

from sqlalchemy import Column
from sqlalchemy import DateTime
from sqlalchemy import Integer
from sqlalchemy import String

from kitten import conf
from kitten.db import Session
from kitten.db import Base
from kitten.client import KittenClient
from kitten.server import RequestError
from kitten.validation import Validator
from jsonschema.exceptions import ValidationError

class Node(Base):
    __tablename__ = 'node'

    id = Column(Integer(), primary_key=True)
    address = Column(String(255))
    created = Column(DateTime, default=datetime.datetime.now)
    last_seen = Column(DateTime, default=datetime.datetime.now)

    client = KittenClient()
    log = logbook.Logger('Node')

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

    def __str__(self):  # pragma: nocover
        return '<Node: {0.address}>'.format(self)

    @staticmethod
    def create(address):
        """
        Create a new node

        Try connecting to the node to see if it is eligible for adding.

        """

        session = Session()

        # If no port is specified, make sure to add the default.
        if not re.search(r':\d+', address):
            address += ':{0}'.format(conf.DEFAULT_PORT)

        con = Node(address)
        q = session.query(Node).filter(Node.address == address)

        if session.query(q.exists()).scalar():
            con.log.error('{0} already exists.'.format(con))
            return

        if con.ping():
            session.add(con)
            session.commit()
            con.log.info('{0} added.'.format(con))
        else:
            con.log.error('Could not connect to {0}.'.format(con))

        session.close()

    @staticmethod
    def list():
        """
        Return a list of all nodes

        """

        session = Session()
        return session.query(Node).all()

    def message(self, data):
        """
        Send a message to the node

        """

        request = {
            'paradigm': 'node',
            # 'address': self.address,
        }
        request.update(data)

        response = self.client.send(self.address, request)
        return response

    def ping(self):
        """
        Send a quick ping heartbeat to the node

        Returns boolean success.

        """

        try:
            response = self.message({'method': 'ping'})
            return response['pong']
        except (KeyError, RequestError, ValidationError):
            self.log.exception('Ping failed')
            return False

    def repr(self):  # pragma: nocover
        return self.__str__()

class NodeValidator(Validator):
    def ping_request(self):
        return {}  # Yay no extra fields

class NodeParadigm(object):
    validator = NodeValidator()

    def setup(self):  # pragma: nocover
        pass

    def ping(self, request):
        """
        Handle a remote ping request

        """

        return {
            'pong': True
        }

def setup_parser(subparsers):
    con = subparsers.add_parser('node', help="List, add or modify nodes.")
    sub = con.add_subparsers(help='Node commands', dest="sub")

    list_ = sub.add_parser('list', help='List nodes (default)')
    list_.add_argument('--filter', type=str)

    add = sub.add_parser('add', help='Add a node')
    add.add_argument('address', type=str)

    remove = sub.add_parser('remove', help='Remove a node')
    remove.add_argument('name', type=str)

    return execute_parser

def execute_parser(ns):
    if not ns.sub or ns.sub == "list":
        src = Node.list()

        # If a filter is specified, apply it to the display name
        if hasattr(ns, 'filter') and ns.filter:
            src = list(filter(lambda x: ns.filter in x.address, src))

        for con in src:
            print(con.repr())

    elif ns.sub == 'add':
        Node.create(ns.address)
<commit_msg>Add NodeParadigm to Node<commit_after>import re
import datetime
import logbook

from sqlalchemy import Column
from sqlalchemy import DateTime
from sqlalchemy import Integer
from sqlalchemy import String

from kitten import conf
from kitten.db import Session
from kitten.db import Base
from kitten.client import KittenClient
from kitten.server import RequestError
from kitten.validation import Validator
from jsonschema.exceptions import ValidationError

class NodeValidator(Validator):
    def ping_request(self):
        return {}  # Yay no extra fields

class NodeParadigm(object):
    validator = NodeValidator()

    def setup(self):  # pragma: nocover
        pass

    def ping(self, request):
        """
        Handle a remote ping request

        """

        return {
            'pong': True
        }

class Node(Base):
    __tablename__ = 'node'
    paradigm = NodeParadigm()

    id = Column(Integer(), primary_key=True)
    address = Column(String(255))
    created = Column(DateTime, default=datetime.datetime.now)
    last_seen = Column(DateTime, default=datetime.datetime.now)

    client = KittenClient()
    log = logbook.Logger('Node')

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

    def __str__(self):  # pragma: nocover
        return '<Node: {0.address}>'.format(self)

    @staticmethod
    def create(address):
        """
        Create a new node

        Try connecting to the node to see if it is eligible for adding.

        """

        session = Session()

        # If no port is specified, make sure to add the default.
        if not re.search(r':\d+', address):
            address += ':{0}'.format(conf.DEFAULT_PORT)

        con = Node(address)
        q = session.query(Node).filter(Node.address == address)

        if session.query(q.exists()).scalar():
            con.log.error('{0} already exists.'.format(con))
            return

        if con.ping():
            session.add(con)
            session.commit()
            con.log.info('{0} added.'.format(con))
        else:
            con.log.error('Could not connect to {0}.'.format(con))

        session.close()

    @staticmethod
    def list():
        """
        Return a list of all nodes

        """

        session = Session()
        return session.query(Node).all()

    def message(self, data):
        """
        Send a message to the node

        """

        request = {
            'paradigm': 'node',
            # 'address': self.address,
        }
        request.update(data)

        response = self.client.send(self.address, request)
        return response

    def ping(self):
        """
        Send a quick ping heartbeat to the node

        Returns boolean success.

        """

        try:
            response = self.message({'method': 'ping'})
            return response['pong']
        except (KeyError, RequestError, ValidationError):
            self.log.exception('Ping failed')
            return False

    def repr(self):  # pragma: nocover
        return self.__str__()

def setup_parser(subparsers):
    con = subparsers.add_parser('node', help="List, add or modify nodes.")
    sub = con.add_subparsers(help='Node commands', dest="sub")

    list_ = sub.add_parser('list', help='List nodes (default)')
    list_.add_argument('--filter', type=str)

    add = sub.add_parser('add', help='Add a node')
    add.add_argument('address', type=str)

    remove = sub.add_parser('remove', help='Remove a node')
    remove.add_argument('name', type=str)

    return execute_parser

def execute_parser(ns):
    if not ns.sub or ns.sub == "list":
        src = Node.list()

        # If a filter is specified, apply it to the display name
        if hasattr(ns, 'filter') and ns.filter:
            src = list(filter(lambda x: ns.filter in x.address, src))

        for con in src:
            print(con.repr())

    elif ns.sub == 'add':
        Node.create(ns.address)
<|endoftext|>


<commit_before>from feedly.feeds.base import UserBaseFeed
from feedly.tasks import follow_many, unfollow_many
from feedly.tasks import fanout_operation
from feedly.tasks import fanout_operation_hi_priority
from feedly.tasks import fanout_operation_low_priority
from feedly.utils import chunks
from feedly.utils.timing import timer
import logging
from feedly.feeds.redis import RedisFeed

logger = logging.getLogger(__name__)

def add_operation(feed, activities, trim=True, batch_interface=None):
    '''
    Add the activities to the feed
    functions used in tasks need to be at the main level of the module
    '''
    t = timer()
    msg_format = 'running %s.add_many operation for %s activities batch interface %s and trim %s'
    logger.debug(msg_format, feed, len(activities), batch_interface, trim)
    feed.add_many(activities, batch_interface=batch_interface, trim=trim)
    logger.debug('add many operation took %s seconds', t.next())

def remove_operation(feed, activities, trim=True, batch_interface=None):
    '''
    Remove the activities from the feed
    functions used in tasks need to be at the main level of the module
    '''
    t = timer()
    msg_format = 'running %s.remove_many operation for %s activities batch interface %s'
    logger.debug(msg_format, feed, len(activities), batch_interface)
    feed.remove_many(activities, trim=trim, batch_interface=batch_interface)
    logger.debug('remove many operation took %s seconds', t.next())

class FanoutPriority(object):
    HIGH = 'HIGH'
    LOW = 'LOW'

class Feedly(object):

    '''
    The Feedly class handles the fanout from a user's activity
    to all their follower's feeds

    .. note::

        Fanout is the process which pushes a little bit of data to all of your
        followers in many small and asynchronous tasks.

    To write your own Feedly class you will need to implement

    - get_user_follower_ids
    - feed_classes
    - user_feed_class

    **Example** ::

        from feedly.feed_managers.base import Feedly

        class PinFeedly(Feedly):
            # customize the feed classes we write to
            feed_classes = dict(
                normal=PinFeed,
                aggregated=AggregatedPinFeed
            )
            # customize the user feed class
            user_feed_class = UserPinFeed

            # define how feedly can get the follower ids
            def get_user_follower_ids(self, user_id):
                return Follow.objects.filter(target=user_id).values_list('user_id', flat=True)

            # utility functions to easy integration for your project
            def add_pin(self, pin):
                activity = pin.create_activity()
                # add user activity adds it to the user feed, and starts the fanout
                self.add_user_activity(pin.user_id, activity)

            def remove_pin(self, pin):
                activity = pin.create_activity()
                # removes the pin from the user's followers feeds
                self.remove_user_activity(pin.user_id, activity)

    '''
    # : a dictionary with the feeds to fanout to
    # : for example feed_classes = dict(normal=PinFeed, aggregated=AggregatedPinFeed)
    feed_classes = dict(
        normal=RedisFeed
    )
    # : the user feed class (it stores the latest activity by one user)
    user_feed_class = UserBaseFeed

    # : the number of activities which enter your feed when you follow someone
    follow_activity_limit = 5000
    # : the number of users which are handled in one asynchronous task
    # : when doing the fanout
    fanout_chunk_size = 100

    # maps between priority and fanout tasks
    priority_fanout_task = {
        FanoutPriority.HIGH: fanout_operation_hi_priority,
        FanoutPriority.LOW: fanout_operation_low_priority
    }

    def get_user_follower_ids(self, user_id):
        '''
        Returns a dict of users ids which follow the given user grouped by
        priority/importance

        eg.
        {'HIGH': [...], 'LOW': [...]}

        :param user_id: the user id for which to get the follower ids
        '''
        raise NotImplementedError()

    def add_user_activity(self, user_id, activity):
        '''
        Store the new activity and then fanout to user followers

        This function will
        - store the activity in the activity storage
        - store it in the user feed (list of activities for one user)
        - fanout for all feed_classes

        :param user_id: the id of the user
        :param activity: the activity which to add
        '''
        # add into the global activity cache (if we are using it)
        self.user_feed_class.insert_activity(activity)
        # now add to the user's personal feed
        user_feed = self.get_user_feed(user_id)
        user_feed.add(activity)
        operation_kwargs = dict(activities=[activity], trim=True)

        for priority_group, follower_ids in self.get_user_follower_ids(user_id=user_id).items():
            # create the fanout tasks
            for feed_class in self.feed_classes.values():
                self.create_fanout_tasks(
                    follower_ids,
                    feed_class,
                    add_operation,
                    operation_kwargs=operation_kwargs,
                    fanout_priority=priority_group
                )

    def remove_user_activity(self, user_id, activity):
        '''
        Remove the activity and then fanout to user followers

        :param user_id: the id of the user
        :param activity: the activity which to add
        '''
        # we don't remove from the global feed due to race conditions
        # but we do remove from the personal feed
        user_feed = self.get_user_feed(user_id)
        user_feed.remove(activity)

        # no need to trim when removing items
        operation_kwargs = dict(activities=[activity], trim=False)

        for priority_group, follower_ids in self.get_user_follower_ids(user_id=user_id).items():
            for feed_class in self.feed_classes.values():
                self.create_fanout_tasks(
                    follower_ids,
                    feed_class,
                    remove_operation,
                    operation_kwargs=operation_kwargs,
                    fanout_priority=priority_group
                )

    def get_feeds(self, user_id):
        '''
        get the feed that contains the sum of all activity
        from feeds :user_id is subscribed to

        :returns dict: a dictionary with the feeds we're pushing to
        '''
        return dict([(k, feed(user_id)) for k, feed in self.feed_classes.items()])

    def get_user_feed(self, user_id):
        '''
        feed where activity from :user_id is saved

        :param user_id: the id of the user
        '''
        return self.user_feed_class(user_id)

    def update_user_activities(self, activities):
        '''
        Update the user activities
        :param activities: the activities to update
        '''
        for activity in activities:
            self.add_user_activity(activity.actor_id, activity)

    def update_user_activity(self, activity):
        self.update_user_activities([activity])

    def follow_feed(self, feed, activities):
        '''
        copies source_feed entries into feed
        it will only copy follow_activity_limit activities

        :param feed: the feed to copy to
        :param activities: the activities to copy into the feed
        '''
        if activities:
            return feed.add_many(activities)

    def unfollow_feed(self, feed, source_feed):
        '''
        removes entries originating from the source feed form the feed class
        this will remove all activities, so this could take a wh
        :param feed: the feed to copy to
        :param source_feed: the feed with a list of activities to remove
        '''
        activities = source_feed[:]  # need to slice
        if activities:
            return feed.remove_many(activities)

    def follow_user(self, user_id, target_user_id, async=True):
        '''
        user_id starts following target_user_id

        :param user_id: the user which is doing the following/unfollowing
        :target_user_id: the user which is being unfollowed
        '''
        source_feed = self.get_user_feed(target_user_id)
        # fetch the activities only once
        activities = source_feed[:self.follow_activity_limit]
        for user_feed in self.get_feeds(user_id).values():
            self.follow_feed(user_feed, activities)

    def unfollow_user(self, user_id, target_user_id, async=True):
        '''
        unfollows the user

        :param user_id: the user which is doing the following/unfollowing
        :target_user_id: the user which is being unfollowed
        '''
        if async:
            unfollow_many_fn = unfollow_many.delay
        else:
            unfollow_many_fn = unfollow_many

        unfollow_many_fn(self, user_id, [target_user_id])

    def follow_many_users(self, user_id, target_ids, async=True):
        '''
        copies feeds for target_ids in user_id

        :param user_id: the user which is doing the following/unfollowing
        :param target_ids: the user to follow
        :param async: controls if the operation should be done via celery
        '''
        if async:
            follow_many_fn = follow_many.delay
        else:
            follow_many_fn = follow_many

        follow_many_fn(
            self,
            user_id,
            target_ids,
            self.follow_activity_limit
        )

    def get_fanout_task(self, priority=None, feed_class=None):
        '''
        Returns the fanout task taking priority in account.

        :param priority: the priority of the task
        :param feed_class: the feed_class the task will write to
        '''
        return self.priority_fanout_task.get(priority, fanout_operation)

    def create_fanout_tasks(self, follower_ids, feed_class, operation, operation_kwargs=None, fanout_priority=None):
        '''
        Creates the fanout task for the given activities and feed classes
        followers

        It takes the following ids and distributes them per fanout_chunk_size
        into smaller tasks

        :param follower_ids: specify the list of followers
        :param feed_class: the feed classes to run the operation on
        :param operation: the operation function applied to all follower feeds
        :param operation_kwargs: kwargs passed to the operation
        :param fanout_priority: the priority set to this fanout
        '''
        fanout_task = self.get_fanout_task(
            fanout_priority, feed_class=feed_class)
        if not fanout_task:
            return []
        chunk_size = self.fanout_chunk_size
        user_ids_chunks = list(chunks(follower_ids, chunk_size))
        msg_format = 'spawning %s subtasks for %s user ids in chunks of %s users'
        logger.info(
            msg_format, len(user_ids_chunks), len(follower_ids), chunk_size)
        tasks = []
        # now actually create the tasks
        for ids_chunk in user_ids_chunks:
            task = fanout_task.delay(
                feed_manager=self,
                feed_class=feed_class,
                user_ids=ids_chunk,
                operation=operation,
                operation_kwargs=operation_kwargs
            )
            tasks.append(task)
        return tasks

    def fanout(self, user_ids, feed_class, operation, operation_kwargs):
        '''
        This functionality is called from within feedly.tasks.fanout_operation

        :param user_ids: the list of user ids which feeds we should apply the
        operation against
        :param feed_class: the feed to run the operation on
        :param operation: the operation to run on the feed
        :param operation_kwargs: kwargs to pass to the operation

        '''
        separator = '===' * 10
        logger.info('%s starting fanout %s', separator, separator)
        batch_context_manager = feed_class.get_timeline_batch_interface()
        msg_format = 'starting batch interface for feed %s, fanning out to %s users'
        with batch_context_manager as batch_interface:
            logger.info(msg_format, feed_class, len(user_ids))
            operation_kwargs['batch_interface'] = batch_interface
            for user_id in user_ids:
                logger.debug('now handling fanout to user %s', user_id)
                feed = feed_class(user_id)
                operation(feed, **operation_kwargs)
        logger.info('finished fanout for feed %s', feed_class)

    def batch_import(self, user_id, activities, fanout=True, chunk_size=500):
        '''
        Batch import all of the users activities and distributes
        them to the users followers

        **Example**::

            activities = [long list of activities]
            feedly.batch_import(13, activities, 500)

        :param user_id: the user who created the activities
        :param activities: a list of activities from this user
        :param fanout: if we should run the fanout or not
        :param chunk_size: per how many activities to run the batch operations

        '''
        activities = list(activities)
        # skip empty lists
        if not activities:
            return
        logger.info('running batch import for user %s', user_id)

        user_feed = self.get_user_feed(user_id)
        if activities[0].actor_id != user_id:
            raise ValueError('Send activities for only one user please')

        activity_chunks = list(chunks(activities, chunk_size))
        logger.info('processing %s items in %s chunks of %s',
                    len(activities), len(activity_chunks), chunk_size)

        for index, activity_chunk in enumerate(activity_chunks):
            # first insert into the global activity storage
            self.user_feed_class.insert_activities(activity_chunk)
            logger.info(
                'inserted chunk %s (length %s) into the global activity store', index, len(activity_chunk))
            # next add the activities to the users personal timeline
            user_feed.add_many(activity_chunk, trim=False)
            logger.info(
                'inserted chunk %s (length %s) into the user feed', index, len(activity_chunk))
            # now start a big fanout task
            if fanout:
                logger.info('starting task fanout for chunk %s', index)
                follower_ids_by_prio = self.get_user_follower_ids(
                    user_id=user_id)
                # create the fanout tasks
                operation_kwargs = dict(activities=activity_chunk, trim=False)
                for feed_class in self.feed_classes.values():
                    for priority_group, fids in follower_ids_by_prio.items():
                        self.create_fanout_tasks(
                            fids,
                            feed_class,
                            add_operation,
                            fanout_priority=priority_group,
                            operation_kwargs=operation_kwargs
                        )
<commit_msg>trailing whitespace<commit_after>from feedly.feeds.base import UserBaseFeed
from feedly.tasks import follow_many, unfollow_many
from feedly.tasks import fanout_operation
from feedly.tasks import fanout_operation_hi_priority
from feedly.tasks import fanout_operation_low_priority
from feedly.utils import chunks
from feedly.utils.timing import timer
import logging
from feedly.feeds.redis import RedisFeed

logger = logging.getLogger(__name__)

def add_operation(feed, activities, trim=True, batch_interface=None):
    '''
    Add the activities to the feed
    functions used in tasks need to be at the main level of the module
    '''
    t = timer()
    msg_format = 'running %s.add_many operation for %s activities batch interface %s and trim %s'
    logger.debug(msg_format, feed, len(activities), batch_interface, trim)
    feed.add_many(activities, batch_interface=batch_interface, trim=trim)
    logger.debug('add many operation took %s seconds', t.next())

def remove_operation(feed, activities, trim=True, batch_interface=None):
    '''
    Remove the activities from the feed
    functions used in tasks need to be at the main level of the module
    '''
    t = timer()
    msg_format = 'running %s.remove_many operation for %s activities batch interface %s'
    logger.debug(msg_format, feed, len(activities), batch_interface)
    feed.remove_many(activities, trim=trim, batch_interface=batch_interface)
    logger.debug('remove many operation took %s seconds', t.next())

class FanoutPriority(object):
    HIGH = 'HIGH'
    LOW = 'LOW'

class Feedly(object):

    '''
    The Feedly class handles the fanout from a user's activity
    to all their follower's feeds

    .. note::

        Fanout is the process which pushes a little bit of data to all of your
        followers in many small and asynchronous tasks.

    To write your own Feedly class you will need to implement

    - get_user_follower_ids
    - feed_classes
    - user_feed_class

    **Example** ::

        from feedly.feed_managers.base import Feedly

        class PinFeedly(Feedly):
            # customize the feed classes we write to
            feed_classes = dict(
                normal=PinFeed,
                aggregated=AggregatedPinFeed
            )
            # customize the user feed class
            user_feed_class = UserPinFeed

            # define how feedly can get the follower ids
            def get_user_follower_ids(self, user_id):
                return Follow.objects.filter(target=user_id).values_list('user_id', flat=True)

            # utility functions to easy integration for your project
            def add_pin(self, pin):
                activity = pin.create_activity()
                # add user activity adds it to the user feed, and starts the fanout
                self.add_user_activity(pin.user_id, activity)

            def remove_pin(self, pin):
                activity = pin.create_activity()
                # removes the pin from the user's followers feeds
                self.remove_user_activity(pin.user_id, activity)

    '''
    # : a dictionary with the feeds to fanout to
    # : for example feed_classes = dict(normal=PinFeed, aggregated=AggregatedPinFeed)
    feed_classes = dict(
        normal=RedisFeed
    )
    # : the user feed class (it stores the latest activity by one user)
    user_feed_class = UserBaseFeed

    # : the number of activities which enter your feed when you follow someone
    follow_activity_limit = 5000
    # : the number of users which are handled in one asynchronous task
    # : when doing the fanout
    fanout_chunk_size = 100

    # maps between priority and fanout tasks
    priority_fanout_task = {
        FanoutPriority.HIGH: fanout_operation_hi_priority,
        FanoutPriority.LOW: fanout_operation_low_priority
    }

    def get_user_follower_ids(self, user_id):
        '''
        Returns a dict of users ids which follow the given user grouped by
        priority/importance

        eg.
        {'HIGH': [...], 'LOW': [...]}

        :param user_id: the user id for which to get the follower ids
        '''
        raise NotImplementedError()

    def add_user_activity(self, user_id, activity):
        '''
        Store the new activity and then fanout to user followers

        This function will
        - store the activity in the activity storage
        - store it in the user feed (list of activities for one user)
        - fanout for all feed_classes

        :param user_id: the id of the user
        :param activity: the activity which to add
        '''
        # add into the global activity cache (if we are using it)
        self.user_feed_class.insert_activity(activity)
        # now add to the user's personal feed
        user_feed = self.get_user_feed(user_id)
        user_feed.add(activity)
        operation_kwargs = dict(activities=[activity], trim=True)

        for priority_group, follower_ids in self.get_user_follower_ids(user_id=user_id).items():
            # create the fanout tasks
            for feed_class in self.feed_classes.values():
                self.create_fanout_tasks(
                    follower_ids,
                    feed_class,
                    add_operation,
                    operation_kwargs=operation_kwargs,
                    fanout_priority=priority_group
                )

    def remove_user_activity(self, user_id, activity):
        '''
        Remove the activity and then fanout to user followers

        :param user_id: the id of the user
        :param activity: the activity which to add
        '''
        # we don't remove from the global feed due to race conditions
        # but we do remove from the personal feed
        user_feed = self.get_user_feed(user_id)
        user_feed.remove(activity)

        # no need to trim when removing items
        operation_kwargs = dict(activities=[activity], trim=False)

        for priority_group, follower_ids in self.get_user_follower_ids(user_id=user_id).items():
            for feed_class in self.feed_classes.values():
                self.create_fanout_tasks(
                    follower_ids,
                    feed_class,
                    remove_operation,
                    operation_kwargs=operation_kwargs,
                    fanout_priority=priority_group
                )

    def get_feeds(self, user_id):
        '''
        get the feed that contains the sum of all activity
        from feeds :user_id is subscribed to

        :returns dict: a dictionary with the feeds we're pushing to
        '''
        return dict([(k, feed(user_id)) for k, feed in self.feed_classes.items()])

    def get_user_feed(self, user_id):
        '''
        feed where activity from :user_id is saved

        :param user_id: the id of the user
        '''
        return self.user_feed_class(user_id)

    def update_user_activities(self, activities):
        '''
        Update the user activities
        :param activities: the activities to update
        '''
        for activity in activities:
            self.add_user_activity(activity.actor_id, activity)

    def update_user_activity(self, activity):
        self.update_user_activities([activity])

    def follow_feed(self, feed, activities):
        '''
        copies source_feed entries into feed
        it will only copy follow_activity_limit activities

        :param feed: the feed to copy to
        :param activities: the activities to copy into the feed
        '''
        if activities:
            return feed.add_many(activities)

    def unfollow_feed(self, feed, source_feed):
        '''
        removes entries originating from the source feed form the feed class
        this will remove all activities, so this could take a wh
        :param feed: the feed to copy to
        :param source_feed: the feed with a list of activities to remove
        '''
        activities = source_feed[:]  # need to slice
        if activities:
            return feed.remove_many(activities)

    def follow_user(self, user_id, target_user_id, async=True):
        '''
        user_id starts following target_user_id

        :param user_id: the user which is doing the following/unfollowing
        :target_user_id: the user which is being unfollowed
        '''
        source_feed = self.get_user_feed(target_user_id)
        # fetch the activities only once
        activities = source_feed[:self.follow_activity_limit]
        for user_feed in self.get_feeds(user_id).values():
            self.follow_feed(user_feed, activities)

    def unfollow_user(self, user_id, target_user_id, async=True):
        '''
        unfollows the user

        :param user_id: the user which is doing the following/unfollowing
        :target_user_id: the user which is being unfollowed
        '''
        if async:
            unfollow_many_fn = unfollow_many.delay
        else:
            unfollow_many_fn = unfollow_many

        unfollow_many_fn(self, user_id, [target_user_id])

    def follow_many_users(self, user_id, target_ids, async=True):
        '''
        copies feeds for target_ids in user_id

        :param user_id: the user which is doing the following/unfollowing
        :param target_ids: the user to follow
        :param async: controls if the operation should be done via celery
        '''
        if async:
            follow_many_fn = follow_many.delay
        else:
            follow_many_fn = follow_many

        follow_many_fn(
            self,
            user_id,
            target_ids,
            self.follow_activity_limit
        )

    def get_fanout_task(self, priority=None, feed_class=None):
        '''
        Returns the fanout task taking priority in account.

        :param priority: the priority of the task
        :param feed_class: the feed_class the task will write to
        '''
        return self.priority_fanout_task.get(priority, fanout_operation)

    def create_fanout_tasks(self, follower_ids, feed_class, operation, operation_kwargs=None, fanout_priority=None):
        '''
        Creates the fanout task for the given activities and feed classes
        followers

        It takes the following ids and distributes them per fanout_chunk_size
        into smaller tasks

        :param follower_ids: specify the list of followers
        :param feed_class: the feed classes to run the operation on
        :param operation: the operation function applied to all follower feeds
        :param operation_kwargs: kwargs passed to the operation
        :param fanout_priority: the priority set to this fanout
        '''
        fanout_task = self.get_fanout_task(
            fanout_priority, feed_class=feed_class)
        if not fanout_task:
            return []
        chunk_size = self.fanout_chunk_size
        user_ids_chunks = list(chunks(follower_ids, chunk_size))
        msg_format = 'spawning %s subtasks for %s user ids in chunks of %s users'
        logger.info(
            msg_format, len(user_ids_chunks), len(follower_ids), chunk_size)
        tasks = []
        # now actually create the tasks
        for ids_chunk in user_ids_chunks:
            task = fanout_task.delay(
                feed_manager=self,
                feed_class=feed_class,
                user_ids=ids_chunk,
                operation=operation,
                operation_kwargs=operation_kwargs
            )
            tasks.append(task)
        return tasks

    def fanout(self, user_ids, feed_class, operation, operation_kwargs):
        '''
        This functionality is called from within feedly.tasks.fanout_operation

        :param user_ids: the list of user ids which feeds we should apply the
        operation against
        :param feed_class: the feed to run the operation on
        :param operation: the operation to run on the feed
        :param operation_kwargs: kwargs to pass to the operation

        '''
        separator = '===' * 10
        logger.info('%s starting fanout %s', separator, separator)
        batch_context_manager = feed_class.get_timeline_batch_interface()
        msg_format = 'starting batch interface for feed %s, fanning out to %s users'
        with batch_context_manager as batch_interface:
            logger.info(msg_format, feed_class, len(user_ids))
            operation_kwargs['batch_interface'] = batch_interface
            for user_id in user_ids:
                logger.debug('now handling fanout to user %s', user_id)
                feed = feed_class(user_id)
                operation(feed, **operation_kwargs)
        logger.info('finished fanout for feed %s', feed_class)

    def batch_import(self, user_id, activities, fanout=True, chunk_size=500):
        '''
        Batch import all of the users activities and distributes
        them to the users followers

        **Example**::

            activities = [long list of activities]
            feedly.batch_import(13, activities, 500)

        :param user_id: the user who created the activities
        :param activities: a list of activities from this user
        :param fanout: if we should run the fanout or not
        :param chunk_size: per how many activities to run the batch operations

        '''
        activities = list(activities)
        # skip empty lists
        if not activities:
            return
        logger.info('running batch import for user %s', user_id)

        user_feed = self.get_user_feed(user_id)
        if activities[0].actor_id != user_id:
            raise ValueError('Send activities for only one user please')

        activity_chunks = list(chunks(activities, chunk_size))
        logger.info('processing %s items in %s chunks of %s',
                    len(activities), len(activity_chunks), chunk_size)

        for index, activity_chunk in enumerate(activity_chunks):
            # first insert into the global activity storage
            self.user_feed_class.insert_activities(activity_chunk)
            logger.info(
                'inserted chunk %s (length %s) into the global activity store', index, len(activity_chunk))
            # next add the activities to the users personal timeline
            user_feed.add_many(activity_chunk, trim=False)
            logger.info(
                'inserted chunk %s (length %s) into the user feed', index, len(activity_chunk))
            # now start a big fanout task
            if fanout:
                logger.info('starting task fanout for chunk %s', index)
                follower_ids_by_prio = self.get_user_follower_ids(
                    user_id=user_id)
                # create the fanout tasks
                operation_kwargs = dict(activities=activity_chunk, trim=False)
                for feed_class in self.feed_classes.values():
                    for priority_group, fids in follower_ids_by_prio.items():
                        self.create_fanout_tasks(
                            fids,
                            feed_class,
                            add_operation,
                            fanout_priority=priority_group,
                            operation_kwargs=operation_kwargs
                        )
<|endoftext|>


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

__author__ = 'Daniel Zhang (張道博)'
__copyright__ = 'Copyright (c) 2014, University of Hawaii Smart Energy Project'
__license__ = 'https://raw.github' \
              '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD' \
              '-LICENSE.txt'

from enum import Enum

class MSGAggregationTypes(Enum):
    """
    Types for aggregation.
    """
    weather = 1
    egauge = 2
    circuit = 3
    irradiance = 4

<commit_msg>adding notification history types<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

__author__ = 'Daniel Zhang (張道博)'
__copyright__ = 'Copyright (c) 2014, University of Hawaii Smart Energy Project'
__license__ = 'https://raw.github' \
              '.com/Hawaii-Smart-Energy-Project/Maui-Smart-Grid/master/BSD' \
              '-LICENSE.txt'

from enum import Enum

class MSGAggregationTypes(Enum):
    """
    Types for aggregation.
    """
    weather = 1
    egauge = 2
    circuit = 3
    irradiance = 4

class MSGNotificationHistoryTypes(Enum):
    """
    Types for the notification history.
    """
    msg_db_exporter = 1
    msg_data_aggregator = 2
    msg_export_summary = 3
<|endoftext|>