#!/usr/bin/python

import pythonwhois
import argparse
import sys
import itertools

parser = argparse.ArgumentParser(description='Check whois and dns information.')
parser.add_argument('--domain', nargs=1, required=True, help='the domain name')
parser.add_argument('--nameserver', nargs='+', help='a space delimited list of nameservers')
args = parser.parse_args()

nsargs = args.nameserver
# how many ns servers should we expect?
nsargscount = len(nsargs)

try:
    whois = pythonwhois.get_whois(args.domain[0])
    whoiscount = []

    for key, value in whois.iteritems():
        if key == "nameservers":
            whoiscount.append(value)
except:
    #something happened, raise an erro
    print "3 Unexpected_error: ", sys.exc_info()[0]

# how many nameservers did whois return?
#whoiscount = reduce(lambda x,y: x.extend(y),whoiscount)
#print len(whoiscount)
#print nsargscount
nscount = list(itertools.chain.from_iterable(whoiscount))
#print len(nscount)

if len(nscount) < nsargscount:
   print("Extended_DNS_Check CRITICAL - WHOIS DNS servers do not match provided DNS nameservers.")
   sys.exit(2)
elif len(nscount) > nsargscount:
   print("Extended_DNS_Check WARNING - WHOIS DNS servers do not match. There are more DNS servers from the WHOIS datbase than in your query.")
   sys.exit(1)
else:
   print("Extended_DNS_Check OK - all DNS server match")
   sys.exit(0)