#!/opt/mesosphere/bin/python

import os
import socket
import sys
import random

import dns.query


# Constants
MAX_SERVER_COUNT = 3
SPARTANS = ['198.51.100.1', '198.51.100.2', '198.51.100.3']


if len(sys.argv) != 2:
    print('Usage: gen_resolvconf.py RESOLV_CONF_PATH', file=sys.stderr)
    print('Received: {}'.format(sys.argv), file=sys.stderr)
    sys.exit(-1)
resolvconf_path = sys.argv[1]
dns_test_query = 'ready.spartan'
dns_timeout = 5


def check_server(addr):
    try:
        query = dns.message.make_query(dns_test_query, dns.rdatatype.ANY)
        result = dns.query.udp(query, addr, dns_timeout)
        if len(result.answer) == 0:
            print('Skipping DNS server {}: no records for {}'.format(
                addr, dns_test_query), file=sys.stderr)
        else:
            return True
    except socket.gaierror as ex:
        print(ex, file=sys.stderr)
    except dns.exception.Timeout:
        print('Skipping DNS server {}: no response'.format(
            addr), file=sys.stderr)
    except:
        print("Unexpected error querying DNS for server \"{}\" exception: {}".format(
            addr, sys.exc_info()[1]))

    return False


contents = """# Generated by gen_resolvconf.py. Do not edit.
# Change configuration options by changing DC/OS cluster configuration.
# This file must be overwritten regularly for proper cluster operation around
# master failure.

options timeout:1
options attempts:3

"""

if 'SEARCH' in os.environ:
    contents += "search {}\n".format(os.environ['SEARCH'])

# Check if Spartan is up
spartans_up = []
for ns in SPARTANS:
    if check_server(ns):
        spartans_up.append(ns)

if len(spartans_up) > 0:
    for ns in spartans_up:
        contents += "nameserver {}\n".format(ns)

# If Spartan is not up, fall back, and insert the upstreams
else:
    fallback_servers = os.environ['RESOLVERS'].split(',')
    random.shuffle(fallback_servers)
    for ns in fallback_servers[:MAX_SERVER_COUNT]:
        contents += "nameserver {}\n".format(ns)


# Generate the resolv.conf config
print('Updating {}'.format(resolvconf_path))
with open(resolvconf_path, 'w') as f:
    print(contents, file=sys.stderr)
    f.write(contents)

sys.exit(0)
