wifitables/sample.py - vendor/google/test - Git at Google

 #!/usr/bin/python2.7
 """sample: measure wireless performance and write a report to the filesystem."""

 import atexit
 import collections
 import functools
 import logging
 import os
 import pipes
 import platform
 import subprocess
 import sys
 import time

 from fabric import api
 from fabric import network
 from fabric import utils

 import ifstats
 import ifstats.skids as ifstats_skids
 import iperf
 import isostream
 import options
 import tcpdump

 optspec = """
 sample [options...]
 --
 B,bind=              interface IP to bind during iperf runs
 d,destination=       host to run tests against [192.168.1.143]
 s,steps=             number of steps test was run from [10]
 j,journal=           append to journal tracking a series of test runs
 i,interface=         interface to use for outgoing connections
 m,monitor            try to monitor wireless traffic using `tcpdump`
 M,monitor-interface= override default interface for monitoring traffic.
 no-filter            don't attempt to filter packets seen on monitor interface
 r,remote=            comma-separated list of remote hosts to run tests on
 S,skid=              IP address of the skid used to reach the network
 t,time=              length of time in seconds to run isostream test for [50]
 """


 class FabricLoggingHandler(logging.Handler):

   def emit(self, record):
     msg = self.format(record)
     if record.levelno >= logging.ERROR:
       utils.error(msg)
     elif record.levelno == logging.WARNING:
       utils.warn(msg)
     else:
       utils.puts(msg)


 def _Run(cmd_or_args, dirname='.', local=False):
   if isinstance(cmd_or_args, list):
     cmd = ' '.join([pipes.quote(arg) for arg in cmd_or_args])
   else:
     cmd = cmd_or_args

   if local:
     with api.lcd(dirname):
       return api.local(cmd, capture=True)
   else:
     with api.cd(dirname):
       return api.run(cmd)


 def _SetupTestHost():
   # work around current embedded image's lack of mktemp(1).
   wd = api.run('python -c "import tempfile; print tempfile.mkdtemp()"')
   with api.cd(wd):
     for script in ['timeout']:
       api.put(script, script, mode=0755)
   return wd


 def _CleanupTestHost(dirname):
   api.run('rm -r {}'.format(pipes.quote(dirname)))
   network.disconnect_all()


 def main():
   # Structured logging. Helps to pick up Paramiko (SSH transport layer) errors.
   root = logging.getLogger()
   root.setLevel(logging.INFO)
   root.addHandler(FabricLoggingHandler())

   o = options.Options(optspec)
   (opt, _, extra) = o.parse(sys.argv[1:])
   if extra:
     o.fatal('did not understand supplied extra arguments.')

   # Build a directory for the report to live in.
   report_name = 'wifi-{}-{:04}'.format(time.time(), opt.steps)
   if opt.journal:
     dest_dir = os.path.join(os.path.dirname(opt.journal), report_name)
   else:
     dest_dir = report_name

   if not os.path.exists(dest_dir):
     os.makedirs(dest_dir)

   logging.info('Report being written to %s', dest_dir)

   # we run diagnostics, write their output to files, and gather the files into
   # a report that we present at the end of the run.
   cache = collections.defaultdict(dict)
   local_run = functools.partial(_Run, local=True)

   # always gather network statistics for the system we're on
   system = platform.system()
   lifs = ifstats.InterfaceStats(system=system, runner=local_run)

   for host, output in api.execute(lifs.Gather).items():
     cache[host].update(output)
   lresults = ifstats.Parse(system, cache['<local-only>'])
   bssid = lresults.get('link', {}).get('BSSID')

   if opt.skid:
     logging.info('Getting stats from skid %s', opt.skid)
     try:
       status_wireless = ifstats_skids.GetStatusWireless(opt.skid)
       cache[opt.skid]['status_wireless'] = status_wireless
     except IOError as e:
       logging.warning('Got IOError while connecting to skid %s: %s',
                       opt.skid, e)

   if opt.remote:
     # update Fabric env for embedded systems
     api.env.update({
         'always_use_pty': False,
         'key_filename': os.path.expanduser('~/.ssh/bruno-sshkey'),
         'user': 'root',
         'shell': 'sh -l -c',
     })
     hosts = opt.remote.split(',')

     wd = api.execute(_SetupTestHost, hosts=hosts).values()[0]
     atexit.register(api.execute, _CleanupTestHost, wd, hosts=hosts)

     ifsystem = 'Linux'
     run = functools.partial(_Run, dirname=wd, local=False)

     # Gather network statistics from remote system
     ifs = ifstats.InterfaceStats(system=ifsystem, runner=run,
                                  interface=opt.interface)
     for host, output in api.execute(ifs.Gather, hosts=hosts).items():
       cache[host].update(output)

     # try to compute a bind address for iperf automatically.
     results = ifstats.Parse(system, cache[opt.remote])
     addr = results.get('addr', {}).get(opt.interface, '')
   else:
     ifsystem = system
     run = local_run
     hosts = []

     # try to compute a bind address for iperf automatically.
     addr = lresults.get('addr', {}).get(opt.interface, '')

   # because ip addr usually includes a subnet mask, which will prevent iperf
   # from binding to the address
   mask = addr.find('/')
   if mask > -1:
     addr = addr[:mask]

   if opt.monitor:
     defaults = {
         # on a modern MacBook, en0 is the AirPort and can both monitor and send
         'Darwin': 'en0',
         # on Linux, separate client and monitor interfaces are needed; and the
         # monitor interface must be manually created
         'Linux': 'moni0',  # from the `iw` documentation
     }

     if opt.monitor_interface:
       monif = opt.monitor_interface
     else:
       monif = defaults.get(system)

     if not monif:
       logging.error('Cannot gather tcpdump: no monitor interface specified '
                     'and autodetect failed.')

     tcpdump_proc, tcpdump_stderr = tcpdump.Tcpdump(monif, bssid, dest_dir,
                                                    opt.filter)
     logging.info('Gathering tcpdump in background as %d', tcpdump_proc.pid)

   if opt.bind and not addr:
     addr = opt.bind

   ips = iperf.Iperf(runner=run, bind=addr)

   # `iperf` returns 56 if it can't reach the server, or 57 if it doesn't receive
   # a final report from it on Linux; don't abort in these cases
   with api.settings(warn_only=True):
     for host, output in api.execute(ips.RunTestSeries, opt.destination,
                                     hosts=hosts).items():
       cache[host].update(output)

   # `isostream` won't end on its own, so we wrap it with `timeout` and accept a
   # return code of 124 (timed out) as well.
   with api.settings(ok_ret_codes=[0, 124]):
     for host, output in api.execute(isostream.RunIsostreamTest, run,
                                     opt.destination, time=opt.time,
                                     hosts=hosts).items():
       cache[host]['isostream'] = output

   if opt.monitor:
     try:
       tcpdump_proc.terminate()
     except OSError:
       subprocess.check_call(['sudo', 'kill', str(tcpdump_proc.pid)])

     tcpdump_stderr.flush()
     tcpdump_stderr.close()

   if opt.journal:
     with open(opt.journal, 'a') as journal:
       print >> journal, report_name

   # TODO(willangley): write out correctly if more than one host
   for host, items in cache.items():
     for name, value in items.items():
       with open(os.path.join(dest_dir, name), 'w+b') as outfile:
         outfile.write(value)


 if __name__ == '__main__':
   main()
	#!/usr/bin/python2.7
	"""sample: measure wireless performance and write a report to the filesystem."""

	import atexit
	import collections
	import functools
	import logging
	import os
	import pipes
	import platform
	import subprocess
	import sys
	import time

	from fabric import api
	from fabric import network
	from fabric import utils

	import ifstats
	import ifstats.skids as ifstats_skids
	import iperf
	import isostream
	import options
	import tcpdump

	optspec = """
	sample [options...]
	--
	B,bind= interface IP to bind during iperf runs
	d,destination= host to run tests against [192.168.1.143]
	s,steps= number of steps test was run from [10]
	j,journal= append to journal tracking a series of test runs
	i,interface= interface to use for outgoing connections
	m,monitor try to monitor wireless traffic using `tcpdump`
	M,monitor-interface= override default interface for monitoring traffic.
	no-filter don't attempt to filter packets seen on monitor interface
	r,remote= comma-separated list of remote hosts to run tests on
	S,skid= IP address of the skid used to reach the network
	t,time= length of time in seconds to run isostream test for [50]
	"""


	class FabricLoggingHandler(logging.Handler):

	def emit(self, record):
	msg = self.format(record)
	if record.levelno >= logging.ERROR:
	utils.error(msg)
	elif record.levelno == logging.WARNING:
	utils.warn(msg)
	else:
	utils.puts(msg)


	def _Run(cmd_or_args, dirname='.', local=False):
	if isinstance(cmd_or_args, list):
	cmd = ' '.join([pipes.quote(arg) for arg in cmd_or_args])
	else:
	cmd = cmd_or_args

	if local:
	with api.lcd(dirname):
	return api.local(cmd, capture=True)
	else:
	with api.cd(dirname):
	return api.run(cmd)


	def _SetupTestHost():
	# work around current embedded image's lack of mktemp(1).
	wd = api.run('python -c "import tempfile; print tempfile.mkdtemp()"')
	with api.cd(wd):
	for script in ['timeout']:
	api.put(script, script, mode=0755)
	return wd


	def _CleanupTestHost(dirname):
	api.run('rm -r {}'.format(pipes.quote(dirname)))
	network.disconnect_all()


	def main():
	# Structured logging. Helps to pick up Paramiko (SSH transport layer) errors.
	root = logging.getLogger()
	root.setLevel(logging.INFO)
	root.addHandler(FabricLoggingHandler())

	o = options.Options(optspec)
	(opt, _, extra) = o.parse(sys.argv[1:])
	if extra:
	o.fatal('did not understand supplied extra arguments.')

	# Build a directory for the report to live in.
	report_name = 'wifi-{}-{:04}'.format(time.time(), opt.steps)
	if opt.journal:
	dest_dir = os.path.join(os.path.dirname(opt.journal), report_name)
	else:
	dest_dir = report_name

	if not os.path.exists(dest_dir):
	os.makedirs(dest_dir)

	logging.info('Report being written to %s', dest_dir)

	# we run diagnostics, write their output to files, and gather the files into
	# a report that we present at the end of the run.
	cache = collections.defaultdict(dict)
	local_run = functools.partial(_Run, local=True)

	# always gather network statistics for the system we're on
	system = platform.system()
	lifs = ifstats.InterfaceStats(system=system, runner=local_run)

	for host, output in api.execute(lifs.Gather).items():
	cache[host].update(output)
	lresults = ifstats.Parse(system, cache['<local-only>'])
	bssid = lresults.get('link', {}).get('BSSID')

	if opt.skid:
	logging.info('Getting stats from skid %s', opt.skid)
	try:
	status_wireless = ifstats_skids.GetStatusWireless(opt.skid)
	cache[opt.skid]['status_wireless'] = status_wireless
	except IOError as e:
	logging.warning('Got IOError while connecting to skid %s: %s',
	opt.skid, e)

	if opt.remote:
	# update Fabric env for embedded systems
	api.env.update({
	'always_use_pty': False,
	'key_filename': os.path.expanduser('~/.ssh/bruno-sshkey'),
	'user': 'root',
	'shell': 'sh -l -c',
	})
	hosts = opt.remote.split(',')

	wd = api.execute(_SetupTestHost, hosts=hosts).values()[0]
	atexit.register(api.execute, _CleanupTestHost, wd, hosts=hosts)

	ifsystem = 'Linux'
	run = functools.partial(_Run, dirname=wd, local=False)

	# Gather network statistics from remote system
	ifs = ifstats.InterfaceStats(system=ifsystem, runner=run,
	interface=opt.interface)
	for host, output in api.execute(ifs.Gather, hosts=hosts).items():
	cache[host].update(output)

	# try to compute a bind address for iperf automatically.
	results = ifstats.Parse(system, cache[opt.remote])
	addr = results.get('addr', {}).get(opt.interface, '')
	else:
	ifsystem = system
	run = local_run
	hosts = []

	# try to compute a bind address for iperf automatically.
	addr = lresults.get('addr', {}).get(opt.interface, '')

	# because ip addr usually includes a subnet mask, which will prevent iperf
	# from binding to the address
	mask = addr.find('/')
	if mask > -1:
	addr = addr[:mask]

	if opt.monitor:
	defaults = {
	# on a modern MacBook, en0 is the AirPort and can both monitor and send
	'Darwin': 'en0',
	# on Linux, separate client and monitor interfaces are needed; and the
	# monitor interface must be manually created
	'Linux': 'moni0', # from the `iw` documentation
	}

	if opt.monitor_interface:
	monif = opt.monitor_interface
	else:
	monif = defaults.get(system)

	if not monif:
	logging.error('Cannot gather tcpdump: no monitor interface specified '
	'and autodetect failed.')

	tcpdump_proc, tcpdump_stderr = tcpdump.Tcpdump(monif, bssid, dest_dir,
	opt.filter)
	logging.info('Gathering tcpdump in background as %d', tcpdump_proc.pid)

	if opt.bind and not addr:
	addr = opt.bind

	ips = iperf.Iperf(runner=run, bind=addr)

	# `iperf` returns 56 if it can't reach the server, or 57 if it doesn't receive
	# a final report from it on Linux; don't abort in these cases
	with api.settings(warn_only=True):
	for host, output in api.execute(ips.RunTestSeries, opt.destination,
	hosts=hosts).items():
	cache[host].update(output)

	# `isostream` won't end on its own, so we wrap it with `timeout` and accept a
	# return code of 124 (timed out) as well.
	with api.settings(ok_ret_codes=[0, 124]):
	for host, output in api.execute(isostream.RunIsostreamTest, run,
	opt.destination, time=opt.time,
	hosts=hosts).items():
	cache[host]['isostream'] = output

	if opt.monitor:
	try:
	tcpdump_proc.terminate()
	except OSError:
	subprocess.check_call(['sudo', 'kill', str(tcpdump_proc.pid)])

	tcpdump_stderr.flush()
	tcpdump_stderr.close()

	if opt.journal:
	with open(opt.journal, 'a') as journal:
	print >> journal, report_name

	# TODO(willangley): write out correctly if more than one host
	for host, items in cache.items():
	for name, value in items.items():
	with open(os.path.join(dest_dir, name), 'w+b') as outfile:
	outfile.write(value)


	if __name__ == '__main__':
	main()