Home | History | Annotate | Download | only in network_emulator 
      1 #!/usr/bin/env python
      2 #  Copyright (c) 2012 The WebRTC project authors. All Rights Reserved.
      3 #
      4 #  Use of this source code is governed by a BSD-style license
      5 #  that can be found in the LICENSE file in the root of the source
      6 #  tree. An additional intellectual property rights grant can be found
      7 #  in the file PATENTS.  All contributing project authors may
      8 #  be found in the AUTHORS file in the root of the source tree.
      9 
     10 """Script for constraining traffic on the local machine."""
     11 
     12 import ctypes
     13 import logging
     14 import os
     15 import subprocess
     16 import sys
     17 
     18 
     19 class NetworkEmulatorError(BaseException):
     20   """Exception raised for errors in the network emulator.
     21 
     22   Attributes:
     23     fail_msg: User defined error message.
     24     cmd: Command for which the exception was raised.
     25     returncode: Return code of running the command.
     26     stdout: Output of running the command.
     27     stderr: Error output of running the command.
     28   """
     29 
     30   def __init__(self, fail_msg, cmd=None, returncode=None, output=None,
     31                error=None):
     32     BaseException.__init__(self, fail_msg)
     33     self.fail_msg = fail_msg
     34     self.cmd = cmd
     35     self.returncode = returncode
     36     self.output = output
     37     self.error = error
     38 
     39 
     40 class NetworkEmulator(object):
     41   """A network emulator that can constrain the network using Dummynet."""
     42 
     43   def __init__(self, connection_config, port_range):
     44     """Constructor.
     45 
     46     Args:
     47         connection_config: A config.ConnectionConfig object containing the
     48             characteristics for the connection to be emulation.
     49         port_range: Tuple containing two integers defining the port range.
     50     """
     51     self._pipe_counter = 0
     52     self._rule_counter = 0
     53     self._port_range = port_range
     54     self._connection_config = connection_config
     55 
     56   def emulate(self, target_ip):
     57     """Starts a network emulation by setting up Dummynet rules.
     58 
     59     Args:
     60         target_ip: The IP address of the interface that shall be that have the
     61             network constraints applied to it.
     62     """
     63     receive_pipe_id = self._create_dummynet_pipe(
     64         self._connection_config.receive_bw_kbps,
     65         self._connection_config.delay_ms,
     66         self._connection_config.packet_loss_percent,
     67         self._connection_config.queue_slots)
     68     logging.debug('Created receive pipe: %s', receive_pipe_id)
     69     send_pipe_id = self._create_dummynet_pipe(
     70         self._connection_config.send_bw_kbps,
     71         self._connection_config.delay_ms,
     72         self._connection_config.packet_loss_percent,
     73         self._connection_config.queue_slots)
     74     logging.debug('Created send pipe: %s', send_pipe_id)
     75 
     76     # Adding the rules will start the emulation.
     77     incoming_rule_id = self._create_dummynet_rule(receive_pipe_id, 'any',
     78                                                   target_ip, self._port_range)
     79     logging.debug('Created incoming rule: %s', incoming_rule_id)
     80     outgoing_rule_id = self._create_dummynet_rule(send_pipe_id, target_ip,
     81                                                   'any', self._port_range)
     82     logging.debug('Created outgoing rule: %s', outgoing_rule_id)
     83 
     84   @staticmethod
     85   def check_permissions():
     86     """Checks if permissions are available to run Dummynet commands.
     87 
     88     Raises:
     89       NetworkEmulatorError: If permissions to run Dummynet commands are not
     90       available.
     91     """
     92     try:
     93       if os.getuid() != 0:
     94         raise NetworkEmulatorError('You must run this script with sudo.')
     95     except AttributeError:
     96 
     97     # AttributeError will be raised on Windows.
     98       if ctypes.windll.shell32.IsUserAnAdmin() == 0:
     99         raise NetworkEmulatorError('You must run this script with administrator'
    100                                    ' privileges.')
    101 
    102   def _create_dummynet_rule(self, pipe_id, from_address, to_address,
    103                             port_range):
    104     """Creates a network emulation rule and returns its ID.
    105 
    106     Args:
    107         pipe_id: integer ID of the pipe.
    108         from_address: The IP address to match source address. May be an IP or
    109           'any'.
    110         to_address: The IP address to match destination address. May be an IP or
    111           'any'.
    112         port_range: The range of ports the rule shall be applied on. Must be
    113           specified as a tuple of with two integers.
    114     Returns:
    115         The ID of the rule, starting at 100. The rule ID increments with 100 for
    116         each rule being added.
    117     """
    118     self._rule_counter += 100
    119     add_part = ['add', self._rule_counter, 'pipe', pipe_id,
    120                 'ip', 'from', from_address, 'to', to_address]
    121     _run_ipfw_command(add_part + ['src-port', '%s-%s' % port_range],
    122                             'Failed to add Dummynet src-port rule.')
    123     _run_ipfw_command(add_part + ['dst-port', '%s-%s' % port_range],
    124                             'Failed to add Dummynet dst-port rule.')
    125     return self._rule_counter
    126 
    127   def _create_dummynet_pipe(self, bandwidth_kbps, delay_ms, packet_loss_percent,
    128                             queue_slots):
    129     """Creates a Dummynet pipe and return its ID.
    130 
    131     Args:
    132         bandwidth_kbps: Bandwidth.
    133         delay_ms: Delay for a one-way trip of a packet.
    134         packet_loss_percent: Float value of packet loss, in percent.
    135         queue_slots: Size of the queue.
    136     Returns:
    137         The ID of the pipe, starting at 1.
    138     """
    139     self._pipe_counter += 1
    140     cmd = ['pipe', self._pipe_counter, 'config',
    141            'bw', str(bandwidth_kbps/8) + 'KByte/s',
    142            'delay', '%sms' % delay_ms,
    143            'plr', (packet_loss_percent/100.0),
    144            'queue', queue_slots]
    145     error_message = 'Failed to create Dummynet pipe. '
    146     if sys.platform.startswith('linux'):
    147       error_message += ('Make sure you have loaded the ipfw_mod.ko module to '
    148                         'your kernel (sudo insmod /path/to/ipfw_mod.ko).')
    149     _run_ipfw_command(cmd, error_message)
    150     return self._pipe_counter
    151 
    152 def cleanup():
    153   """Stops the network emulation by flushing all Dummynet rules.
    154 
    155   Notice that this will flush any rules that may have been created previously
    156   before starting the emulation.
    157   """
    158   _run_ipfw_command(['-f', 'flush'],
    159                           'Failed to flush Dummynet rules!')
    160   _run_ipfw_command(['-f', 'pipe', 'flush'],
    161                           'Failed to flush Dummynet pipes!')
    162 
    163 def _run_ipfw_command(command, fail_msg=None):
    164   """Executes a command and prefixes the appropriate command for
    165      Windows or Linux/UNIX.
    166 
    167   Args:
    168     command: Command list to execute.
    169     fail_msg: Message describing the error in case the command fails.
    170 
    171   Raises:
    172     NetworkEmulatorError: If command fails a message is set by the fail_msg
    173     parameter.
    174   """
    175   if sys.platform == 'win32':
    176     ipfw_command = ['ipfw.exe']
    177   else:
    178     ipfw_command = ['sudo', '-n', 'ipfw']
    179 
    180   cmd_list = ipfw_command[:] + [str(x) for x in command]
    181   cmd_string = ' '.join(cmd_list)
    182   logging.debug('Running command: %s', cmd_string)
    183   process = subprocess.Popen(cmd_list, stdout=subprocess.PIPE,
    184                              stderr=subprocess.PIPE)
    185   output, error = process.communicate()
    186   if process.returncode != 0:
    187     raise NetworkEmulatorError(fail_msg, cmd_string, process.returncode, output,
    188                                error)
    189   return output.strip()
    190