xref: /aosp_15_r20/external/webrtc/tools_webrtc/network_emulator/network_emulator.py (revision d9f758449e529ab9291ac668be2861e7a55c2422)
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