#!/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)

  # On a Mac, announce the end of tests.
  subprocess.check_call(['which say && say "Done"'], shell=True)


if __name__ == '__main__':
  main()