#!/usr/bin/python

import subprocess
import json
import argparse
import sys

from datetime import datetime

parser = argparse.ArgumentParser()
parser.add_argument('domain', help="Required. Domain to check")
parser.add_argument('-c', '--critical', dest='critical', default=50, help="Critical threshold")
parser.add_argument('-w', '--warning', dest='warning', default=20, help="Warning threshold")
parser.add_argument('-i', '--ignore', dest='ignore', default=5, help="Ignore queues from the last X minutes (default: 5)")
args = parser.parse_args()

now = datetime.now()
p = subprocess.Popen(['/usr/sbin/postqueue', '-j'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.STDOUT)
output = str(p.stdout.read(), "utf-8").splitlines()
mail_queues = {}
for line in output:
  j = json.loads(line)
  this_queue = []
  if j["queue_name"] == 'active':
    continue
  queue_old = (now - datetime.fromtimestamp(j["arrival_time"])).total_seconds()/60
  if queue_old < args.ignore:
    # Not old enough
    continue
  for recipient in j['recipients']:
    domain = recipient['address'].split("@")[-1]
    if domain not in this_queue:
      this_queue.append(domain)
      if domain not in mail_queues:
        mail_queues[domain] = 1
      else:
        mail_queues[domain] += 1


delayed = 0
if args.domain in mail_queues:
  delayed = mail_queues[args.domain]

if delayed >= args.warning:
  print("WARNING: Queue length for %s >= %s (%s)" % (args.domain, args.warning, delayed))
  sys.exit(1)
if delayed >= args.critical:
  print("CRITICAL: Queue length for %s >= %s (%s)" % (args.domain, args.critical, delayed))
  sys.exit(2)
print("OK: Queue length for %s < %s (%s)" % (args.domain, args.warning, delayed))
sys.exit(0)