import subprocess
import re # YAY REVERSE ENGINEER... wait its just regex
import readline
#Made this work with linux magic, test

#Define Colours
class bcolors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

#FIXME: Have any CNAMES returned do a lookup and print it after the record.

inInput = ""
partDomain = ""
running = 1
pingOn = "P"
whoOn = "W"
phrases = ["ping", "whois", "quit", "clear", "help", "nameserver"]
   
def printHelp():
    print bcolors.HEADER + "========================================================"
    print "=                 help - This screen                   ="
    print "=              clear - Clears the screen               ="
    print "=    ping - Enables/Disables pinging the IP Address    ="
    print "========================================================"
    print "========================================================"

def compareSOA(inNameservers, inDomain, origSOA):
    dnsServer = '8.8.8.8'                
    soaRecords = []
    count = 0
    outSOA = ''
    returnmsg = "SOA Good"
    for nameserver in inNameservers:
        if nameserver != '':
            dnsServer = "@" + nameserver
        try:
            outSOA = subprocess.check_output(["dig", "SOA", dnsServer, "+short", inDomain])
        except subprocess.CalledProcessError as e:
            output = e.output
            returncode = e.returncode


        outSOAa = outSOA.split(' ')
        #print outSOAa[2]
        if len(outSOAa) != 1:
            soaRecords.append(outSOAa[2])
            for soa in soaRecords:
                if soa != soaRecords[0]:
                    return "SOA Records don't Match"
    return "SOA good"        

    #print soaRecords

def checkIfIP(in1):
    count = 0
    if len(in1) != 4:
        return 0
    for dom in in1:
        if dom.isdigit():
            count += 1
        else:
            return 0
    if count == 4:
        return 1


def checkPhrase(inputs):
    for word in phrases:
        if inputs == word:
            options(inputs)
            return 1
    return 0

def options(inputs):
    global whoOn
    global pingOn

    if inputs == "ping":
        if pingOn == "P":
            pingOn = "p"
        else:
            pingOn = "P"
    if inputs == "whois":
        if whoOn == "W":
            whoOn = "w"
        else:
            whoOn = "W"
    if inputs == "clear":
        subprocess.call("clear")
    if inputs == "help":
        printHelp()

def iterate(inArray, pre):
    count = 0
    while count != len(inArray) - 1:
        print pre, inArray[count]
        count = count + 1
def parseWhois(whioisInfo):
    NS = re.compile("^[Nn]")
    D = re.compile("^[Dd]")
    R = re.compile("^[Rr]")
    U = re.compile("^[Uu]")

    for str in whioisInfo:
        if NS.match(str) or D.match(str) or R.match(str) or U.match(str):
            print str

def whois(inDomain):
    lookup = subprocess.check_output(["whois", inDomain])
    print "========================================================"
    test = lookup.split("\n")
    #iterate(test, "")
    parseWhois(test)

def dig(inDomain, record):
    dig = subprocess.check_output(["dig", record, "+short", inDomain])
    print "========================================================" 
    test = dig.split("\n")
    iterate(test, record)
    return test

def ipLookup(inDomain):
    host = "N/A"
    ipLook = subprocess.check_output(["geoiplookup", inDomain])
    try:
        host = subprocess.check_output(["host", inDomain])
        returncode = 0
    except subprocess.CalledProcessError as e:
        output = e.output
        returncode = e.returncode

    print '\n', ipLook, '\n', "PTR: ", host 

def ipPing(inDomain):
    ping = subprocess.call(["ping", "-c 2", inDomain])
    #iterate(ping, "")

def getInput():
    global inInput 
    global partDomain
    global running
    inInput = raw_input("["+ pingOn +"][" + whoOn + "] Please enter a domain/IP: ")
    partDomain = inInput.split(".")
    if inInput == "quit":
        running = 0

#Main Loop
while running == 1:
    getInput()
    if checkPhrase(inInput) == 1 or len(partDomain) == 1:
        pass
    else:
        if running == 0:
            pass
        else:
            if checkIfIP(partDomain):
                print "========================================================"
                ipLookup(inInput)
                if pingOn == "P":
                    print "========================================================"
                    ipPing(inInput)
                else:
                    
                    
            else:
                if whoOn == "W":
                    whois(inInput)
                #Condense dig calls
                dig(inInput, "A")
                dig(inInput, "MX")
                dig(inInput, "TXT")
                test = dig(inInput, "NS")
                dig(inInput, "SRV")
                soaRec = dig(inInput, "SOA")
                print compareSOA(test, inInput, soaRec)
                print '\n'