#!/usr/bin/python

import subprocess
import time
import argparse
import sys
import tempfile
import smtplib
import time
import re
from email.mime.text import MIMEText

def getDiskUtil(disk):
  ''' The average time (in milliseconds) for I/O requests issued to the device to be served.'''
  disk = disk[:-1]
  iostat = subprocess.check_output(["iostat","-x",disk, "1", "3"])
  for line in iostat.split('\n'):
    if disk.split('/')[2] in line:
      await = line.split()[9]
  return float(await)

def sendAlert(mount, usage, await):
  tf = tempfile.NamedTemporaryFile(delete=True)
  with open(tf.name, "w") as fp:
    fp.write("Disk alert on volume {} usage {} await {}\n".format(mount, usage, await))
    fp.close()

    with open(tf.name, "r") as fp:
      msg = MIMEText(fp.read())
      msg['Subject'] = 'Disk utilization alert'
      msg['From'] = "alert@testing.aws.orchardplatform.com"
      msg['To'] = "rodrick.brown@gmail.com"
      s = smtplib.SMTP('localhost')
      s.sendmail("alert@testing.aws.orchardplatform.com", ["rodrick.brown@gmail.com"], msg.as_string())
      s.quit()

def logWriter(logfile,mount,usage,await):
  timestamp = time.strftime("%d/%m/%y %H:%M:%S")
  with open(logfile.name, "a+") as fp:
    timestamp = time.strftime("%d/%m/%y %H:%M:%S")
    fp.write("Volume {} went critical usage: {} await: {} at {}\n".format(mount,usage,await,timestamp))

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='Process disk checks')
  parser.add_argument('--diskusage', action="store", dest="diskusage", type=float)
  parser.add_argument('--iorequest', action="store", dest="iorequest", type=float)

  args = parser.parse_args()
  if not args.diskusage:
    sys.exit(-1)

  logfile = tempfile.NamedTemporaryFile(delete=False)
  print("--- Writing alerts to {} ---".format(logfile.name))
  while True:
    df = subprocess.check_output(["df", "-m"])
    alert_triggers = set()

    for line in df.split('\n'):
      if line.startswith('/dev'):
        device = line.split()[0]
        usage = line.split()[4]
        mount = line.split()[5]
        usage_value = float(usage[:-1])

        await = getDiskUtil(device)
        if await > float(args.iorequest):
          print("io threshold triggered {} {}".format(await, device))
          logWriter(logfile,mount,usage,await)
          alert_triggers.add(mount)
          sendAlert(mount, usage, await)

        if usage_value > float(args.diskusage):
           print("disk usage threshold triggered: {} {} {}".format(usage_value, args.diskusage, device))
           logWriter(logfile,mount,usage,await)
           alert_triggers.add(mount)
           sendAlert(mount, usage, await)

        else:
           print("mount - {} usage: {} await: {} OK".format(mount, usage, await))
           if mount in alert_triggers:
             with open(logfile.name,"a+") as fp:
               fp.write("mount - {} usage: {} await: {} recovered".format(mount,usage,await))
             logWriter(logfile,mount,usage,await)
    time.sleep(3)

[rodrick@localhost ~]$ ./df.py --diskusage 20 --iorequest .1
--- Writing alerts to /tmp/tmpXQpZi6 ---
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
io threshold triggered 1.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1
disk usage threshold triggered: 43.0 20.0 /dev/xvda1