OLD | NEW |
| (Empty) |
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._CreateDummynetPipe( | |
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._CreateDummynetPipe( | |
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._CreateDummynetRule(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._CreateDummynetRule(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 CheckPermissions(): | |
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 _CreateDummynetRule(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 _RunIpfwCommand(add_part + ['src-port', '%s-%s' % port_range], | |
122 'Failed to add Dummynet src-port rule.') | |
123 _RunIpfwCommand(add_part + ['dst-port', '%s-%s' % port_range], | |
124 'Failed to add Dummynet dst-port rule.') | |
125 return self._rule_counter | |
126 | |
127 def _CreateDummynetPipe(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 _RunIpfwCommand(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 _RunIpfwCommand(['-f', 'flush'], | |
159 'Failed to flush Dummynet rules!') | |
160 _RunIpfwCommand(['-f', 'pipe', 'flush'], | |
161 'Failed to flush Dummynet pipes!') | |
162 | |
163 def _RunIpfwCommand(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() | |
OLD | NEW |