# Script to check the details of an ah-hosted domain.
# Check NS, www adress, domain A records, MX records

import argparse

import dns.resolver

AH_SERVERS = {
    "159.203.100.138":"platypus-2",
    "128.199.219.144":"wp-01.ace-hosting.com.au",
    "172.104.183.16:":"kultarra.ace-hosting.com.au/wp-02.ace-hosting.com.au",
    "198.199.121.45":"wp-03.ace-hosting.com.au",
    "112.213.34.61":"wp-04.ace-hosting.com.au",
    "170.64.172.115":"wp-05.ace-hosting.com.au",
    "43.229.63.127":"wp-06.ace-hosting.com.au",
    "172.104.61.71": "makinmattresses.com.au",
    }


def get_ah_server_name(ip):
    try:
        name = AH_SERVERS[ip]
    except KeyError:
        name = "is not an Ace Server"

    return name

def parse_arguments():
    parser = argparse.ArgumentParser()
    parser.add_argument("-d", "--domain", type=str, required=True, help="The domain to check")
    parser.add_argument("-f", "--output-format", type=str, choices=["text", "json"], default="text", help="Output format. Text of json")

    args = parser.parse_args()

    return args

def lookup_a_record_for_mx(mx):
    answer = dns.resolver.resolve(mx, "A")[0]
    return get_ah_server_name(answer.address)

def main():
    # Use a breakpoint in the code line below to debug your script.
    args = parse_arguments()
    try:
        answers = dns.resolver.resolve(args.domain, "NS")
        for rdata in answers:
            ip = rdata.to_text()
            print(f"NS record for {args.domain} = {ip}")
    except (dns.resolver.NoAnswer, dns.resolver.NXDOMAIN):
        print(f"NS record for {args.domain} Not found")

    try:
        answers = dns.resolver.resolve(args.domain, "A")
        for rdata in answers:
            ip = rdata.to_text()
            print(f"A record for {args.domain} = {ip}({get_ah_server_name(ip)})")
    except (dns.resolver.NoAnswer, dns.resolver.NXDOMAIN):
        print(f"A record for {args.domain} Not found")

    try:
        answers = dns.resolver.resolve("www." + args.domain, "A")
        for rdata in answers:
            for rdata in answers:
                ip = rdata.to_text()
                print(f"A record for {'www.' + args.domain} = {ip}({get_ah_server_name(ip)})")
    except (dns.resolver.NoAnswer, dns.resolver.NXDOMAIN):
        print(f"A record for {'www.' + args.domain} not found")

    try:
        answers = dns.resolver.resolve(args.domain, "MX")
        for rdata in answers:
            host = rdata.exchange
            # preference = rdata.preference

            if "google" in str(host):
                print ("MX: google")
            elif "outlook" in str(host):
                print ("MX: outlook")
            else:
                print(f"MX: {host}({lookup_a_record_for_mx(host)})")
    except (dns.resolver.NoAnswer, dns.resolver.NXDOMAIN):
        print(f"MX record for {args.domain} not found")


# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    main()