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
12import ctypes
13import logging
14import os
15import subprocess
16import sys
17
18
19class 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
40class 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
152def 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
163def _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