#!/usr/bin/env python3
#
# Copyright (C) 2018 VyOS maintainers and contributors
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
#

import sys
import os
import time

from vyos.config import Config
from vyos.util import ConfigError

config_file = r'/etc/dnsmasq.d/vyos.conf'

# borrowed from: https://github.com/donjajo/py-world/blob/master/resolvconfReader.py, THX!
def get_resolvers(file):
    resolvers = []
    try:
        with open(file, 'r') as resolvconf:
            for line in resolvconf.readlines():
                line = line.split('#',1)[0];
                line = line.rstrip();
                if 'nameserver' in line:
                    resolvers.append(line.split()[1])
        return resolvers
    except IOError as error:
        return error.strerror

def get_config():
    dns = {}
    conf = Config()
    conf.set_level('service dns forwarding')
    if not conf.exists(''):
        return dns

    if conf.exists('cache-size'):
        cache = conf.return_value('cache-size')
        dns.setdefault('cache-size', cache)

    if conf.exists('dhcp'):
        dns.setdefault('dhcp', [])
        interfaces = []
        interfaces = conf.return_values('dhcp')
        for interface in interfaces:
            resolvers = get_resolvers("/etc/resolv.conf.dhclient-new-{0}".format(interface))
            dhcp = {
                "interface": interface,
                "resolvers": resolvers
            }
            dns['dhcp'].append(dhcp)

    if conf.exists('domain'):
        dns.setdefault('domain', [])
        for node in conf.list_nodes('domain'):
            server = conf.return_value("domain {0} server".format(node))
            domain = {
                "name": node,
                "server": server
            }
            dns['domain'].append(domain)

    if conf.exists('ignore-hosts-file'):
        dns.setdefault('ignore-hosts-file', True)

    if conf.exists('listen-on'):
        interfaces = []
        interfaces = conf.return_values('listen-on')
        dns.setdefault('listen-on', interfaces)

    if conf.exists('name-server'):
        nameservers = []
        nameservers = conf.return_values('name-server')
        dns.setdefault('name-server', nameservers)

    if conf.exists('query-all-servers'):
        dns.setdefault('query-all-servers', True)

    if conf.exists('system'):
        conf.set_level('system')
        nameservers = []
        nameservers = conf.return_values('name-server')
        if len(nameservers) == 0:
            print("DNS forwarding warning: No name-servers set under 'system name-server'\n")
        dns.setdefault('system-name-server', True)

    return dns

def verify(dns):
    if len(dns) > 0:
        if 'listen-on' not in dns.keys():
            raise ConfigError("Error: DNS forwarding requires a configured listen interface!")

    return None

def generate(dns):
    config_header = '### Autogenerated by vyos-config-dns-forwarding.py on {tm} ###\n'.format(tm=time.strftime("%a, %d %b %Y %H:%M:%S", time.localtime()))

    # write new configuration file
    f = open(config_file, 'w')
    f.write(config_header)
    f.write("log-facility=/var/log/dnsmasq.log\n")
    f.write("no-poll\n")
    f.write("edns-packet-max=4096\n")
    f.write("bind-interfaces\n")

    if 'listen-on' in dns.keys():
        for interface in dns['listen-on']:
            f.write("interface={0}\n".format(interface))

    if 'dhcp' in dns.keys():
        for dhcp in dns['dhcp']:
            for resolver in dhcp['resolvers']:
                f.write("server={0}\t# dhcp {1}\n".format(resolver, dhcp['interface']))

    if 'domain' in dns.keys():
        for domain in dns['domain']:
            f.write("server=/{0}/{1}\t# domain-override\n".format(domain['name'], domain['server']))

    if 'cache-size' in dns.keys():
        f.write("cache-size={0}\n".format(dns['cache-size']))

    if 'ignore-hosts-file' in dns.keys():
        f.write("no-hosts\n")

    if 'name-server' in dns.keys():
        for nameserver in dns['name-server']:
            f.write("server={0}\t# statically configured\n".format(nameserver))

    if 'query-all-servers' in dns.keys():
        f.write("all-servers\n")

    if 'system-name-server' in dns.keys():
        # Read the IP addresses of the upstream nameservers from /etc/resolv.conf
        f.write("resolv-file=/etc/resolv.conf\n")

    f.close()
    return None

def apply(dns):
    if len(dns) == 0:
        cmd = "sudo systemctl stop dnsmasq"
    else:
        cmd = "sudo systemctl restart dnsmasq"

    os.system(cmd)
    return None

if __name__ == '__main__':
    try:
        c = get_config()
        verify(c)
        generate(c)
        apply(c)
    except ConfigError as e:
        print(e)
        sys.exit(1)