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