import sys
import time
import datetime
import dns
import dns.name
import dns.query

def myDig(domain, server_address):
    queryName = dns.name.from_text(domain) # Get domain for query
    query = dns.message.make_query(queryName, dns.rdatatype.A)
    queryData = dns.query.udp(query, server_address)
    tld_server = '' # ROOT SERVER -> TLD -> ANS -> CLIENT
    for additionalData in queryData.additional:
        currentData = str(additionalData)
        if 'A' in currentData:
            tld_server = currentData[(currentData.index('A') + 2): ]
            break
    queryData = dns.query.udp(query, tld_server)
    authoritative_name_server = ''
    for additionalData in queryData.additional:
        currentData = str(additionalData)
        if "IN A " in currentData:
            splitCurrentData = currentData.split(" ")
            authoritative_name_server = splitCurrentData[4]
            break
    if(authoritative_name_server):
        queryData = dns.query.udp(query, authoritative_name_server)
        stringQueryData = str(queryData)
        queryDataAnswerData = stringQueryData.split("\n")
        answerSection = queryDataAnswerData[queryDataAnswerData.index(";ANSWER") + 1]
        splitAnswerSection = answerSection.split(" ")
        print("QUESTION SECTION: ")
        print(domain + "   " + "IN A")
        print("ANSWER SECTION: ")
        print(splitAnswerSection[0] + "  " + splitAnswerSection[1]
        + "  " + splitAnswerSection[2] + "  " + splitAnswerSection[3]
        + "  " + splitAnswerSection[4])

        # for answerData in queryData.answer:
        #     queryDataAnswerData.append(answerData)
        # print(queryDataAnswerData)

def myDigTimer(domain):
    start_time = time.time()
    myDig(domain, "198.41.0.4")
    end_time = (time.time() - start_time) * 1000
    print("Query time: " + str(end_time) + " ms")
    localtime = time.asctime( time.localtime(time.time())) # From TutorialsPoint
    print("WHEN: " + str(localtime))

if __name__ == "__main__":
    myDigTimer(sys.argv[1])