#!/usr/bin/env python
import argparse
from sys import exit
import envoy

OK=0
WARN=1
CRIT=2
UNKN=3

def clean_dig_results(input):
    results=[]

    # Clean results from dig
    for line in input.splitlines():
        line   = line.lower()
        result = line.rstrip(' \t.')
        results.append(result)

    return results


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("-H", "--host",    dest="host", required=True)
    parser.add_argument("-v", "--verbose", dest="verbose", default=False, action="store_true")
    args = parser.parse_args()

    if args.verbose:
        print "Checking: %s" % args.host

    dig = envoy.run('dig +short %s' % args.host)
    if dig.status_code == 9:
        print "UNKN: DNS server timeout"
        exit(UNKN)

    if dig:
        ips = clean_dig_results(dig.std_out)
        if len(ips) == 0:
            print "CRIT: Can't find DNS entry for '%s'" % args.host
            exit(CRIT)

        for ip in ips:
            # Use dig to get reverse records
            dig_x = envoy.run('dig +short -x %s' % ip)
            if dig_x.status_code == 9:
                print "UNKN: DNS server timeout"
                exit(UNKN)

            results = clean_dig_results(dig_x.std_out)

            if len(results) > 1:
                # Format results as quoted, comma seperated list
                formatted_results = ','.join(map(lambda x: repr(x), results))
                print "CRIT: Host '%s' has too many records - [%s]" % (args.host, formatted_results)
                exit(CRIT)
            elif len(results) == 0:
                print "CRIT: Host '%s' is missing reverse dns" % args.host
                exit(CRIT)
            elif results[0] != args.host:
                print "CRIT: Reverse dns result '%s' does not match forward dns '%s'" % (results[0], args.host)
                exit(CRIT)
    else:
        print "CRIT: Can't get DNS for '%s'" % (args.host)
        exit(CRIT)

    formatted_ips = ','.join(map(lambda x: repr(x), ips))
    print "OK: Reverse DNS for '%s' matches: [%s]" % (args.host, formatted_ips)
    exit(OK)

if __name__ == "__main__":
    main()
