diff --git a/tools/network_emulator/emulate.py b/tools/network_emulator/emulate.py index e25670574..af22d86bb 100755 --- a/tools/network_emulator/emulate.py +++ b/tools/network_emulator/emulate.py @@ -9,15 +9,16 @@ """Script for constraining traffic on the local machine.""" + import logging import optparse -import os import socket import sys import config import network_emulator + _DEFAULT_LOG_LEVEL = logging.INFO # Default port range to apply network constraints on. @@ -68,7 +69,7 @@ def _parse_args(): 'ID Name Receive Send Queue Delay loss \n' '-- ---- --------- -------- ----- ------- ------\n' '%s\n' % presets_string)) - parser.add_option('-p', '--preset', type='int', default=2, + parser.add_option('-p', '--preset', type='int', default=_DEFAULT_PRESET_ID, help=('ConnectionConfig configuration, specified by ID. ' 'Default: %default')) parser.add_option('-r', '--receive-bw', type='int', @@ -98,11 +99,11 @@ def _parse_args(): options = parser.parse_args()[0] - # Find preset by ID, if specified: + # Find preset by ID, if specified. if options.preset and not _PRESETS_DICT.has_key(options.preset): parser.error('Invalid preset: %s' % options.preset) - # Simple validation of the IP address, if supplied: + # Simple validation of the IP address, if supplied. if options.target_ip: try: socket.inet_aton(options.target_ip) @@ -133,27 +134,21 @@ def _set_logger(verbose): def _main(): - """Checks arguments, permissions and runs a network emulation.""" - if os.name != 'posix': - print >> sys.stderr, 'This script is only supported on Linux and Mac.' - return 1 - options = _parse_args() # Build a configuration object. Override any preset configuration settings if # a value of a setting was also given as a flag. connection_config = _PRESETS_DICT[options.preset] - if options.receive_bw: + if options.receive_bw is not _DEFAULT_PRESET.receive_bw_kbps: connection_config.receive_bw_kbps = options.receive_bw - if options.send_bw: + if options.send_bw is not _DEFAULT_PRESET.send_bw_kbps: connection_config.send_bw_kbps = options.send_bw - if options.delay: + if options.delay is not _DEFAULT_PRESET.delay_ms: connection_config.delay_ms = options.delay - if options.packet_loss: + if options.packet_loss is not _DEFAULT_PRESET.packet_loss_percent: connection_config.packet_loss_percent = options.packet_loss - if options.queue: + if options.queue is not _DEFAULT_PRESET.queue_slots: connection_config.queue_slots = options.queue - emulator = network_emulator.NetworkEmulator(connection_config, options.port_range) try: diff --git a/tools/network_emulator/network_emulator.py b/tools/network_emulator/network_emulator.py index 9ae1bdb5b..00b758207 100644 --- a/tools/network_emulator/network_emulator.py +++ b/tools/network_emulator/network_emulator.py @@ -9,6 +9,7 @@ """Script for constraining traffic on the local machine.""" +import ctypes import logging import os import subprocess @@ -88,21 +89,26 @@ class NetworkEmulator(object): NetworkEmulatorError: If permissions to run Dummynet commands are not available. """ - if os.geteuid() != 0: - _run_shell_command( - ['sudo', '-n', 'ipfw', '-h'], - msg=('Cannot run \'ipfw\' command. This script must be run as ' - 'root or have password-less sudo access to this command.')) + try: + if os.getuid() != 0: + raise NetworkEmulatorError('You must run this script with sudo.') + except AttributeError: - @staticmethod - def cleanup(): + # AttributeError will be raised on Windows. + if ctypes.windll.shell32.IsUserAnAdmin() == 0: + raise NetworkEmulatorError('You must run this script with administrator' + 'privileges.') + + def cleanup(self): """Stops the network emulation by flushing all Dummynet rules. Notice that this will flush any rules that may have been created previously before starting the emulation. """ - _run_shell_command(['sudo', 'ipfw', '-f', 'flush'], - 'Failed to flush Dummynet rules!') + self._run_ipfw_command(['-f', 'flush'], + 'Failed to flush Dummynet rules!') + self._run_ipfw_command(['-f', 'pipe', 'flush'], + 'Failed to flush Dummynet pipes!') def _create_dummynet_rule(self, pipe_id, from_address, to_address, port_range): @@ -121,12 +127,12 @@ class NetworkEmulator(object): each rule being added. """ self._rule_counter += 100 - add_part = ['sudo', 'ipfw', 'add', self._rule_counter, 'pipe', pipe_id, + add_part = ['add', self._rule_counter, 'pipe', pipe_id, 'ip', 'from', from_address, 'to', to_address] - _run_shell_command(add_part + ['src-port', '%s-%s' % port_range], - 'Failed to add Dummynet src-port rule.') - _run_shell_command(add_part + ['dst-port', '%s-%s' % port_range], - 'Failed to add Dummynet dst-port rule.') + self._run_ipfw_command(add_part + ['src-port', '%s-%s' % port_range], + 'Failed to add Dummynet src-port rule.') + self._run_ipfw_command(add_part + ['dst-port', '%s-%s' % port_range], + 'Failed to add Dummynet dst-port rule.') return self._rule_counter def _create_dummynet_pipe(self, bandwidth_kbps, delay_ms, packet_loss_percent, @@ -142,7 +148,7 @@ class NetworkEmulator(object): The ID of the pipe, starting at 1. """ self._pipe_counter += 1 - cmd = ['sudo', 'ipfw', 'pipe', self._pipe_counter, 'config', + cmd = ['pipe', self._pipe_counter, 'config', 'bw', str(bandwidth_kbps/8) + 'KByte/s', 'delay', '%sms' % delay_ms, 'plr', (packet_loss_percent/100.0), @@ -150,10 +156,13 @@ class NetworkEmulator(object): error_message = 'Failed to create Dummynet pipe. ' if sys.platform.startswith('linux'): error_message += ('Make sure you have loaded the ipfw_mod.ko module to ' - 'your kernel (sudo insmod /path/to/ipfw_mod.ko)') - _run_shell_command(cmd, error_message) + 'your kernel (sudo insmod /path/to/ipfw_mod.ko).') + self._run_ipfw_command(cmd, error_message) return self._pipe_counter + def _run_ipfw_command(self, command, msg=None): + """Executes a command and it prefixes the appropriate command for + Windows or Linux/UNIX. def _run_shell_command(command, msg=None): """Executes a command. @@ -162,20 +171,22 @@ def _run_shell_command(command, msg=None): command: Command list to execute. msg: Message describing the error in case the command fails. - Returns: - The standard output from running the command. - - Raises: - NetworkEmulatorError: If command fails. Message is set by the msg + Raises: + NetworkEmulatorError: If command fails. Message is set by the msg parameter. - """ - cmd_list = [str(x) for x in command] - cmd = ' '.join(cmd_list) - logging.debug('Running command: %s', cmd) + """ + if sys.platform == 'win32': + ipfw_command = ['ipfw.exe'] + else: + ipfw_command = ['sudo', '-n', 'ipfw'] - process = subprocess.Popen(cmd_list, stdout=subprocess.PIPE, - stderr=subprocess.PIPE) - output, error = process.communicate() - if process.returncode != 0: - raise NetworkEmulatorError(msg, cmd, process.returncode, output, error) - return output.strip() + cmd_list = ipfw_command[:] + [str(x) for x in command] + cmd_string = ' '.join(cmd_list) + logging.debug('Running command: %s', cmd_string) + process = subprocess.Popen(cmd_list, stdout=subprocess.PIPE, + stderr=subprocess.PIPE) + output, error = process.communicate() + if process.returncode != 0: + raise NetworkEmulatorError(msg, cmd_string, process.returncode, output, + error) + return output.strip()