##      common/resolver.py
##
## Copyright (C) 2006 Dimitur Kirov <dkirov@gmail.com>
##
## This file is part of Gajim.
##
## Gajim is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published
## by the Free Software Foundation; version 3 only.
##
## Gajim 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 Gajim.  If not, see <http://www.gnu.org/licenses/>.
##

import sys
import os
import re
import logging
import functools
log = logging.getLogger('gajim.c.resolver')

if __name__ == '__main__':
    sys.path.append('..')
    from common import i18n
    import common.configpaths
    common.configpaths.gajimpaths.init(None)

from common import helpers
from nbxmpp.idlequeue import IdleCommand
from gi.repository import Gio, GLib


# it is good to check validity of arguments, when calling system commands
ns_type_pattern = re.compile('^[a-z]+$')

# match srv host_name
host_pattern = re.compile('^[a-z0-9\-._]*[a-z0-9]\.[a-z]{2,}$')

try:
    #raise ImportError("Manually disabled libasync")
    import libasyncns
    USE_LIBASYNCNS = True
    log.info("libasyncns-python loaded")
except ImportError:
    USE_LIBASYNCNS = False
    log.debug("Import of libasyncns-python failed, getaddrinfo will block", exc_info=True)


def get_resolver(idlequeue):
    if USE_LIBASYNCNS:
        log.info('Using LibAsyncNSResolver')
        return LibAsyncNSResolver()
    else:
        if helpers.is_in_path('host'):
            log.info('Using HostResolver')
            return HostResolver(idlequeue)
        log.info('Using NSLookupResolver')
        return NSLookupResolver(idlequeue)

class CommonResolver():
    def __init__(self):
        # dict {"host+type" : list of records}
        self.resolved_hosts = {}
        # dict {"host+type" : list of callbacks}
        self.handlers = {}

    def resolve(self, host, on_ready, type_='srv'):
        host = host.lower()
        log.debug('resolve %s type=%s' % (host, type_))
        assert(type_ in ['srv', 'txt'])
        if not host:
            # empty host, return empty list of srv records
            on_ready([])
            return
        if host + type_ in self.resolved_hosts:
            # host is already resolved, return cached values
            log.debug('%s already resolved: %s' % (host,
                self.resolved_hosts[host + type_]))
            on_ready(host, self.resolved_hosts[host + type_])
            return
        if host + type_ in self.handlers:
            # host is about to be resolved by another connection,
            # attach our callback
            log.debug('already resolving %s' % host)
            self.handlers[host + type_].append(on_ready)
        else:
            # host has never been resolved, start now
            log.debug('Starting to resolve %s using %s' % (host, self))
            self.handlers[host + type_] = [on_ready]
            self.start_resolve(host, type_)

    def _on_ready(self, host, type_, result_list):
        # practically it is impossible to be the opposite, but who knows :)
        host = host.lower()
        log.debug('Resolving result for %s: %s' % (host, result_list))
        if host + type_ not in self.resolved_hosts:
            self.resolved_hosts[host + type_] = result_list
        if host + type_ in self.handlers:
            for callback in self.handlers[host + type_]:
                callback(host, result_list)
            del(self.handlers[host + type_])

    def start_resolve(self, host, type_):
        pass

# FIXME: API usage is not consistent! This one requires that process is called
class LibAsyncNSResolver(CommonResolver):
    """
    Asynchronous resolver using libasyncns-python. process() method has to be
    called in order to proceed the pending requests. Based on patch submitted by
    Damien Thebault.
    """

    def __init__(self):
        self.asyncns = libasyncns.Asyncns()
        CommonResolver.__init__(self)

    def start_resolve(self, host, type_):
        type_ = libasyncns.ns_t_srv
        if type_ == 'txt': type_ = libasyncns.ns_t_txt
        resq = self.asyncns.res_query(host, libasyncns.ns_c_in, type_)
        resq.userdata = {'host':host, 'type':type_}

    # getaddrinfo to be done
    #def resolve_name(self, dname, callback):
        #resq = self.asyncns.getaddrinfo(dname)
        #resq.userdata = {'callback':callback, 'dname':dname}

    def _on_ready(self, host, type_, result_list):
        if type_ == libasyncns.ns_t_srv: type_ = 'srv'
        elif type_ == libasyncns.ns_t_txt: type_ = 'txt'

        CommonResolver._on_ready(self, host, type_, result_list)

    def process(self):
        try:
            self.asyncns.wait(False)
            resq = self.asyncns.get_next()
        except:
            return True
        if type(resq) == libasyncns.ResQuery:
            # TXT or SRV result
            while resq is not None:
                try:
                    rl = resq.get_done()
                except Exception:
                    rl = []
                hosts = []
                requested_type = resq.userdata['type']
                requested_host = resq.userdata['host']
                if rl:
                    for r in rl:
                        if r['type'] != requested_type:
                            # Answer doesn't contain valid SRV data
                            continue
                        r['prio'] = r['pref']
                        hosts.append(r)
                self._on_ready(host=requested_host, type_=requested_type,
                        result_list=hosts)
                try:
                    resq = self.asyncns.get_next()
                except Exception:
                    resq = None
        elif type(resq) == libasyncns.AddrInfoQuery:
            # getaddrinfo result (A or AAAA)
            rl = resq.get_done()
            resq.userdata['callback'](resq.userdata['dname'], rl)
        return True


class NSLookupResolver(CommonResolver):
    """
    Asynchronous DNS resolver calling nslookup. Processing of pending requests
    is invoked from idlequeue which is watching file descriptor of pipe of
    stdout of nslookup process.
    """

    def __init__(self, idlequeue):
        self.idlequeue = idlequeue
        self.process = False
        CommonResolver.__init__(self)

    def parse_srv_result(self, fqdn, result):
        """
        Parse the output of nslookup command and return list of properties:
        'host', 'port','weight', 'priority'     corresponding to the found srv hosts
        """
        if os.name == 'nt':
            return self._parse_srv_result_nt(fqdn, result)
        elif os.name == 'posix':
            return self._parse_srv_result_posix(fqdn, result)

    def _parse_srv_result_nt(self, fqdn, result):
        # output from win32 nslookup command
        if not result:
            return []
        hosts = []
        result = result.decode(sys.stdout.encoding)
        lines = result.replace('\r', '').split('\n')
        current_host = None
        for line in lines:
            line = line.lstrip()
            if line == '':
                continue
            if line.startswith(fqdn):
                rest = line[len(fqdn):]
                if rest.find('service') > -1:
                    current_host = {}
            elif isinstance(current_host, dict):
                res = line.strip().split('=')
                if len(res) != 2:
                    if len(current_host) == 4:
                        hosts.append(current_host)
                    current_host = None
                    continue
                prop_type = res[0].strip()
                prop_value = res[1].strip()
                if prop_type.find('prio') > -1:
                    try:
                        current_host['prio'] = int(prop_value)
                    except ValueError:
                        continue
                elif prop_type.find('weight') > -1:
                    try:
                        current_host['weight'] = int(prop_value)
                    except ValueError:
                        continue
                elif prop_type.find('port') > -1:
                    try:
                        current_host['port'] = int(prop_value)
                    except ValueError:
                        continue
                elif prop_type.find('host') > -1:
                    # strip '.' at the end of hostname
                    if prop_value[-1] == '.':
                        prop_value = prop_value[:-1]
                    current_host['host'] = prop_value
                if len(current_host) == 4:
                    hosts.append(current_host)
                    current_host = None
        return hosts

    def _parse_srv_result_posix(self, fqdn, result):
        # typical output of bind-tools nslookup command:
        # _xmpp-client._tcp.jabber.org    service = 30 30 5222 jabber.org.
        if not result:
            return []
        ufqdn = helpers.ascii_to_idn(fqdn) # Unicode domain name
        hosts = []
        lines = result.split('\n')
        for line in lines:
            if line == '':
                continue
            domain = None
            if line.startswith(fqdn):
                domain = fqdn # For nslookup 9.5
            elif line.startswith(ufqdn):
                domain = ufqdn # For nslookup 9.6
            if domain:
                rest = line[len(domain):].split('=')
                if len(rest) != 2:
                    continue
                answer_type, props_str = rest
                if answer_type.strip() != 'service':
                    continue
                props = props_str.strip().split(' ')
                if len(props) < 4:
                    continue
                prio, weight, port, host  = props[-4:]
                if host[-1] == '.':
                    host = host[:-1]
                try:
                    prio = int(prio)
                    weight = int(weight)
                    port = int(port)
                except ValueError:
                    continue
                hosts.append({'host': host, 'port': port, 'weight': weight,
                        'prio': prio})
        return hosts

    def _on_ready(self, host, type_, result):
        # nslookup finished, parse the result and call the handlers
        result_list = self.parse_srv_result(host, result)
        CommonResolver._on_ready(self, host, type_, result_list)

    def start_resolve(self, host, type_):
        """
        Spawn new nslookup process and start waiting for results
        """
        ns = NsLookup(self._on_ready, host, type_)
        ns.set_idlequeue(self.idlequeue)
        ns.commandtimeout = 20
        ns.start()


class HostResolver(CommonResolver):
    """
    Asynchronous DNS resolver calling host. Processing of pending requests
    is invoked from idlequeue which is watching file descriptor of pipe of
    stdout of host process.
    """

    def __init__(self, idlequeue):
        self.idlequeue = idlequeue
        self.process = False
        CommonResolver.__init__(self)

    def parse_srv_result(self, fqdn, result):
        """
        Parse the output of host command and return list of properties:
        'host', 'port','weight', 'priority' corresponding to the found srv hosts
        """
        # typical output of host command:
        # _xmpp-client._tcp.jabber.org has SRV record 30 30 5222 jabber.org.
        if not result:
            return []
        ufqdn = helpers.ascii_to_idn(fqdn) # Unicode domain name
        hosts = []
        lines = result.split('\n')
        for line in lines:
            if line == '':
                continue
            domain = None
            if line.startswith(fqdn):
                domain = fqdn # For nslookup 9.5
            elif line.startswith(ufqdn):
                domain = ufqdn # For nslookup 9.6
            if domain:
                # add 4 for ' has' after domain name
                rest = line[len(domain)+4:].split('record')
                if len(rest) != 2:
                    continue
                answer_type, props_str = rest
                if answer_type.strip() != 'SRV':
                    continue
                props = props_str.strip().split(' ')
                if len(props) < 4:
                    continue
                prio, weight, port, host  = props[-4:]
                if host[-1] == '.':
                    host = host[:-1]
                try:
                    prio = int(prio)
                    weight = int(weight)
                    port = int(port)
                except ValueError:
                    continue
                hosts.append({'host': host, 'port': port, 'weight': weight,
                        'prio': prio})
        return hosts

    def _on_ready(self, host, type_, result):
        # host finished, parse the result and call the handlers
        result_list = self.parse_srv_result(host, result)
        CommonResolver._on_ready(self, host, type_, result_list)

    def start_resolve(self, host, type_):
        """
        Spawn new nslookup process and start waiting for results
        """
        ns = Host(self._on_ready, host, type_)
        ns.set_idlequeue(self.idlequeue)
        ns.commandtimeout = 20
        ns.start()

class NsLookup(IdleCommand):
    def __init__(self, on_result, host='_xmpp-client', type_='srv'):
        IdleCommand.__init__(self, on_result)
        self.commandtimeout = 10
        self.host = host.lower()
        self.type_ = type_.lower()
        if not host_pattern.match(self.host):
            # invalid host name
            log.error('Invalid host: %s' % self.host)
            self.canexecute = False
            return
        if not ns_type_pattern.match(self.type_):
            log.error('Invalid querytype: %s' % self.type_)
            self.canexecute = False
            return

    def _compose_command_args(self):
        return ['nslookup', '-type=' + self.type_, self.host]

    def _return_result(self):
        if self.result_handler:
            self.result_handler(self.host, self.type_, self.result)
        self.result_handler = None


class Host(NsLookup):
    def _compose_command_args(self):
        return ['host', '-t', self.type_, self.host]


class GioResolver(CommonResolver):
    """
    Asynchronous resolver using GIO. process() method has to be
    called in order to proceed the pending requests.
    """

    def __init__(self):
        super().__init__()
        self.gio_resolver = Gio.Resolver.get_default()

    def start_resolve(self, host, type_):
        if type_ == 'txt':
            # TXT record resolution isn't used anywhere at the moment so
            # implementing it here isn't urgent
            raise NotImplemented("Gio resolver does not currently implement TXT records")
        else:
            callback = functools.partial(self._on_ready_srv, host)
            type_ = Gio.ResolverRecordType.SRV

        resq = self.gio_resolver.lookup_records_async(host, type_, None, callback)

    def _on_ready_srv(self, host, source_object, result):
        try:
            variant_results = source_object.lookup_records_finish(result)
        except GLib.Error as e:
            if e.domain == 'g-resolver-error-quark':
                result_list = []
                log.warning("Could not resolve host: %s", e.message)
            else:
                raise
        else:
            result_list = [
                {
                    'weight': weight,
                    'prio': prio,
                    'port': port,
                    'host': host,
                }
                for prio, weight, port, host
                in variant_results
            ]
        super()._on_ready(host, 'srv', result_list)


# below lines is on how to use API and assist in testing
if __name__ == '__main__':
    from gi.repository import GLib
    from gi.repository import Gtk
    from nbxmpp import idlequeue

    idlequeue = idlequeue.get_idlequeue()
    resolver = get_resolver(idlequeue)

    def clicked(widget):
        global resolver
        host = text_view.get_text()
        def on_result(host, result_array):
            print('Result:\n' + repr(result_array))
        resolver.resolve(host, on_result)
    win = Gtk.Window()
    win.set_border_width(6)
    text_view = Gtk.Entry()
    text_view.set_text('_xmpp-client._tcp.jabber.org')
    hbox = Gtk.HBox()
    hbox.set_spacing(3)
    but = Gtk.Button(' Lookup SRV ')
    hbox.pack_start(text_view, 5, True, 0)
    hbox.pack_start(but, 0, True, 0)
    but.connect('clicked', clicked)
    win.add(hbox)
    win.show_all()
    GLib.timeout_add(200, idlequeue.process)
    if USE_LIBASYNCNS:
        GLib.timeout_add(200, resolver.process)
    Gtk.main()