#!/usr/bin/python

import adns, sys
from string import join
from rblcheck import RBLQueryEngine

class QueryEngine(RBLQueryEngine):

    def submit_A(self, qname):
        from adns import rr
        if not hasattr(self, 'A_results'):
            self.A_results = {}
            self.PTR_results = {}
        self.A_results[qname] = ()
        self.submit(qname, rr.A, callback=self.A_callback)

    def A_callback(self, answer, qname, rr, flags, extra):
        if answer[0] in (adns.status.ok, adns.status.nodata,
                         adns.status.nxdomain):
            self.A_results[qname] = answer[3]
            for ip in self.A_results[qname]:
                if not self.PTR_results.has_key(ip):
                    self.submit_PTR(ip)
                    self.submit_rbl(ip)
            
    def submit_PTR(self, qname):
        from adns import rr
        if not hasattr(self, 'PTR_results'):
            self.A_results = {}
            self.PTR_results = {}
        self.PTR_results[qname] = ()
        self.submit_reverse(qname, rr.PTRraw, callback=self.PTR_callback)

    def PTR_callback(self, answer, qname, rr, flags, extra):
        if answer[0] in (adns.status.ok, adns.status.nodata,
                         adns.status.nxdomain):
            self.PTR_results[qname] = answer[3]
            for name in self.PTR_results[qname]:
                if not self.A_results.has_key(name):
                    self.submit_A(name)
            
s = adns.init(adns.iflags.noautosys)

server = None
if len(sys.argv) < 2:
    sys.stderr.write("Usage: hostmx <hostname> [<server>]\n")
    sys.exit(1)

hostname = sys.argv[1]
if len(sys.argv) > 2:
    server = sys.argv[2]

if server:
    r = s.synchronous(server, adns.rr.A)
    adns.exception(r[0])
    s_addr = r[3][0]
    config = [ "clearnameservers" ]
    config.append("nameserver %s" % s_addr)
    cfg = join(config, '\n')
    s = adns.init(adns.iflags.noautosys, sys.stderr, cfg)

s = QueryEngine(s)
r = s.synchronous(hostname, adns.rr.MXraw)
try: adns.exception(r[0])
except adns.NXDomain:
    print "%s: non-existent domain" % hostname
    sys.exit(0)
except adns.NoData:
    print "%s: no mail exchangers" % hostname

s.submit_A(hostname)

if r[1]:
    print "%s: alias for %s" % (hostname, r[1])
    s.submit_A(r[1])

for mx in r[3]:
    print "%s MX(%d): %s" % (hostname, mx[0], mx[1])
    s.submit_A(mx[1])

s.finish() # churn

listed = s.rbl_results
hosts = s.A_results
addrs = s.PTR_results

for k,v in hosts.items(): print "%s: %s" % (k, join(v))
for k,v in addrs.items(): print "%s: %s" % (k, join(v+tuple(listed[k])))

