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 platform
import subprocess
import sys
import time

from fabric.api import env, execute, local, run
from fabric.network import disconnect_all

import ifstats
import iperf
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
r,remote=       remote host to run tests on
"""


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:
    with open(opt.journal, 'a') as journal:
      print >> journal, report_name
    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
    env.update({
        'key_filename': os.path.expanduser('~/.ssh/bruno-sshkey'),
        'user': 'root',
        'shell': 'sh -l -c',
    })
    execute_args['host'] = opt.remote

    ifsystem = 'Linux'
    ifrun = run

    atexit.register(disconnect_all)
  else:
    ifsystem = system
    ifrun = functools.partial(local, capture=True)

  ifs = ifstats.InterfaceStats(system=ifsystem,
                               runner=ifrun,
                               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 = 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:
    sudo_tcpdump, mcs_out, mcs_err = tcpdump.MCS(bssid, opt.monitor, dest_dir)
    print 'Gathering tcpdump in background as', sudo_tcpdump.pid

  ips = iperf.Iperf(runner=ifrun, bind=addr)
  env.warn_only = True  # `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
  cache.update(
      execute(ips.RunTestSeries, opt.destination, **execute_args).values()[0]
  )

  if opt.monitor:
    subprocess.check_call(['sudo', 'kill', str(sudo_tcpdump.pid)])
    for stream in [mcs_out, mcs_err]:
      stream.flush()
      stream.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()