#!/usr/bin/python2

'''This is a script which runs an MTR (matt's traceroute)
against a target hosts - meant to be triggered by a
smokeping alert. Output is emailed and saved to log.'''

import argparse, datetime, subprocess, smtplib, socket
from email.mime.multipart import MIMEMultipart
from email.mime.text import MIMEText

MAIL_FROM  = 'SmokeMTR <root@mistral.server-speed.net>'
MAIL_TO    = 'root@mistral.server-speed.net'
LOG_FILE   = '/tmp/smokealert.out'

# leave empty if you don't want it to be used
SENDMAIL   = '/usr/bin/sendmail'
# leave empty if you don't want it to be used
SMTP_HOST  = ''

# mtr binary path
MTR        = '/usr/sbin/mtr'
# how many pings to send
MTR_CYCLES = 20

def parse_args():
  """Define the argparse parameters to accept for smokeping args:
  name-of-alert, target, loss-pattern, rtt-pattern, hostname, raised"""
  parser = argparse.ArgumentParser(description='Example with non-optional arguments')
  parser.add_argument('alert', action="store")
  parser.add_argument('target', action="store")
  parser.add_argument('loss', action="store")
  parser.add_argument('rtt', action="store")
  parser.add_argument('host', action="store")
  parser.add_argument('raised', action="store", default=None, nargs="?")
  return(parser.parse_args())

def get_ips_for_target(host):
    """Get a list of IPs for a host"""

    ips = []

    for entry in socket.getaddrinfo(host, None):
        ip = entry[4][0]
        if ip not in ips:
            ips.append(ip)

    return ips

def mtr_host(parsed):
  """Run a matt's traceroute against the problematic host and append the results to
  a list which will be used to log to file and send email"""

  ips = get_ips_for_target(parsed.host)

  m = []
  m.append( '===== %s =====\n' % str(datetime.datetime.now())[:19] )
  m.append( 'Alert: %s\n' % parsed.alert )
  m.append( 'Loss: %s\n' % parsed.loss )
  m.append( 'Round Trip Time: %s\n' % parsed.rtt )
  m.append( '%s: %s\n\n' % (parsed.target, ", ".join(ips)) )

  cmd = [MTR, '--report-wide', '--report', '--show-ips', '--report-cycles', str(MTR_CYCLES)]

  procs = []
  for ip in ips:
      procs.append({
          "proc": cmd_run(cmd + [ip]),
          "ip": ip,
      })

  for entry in procs:
      entry["proc"].wait()
      m.append( 'MTR to: %s\n' % entry["ip"] )
      m.append(entry["proc"].communicate()[0]+"\n")

  extra = ""
  if parsed.raised == '0':
      extra = "cleared"
  elif parsed.raised == '1':
      extra = "raised"

  mail_alert(subject="[SmokeAlert] %s %s %s" % (parsed.host, parsed.alert, extra), body=''.join(m))

  f = open(LOG_FILE, 'a')
  f.write( "".join(m) + '\n' )
  f.close()

def cmd_run(args, **kwds):
  """cmd_run uses the subprocess lib to run a system command.
  There are easier ways to this but apparently this is the 'best practice'
  method according to: http://goo.gl/LfYhP"""
  kwds.setdefault("stdout", subprocess.PIPE)
  kwds.setdefault("stderr", subprocess.STDOUT)
  p = subprocess.Popen(args, **kwds)
  return p

def mail_alert(subject="Smokeping alert", body=""):
  '''Generate an HTML and TXT encoded mail with given subject and body'''

  # Create message container - the correct MIME type is multipart/alternative.
  msg = MIMEMultipart('alternative')
  msg['Subject'] = subject
  msg['From'] = MAIL_FROM
  msg['To'] = MAIL_TO

  # Create the body of the message (a plain-text and an HTML version).
  text = body
  html = """<pre>%s</pre>""" % body

  # Record the MIME types of both parts - text/plain and text/html.
  part1 = MIMEText(text, 'plain')
  part2 = MIMEText(html, 'html')

  # Attach parts into message container.
  # According to RFC 2046, the last part of a multipart message, in this case
  # the HTML message, is best and preferred.
  msg.attach(part1)
  msg.attach(part2)

  #print(msg.as_string())

  if SENDMAIL != "":
      sendmail = subprocess.Popen([SENDMAIL, MAIL_TO], stdin=subprocess.PIPE)
      sendmail.communicate(msg.as_string())

  if SMTP_HOST != "":
      # Send the message via local SMTP server.
      s = smtplib.SMTP(SMTP_HOST)
      # sendmail function takes 3 arguments: sender's address, recipient's address
      # and message to send - here it is sent as one string.
      s.sendmail(MAIL_FROM, MAIL_TO, msg.as_string())
      s.quit()

if __name__ == '__main__':
  mtr_host(parse_args())