#!/usr/bin/python3

import argparse
import json
import re
import sys
from urllib.parse import urlparse
import DNS  # type: ignore


def get_dns_pool_ips(hostname, protocol, resolver):
    """Return IPs in the DNS pool"""
    dnsreq = DNS.Request(hostname, qtype='A', protocol=protocol,
                         server=resolver)
    dnsres = dnsreq.req()
    return [answer['data'] for answer in dnsres.answers]


def is_in_the_dns_pool(url_prefix, ips, protocol, resolver):
    """Return True iff. any of the specified mirror's IP addresses
    is in the DNS pool"""
    mirror_ips_req = DNS.Request(
        urlparse(url_prefix).hostname,
        qtype='A',
        protocol=protocol,
        server=resolver
    )
    mirror_ips_res = mirror_ips_req.req()
    mirror_ips = [answer['data'] for answer in mirror_ips_res.answers]
    for mirror_ip in mirror_ips:
        if mirror_ip in ips:
            return True
    return False


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="\
    If --url-prefix is used, tell if the specified mirror is in the DNS pool.\
    Otherwise, list the mirrors that are in the DNS pool.\
    This program uses Google's DNS over TCP by default, and is meant to be\
    used with torsocks.\
    ")
    parser.add_argument("--url-prefix",
                        help="URL prefix of the mirror to check")
    parser.add_argument("--resolver", default='8.8.8.8',
                        help="Nameserver to use")
    parser.add_argument("--protocol", default='tcp',
                        help="Protocol to use when querying the nameserver")
    args = parser.parse_args()
    dns_pool_ips = get_dns_pool_ips(
        'dl.amnesia.boum.org',
        protocol=args.protocol,
        resolver=args.resolver
    )
    if args.url_prefix:
        assert re.compile(r'^https?\:\/\/[\w\-\.]+').match(args.url_prefix)
        if is_in_the_dns_pool(
                args.url_prefix,
                dns_pool_ips,
                protocol=args.protocol,
                resolver=args.resolver):
            print(f"Mirror {args.url_prefix} is in the DNS pool")
            sys.exit(0)
        else:
            print(f"Mirror {args.url_prefix} is not in the DNS pool")
            sys.exit(1)
    else:
        with open('mirrors.json', 'r', encoding='utf-8') as f:
            mirrors = [
                mirror
                for mirror in json.loads(f.read())['mirrors']
                if mirror['weight'] > 0
            ]
        for mirror in mirrors:
            if is_in_the_dns_pool(
                    mirror['url_prefix'],
                    dns_pool_ips,
                    protocol=args.protocol,
                    resolver=args.resolver):
                print(mirror)
