#!/usr/bin/env python

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

import multiprocessing
import os
import platform
import re
import shutil
import subprocess
import sys
import tempfile
import time

import options
import report

optspec = """
sample [options...]
--
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 [{0}]
m,monitor=      wireless monitor interface to use [{1}]
"""

DEVNULL = open(os.devnull, 'wb')


def Ping(host):
  code = subprocess.call(['ping', '-c', '1', host], stdout=DEVNULL,
                         stderr=DEVNULL)
  return code


def Iperf(host, path, udp=False, bandwidth=20, bind=None):
  """Run iperf against host and report results."""
  line = ['iperf', '-c', host]
  name = 'iperf'

  if udp:
    line += ['-u', '-b', str(bandwidth * 1000000)]
    name += 'u'
  if bind:
    line += ['-B', bind]

  out = open(os.path.join(path, name), 'w+')
  subprocess.check_call(line, stdout=out)
  return out


def MCSBackground(tcpdump, out):
  """Continually extract wireless MCS from a running `tcpdump` process.

  This function will not return as long as tcpdump is running. You probably want
  to run it in a background activity using multiprocessing.

  Args:
    tcpdump: a tcpdump process that's monitoring a wireless interface
                  and writing textual output to its stdout stream.
    out:          Python file-like object to write MCS information to.
  """
  mcs = re.compile(r'MCS (\d+)')
  x = 0

  for row in iter(tcpdump.stdout.readline, b''):
    x += 1
    match = mcs.search(row)
    if match:
      print >> out, '%2d ' % int(match.group(1)),
    else:
      print >> out, ' . ',

    if x % 25 == 0:
      print >> out


def MCS(bssid, interface, path):
  """Runs tcpdump in the background to gather wireless MCS."""
  print 'Please enter password for `sudo` if prompted.'
  subprocess.call(['sudo', '-v'])

  out = open(os.path.join(path, 'mcs'), 'w+')
  err = open(os.path.join(path, 'mcserr'), 'w+')

  filt = ('(not subtype beacon and not subtype ack) and '
          '(wlan addr1 {0} or wlan addr2 {0} or wlan addr3 {0})'.format(
              bssid, bssid, bssid))

  sudo_tcpdump = subprocess.Popen(['sudo', 'tcpdump', '-Z', os.getlogin(),
                                   '-Ilnei', interface, filt],
                                  stdout=subprocess.PIPE, stderr=err)
  proc = multiprocessing.Process(target=MCSBackground, args=(sudo_tcpdump, out))
  proc.start()

  return sudo_tcpdump, out, err


def IwLink(devname, path):
  out = open(os.path.join(path, 'iwlink'), 'w+')
  subprocess.check_call(['iw', 'dev', devname, 'link'], stdout=out)
  return out


def IwScan(devname, path):
  out = open(os.path.join(path, 'iwscan'), 'w+')
  subprocess.check_call(['iw', 'dev', devname, 'scan', 'dump'], stdout=out)
  if os.fstat(out.file.fileno()).st_size: return out

  subprocess.check_call(['iw', 'dev', devname, 'scan'], stdout=out)
  return out


def IpAddr(path):
  out = open(os.path.join(path, 'ipaddr'), 'w+')
  subprocess.check_call(['ip', '-o', '-f', 'inet', 'addr'], stdout=out)
  return out


def AirportI(path):
  """Gather information about the current wireless network from `airport`."""
  out = open(os.path.join(path, 'airport'), 'w+')
  subprocess.check_call(['airport', '-I'], stdout=out)
  return out


def AirportScan(path):
  """Gather information about other observable networks from `airport`."""
  out = open(os.path.join(path, 'airportscan'), 'w+')
  subprocess.check_call(['airport', '-s'], stdout=out)
  return out


def main():
  system = platform.system()
  defaults = {
      # on a modern MacBook, en0 is the AirPort and can monitor and send at once
      'Darwin': ['en0', 'en0'],
      # on Linux, separate client and monitor interfaces are needed.
      # these defaults are for a TV box, other platforms will be different.
      'Linux': ['wcli0', ''],
  }

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

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

  # 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)
  report_dir = tempfile.mkdtemp(prefix='wifi')

  # we run diagnostics, write their output to files, and gather the files into
  # a report that we present at the end of the run.
  outputs = []
  addr = ''

  if system == 'Darwin':
    ai = AirportI(report_dir)
    ai.seek(0)
    bssid = report.ParseAirportI(ai.read())['BSSID']
    outputs += [ai, AirportScan(report_dir)]
  elif system == 'Linux':
    # It's really likely we're running on a device with more than one interface.
    # Be sure we're using the one that we're trying to test.
    ip = IpAddr(report_dir)
    ip.seek(0)
    addrmap = report.ParseIpAddr(ip.read())
    addr = addrmap.get(opt.interface)
    if not addr:
      raise ValueError('Interface {0} does not have an IPv4 address.'.format(
          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]

    il = IwLink(opt.interface, report_dir)
    il.seek(0)
    bssid = report.ParseIwLink(il.read())['BSSID']

    ic = IwScan(opt.interface, report_dir)
    outputs += [ip, il, ic]
  else:
    raise OSError('This script requires Mac OS X or Linux.')

  if opt.monitor:
    sudo_tcpdump, mcs_out, mcs_err = MCS(bssid, opt.monitor, report_dir)
    print 'Gathering tcpdump in background as', sudo_tcpdump.pid
    outputs += [mcs_out, mcs_err]

  status = Ping(opt.destination)
  if not status:
    it = Iperf(opt.destination, report_dir, bind=addr)

    # Empirically about 1.25x more packets make it through in UDP than TCP.
    # Try to saturate the channel by sending a bit more than that over UDP.
    it.seek(0)
    it_iperf = report.ParseIperfTCP(it.read())

    bandwidth = it_iperf.get('bandwidth', 0.01)
    outputs += [it, Iperf(opt.destination, report_dir, udp=True,
                          bandwidth=bandwidth * 1.5,
                          bind=addr)]
  else:
    print >> sys.stderr, ('Could not ping destination host {0}; '
                          'skipping performance tests').format(opt.destination)

  if opt.monitor:
    subprocess.check_call(['sudo', 'kill', str(sudo_tcpdump.pid)])

  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

  shutil.move(report_dir, dest_dir)
  print 'Report written to', dest_dir
  for name in (o.name for o in outputs):
    print '*', os.path.basename(name)


if __name__ == '__main__':
  main()