wifitables/sample.py

#!/usr/bin/env python

"""sample: measure wireless performance and write a report to the filesystem."""

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

from fabric import api
from fabric import network

import ifstats
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=    wireless interface to use for outgoing connections
m,monitor=      wireless monitor interface to use
no-filter       don't attempt to filter packets seen on monitor interface
r,remote=       remote host to run tests on
t,time=         length of time in seconds to run isostream test for [50]
"""


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():
  system = platform.system()
  defaults = {
      # on a modern MacBook, en0 is the AirPort and can monitor and send at once
      'Darwin': {
          'monitor': 'en0'
      },
      # on Linux, separate client and monitor interfaces are needed; and the
      # monitor interface must be manually created
      'Linux': {
          'monitor': 'moni0'  # from the `iw` documentation
      },
  }

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

  if not defaults.get(system):
    raise OSError('Running on unsupported system {0}; '
                  'supported systems are {1}'.format(system,
                                                     ' '.join(defaults.keys())))

  # Pick the report name before we run it, so it can be consistent across
  # multiple systems.
  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

  print 'Report being written to', dest_dir
  if not os.path.exists(dest_dir):
    os.makedirs(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.
  #
  # entries collects Entry objects that have not been written to disk yet.
  execute_args = {}
  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',
    })
    execute_args['host'] = opt.remote

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

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

  ifs = ifstats.InterfaceStats(system=ifsystem,
                               runner=run,
                               interface=opt.interface)

  # since we're only executing over one host, ignore the return from `execute`
  # that says which host it was for now.
  cache = api.execute(ifs.Gather, **execute_args).values()[0]
  results = ifstats.Parse(ifsystem, cache)

  bssid = results['link']['BSSID']
  addr = results.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:
    tcpdump_proc, tcpdump_stderr = tcpdump.tcpdump(bssid, opt.monitor, dest_dir,
                                                   opt.filter)
    print 'Gathering tcpdump in background as', 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):
    cache.update(
        api.execute(ips.RunTestSeries, opt.destination,
                    **execute_args).values()[0]
    )

  # `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]):
    cache['isostream'] = api.execute(isostream.RunIsostreamTest, run,
                                     opt.destination, time=opt.time,
                                     **execute_args).values()[0]

  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

  for name, value in cache.items():
    with open(os.path.join(dest_dir, name), 'w+b') as outfile:
      outfile.write(value)


if __name__ == '__main__':
  main()