resolve.py - simple recursive DNS resolver

resolve.py

from dnslib import DNSRecord,DNSError,QTYPE

ROOT_NS='198.41.0.4'

def find_cname(a,qtype,cname):
    for r in a.rr:
        if QTYPE[r.rtype] == qtype and r.rname == cname:
            return str(r.rdata)
        elif QTYPE[r.rtype] == 'CNAME' and r.rname == cname:
            return find_cname(a,qtype,str(r.rdata))
    return resolve(str(r.rdata),qtype,ROOT_NS)

def resolve(name,qtype='A',ns=ROOT_NS):
    print(f"dig -r @{ns} {name} {qtype}")
    q = DNSRecord.question(name,qtype)
    a = DNSRecord.parse(q.send(ns))
    if q.header.id != a.header.id:
        raise DNSError('Response transaction id does not match query transaction id')
    for r in a.rr:
        if QTYPE[r.rtype] == qtype and r.rname == name:
            return str(r.rdata)
        elif QTYPE[r.rtype] == 'CNAME' and r.rname == name:
            return find_cname(a,qtype,str(r.rdata))
    for r in a.ar:
        if QTYPE[r.rtype] == 'A':
            return resolve(name,qtype,str(r.rdata))
    for r in a.auth:
        if QTYPE[r.rtype] == 'NS':
            return resolve(name,qtype,resolve(str(r.rdata),'A',ROOT_NS))
    raise ValueError("Cant resolve domain")

if __name__ == '__main__':
    import sys
    def get_arg(i,d):
        try: return sys.argv[i]
        except IndexError: return d
    print(resolve(get_arg(1,"google.com"), get_arg(2,"A"), get_arg(3,ROOT_NS)))