from __future__ import absolute_import

import argparse
import itertools
import logging
import os
import signal
import select
import socket
import sys
import threading
import time

import dns.message
import dns.rdatatype

from pydnstest import scenario


class TestServer:
    """ This simulates UDP DNS server returning scripted or mirror DNS responses. """

    def __init__(self, test_scenario, root_addr, addr_family):
        """ Initialize server instance. """
        self.thread = None
        self.srv_socks = []
        self.client_socks = []
        self.connections = []
        self.active = False
        self.active_lock = threading.Lock()
        self.condition = threading.Condition()
        self.scenario = test_scenario
        self.addr_map = []
        self.start_iface = 2
        self.cur_iface = self.start_iface
        self.kroot_local = root_addr
        self.addr_family = addr_family

    def __del__(self):
        """ Cleanup after deletion. """
        with self.active_lock:
            active = self.active
        if active:
            self.stop()

    def start(self, port=53):
        """ Synchronous start """
        with self.active_lock:
            if self.active:
                raise Exception('TestServer already started')
        with self.active_lock:
            self.active = True
        self.addr, _ = self.start_srv((self.kroot_local, port), self.addr_family)
        self.start_srv(self.addr, self.addr_family, socket.IPPROTO_TCP)
        self._bind_sockets()

    def stop(self):
        """ Stop socket server operation. """
        with self.active_lock:
            self.active = False
        if self.thread:
            self.thread.join()
        for conn in self.connections:
            conn.close()
        for srv_sock in self.srv_socks:
            srv_sock.close()
        for client_sock in self.client_socks:
            client_sock.close()
        self.client_socks = []
        self.srv_socks = []
        self.connections = []
        self.scenario = None

    def address(self):
        """ Returns opened sockets list """
        addrlist = []
        for s in self.srv_socks:
            addrlist.append(s.getsockname())
        return addrlist

    def handle_query(self, client):
        """
        Receive query from client socket and send an answer.

        Returns:
            True if client socket should be closed by caller
            False if client socket should be kept open
        """
        log = logging.getLogger('pydnstest.testserver.handle_query')
        server_addr = client.getsockname()[0]
        query, client_addr = scenario.recvfrom_msg(client)
        if query is None:
            return False
        log.debug('server %s received query from %s: %s', server_addr, client_addr, query)
        response, is_raw_data = self.scenario.reply(query, server_addr)
        if response:
            if not is_raw_data:
                data_to_wire = response.to_wire(max_size=65535)
                log.debug('response: %s', response)
            else:
                data_to_wire = response
                log.debug('raw response not printed')
        else:
            response = dns.message.make_response(query)
            response.set_rcode(dns.rcode.SERVFAIL)
            data_to_wire = response.to_wire()
            self.undefined_answers += 1
            self.scenario.current_step.log.error(
                'server %s has no response for question %s, answering with SERVFAIL',
                server_addr,
                '; '.join([str(rr) for rr in query.question]))

        scenario.sendto_msg(client, data_to_wire, client_addr)
        return True

    def query_io(self):
        """ Main server process """
        self.undefined_answers = 0
        with self.active_lock:
            if not self.active:
                raise Exception("[query_io] Test server not active")
        while True:
            with self.condition:
                self.condition.notify()
            with self.active_lock:
                if not self.active:
                    break
            objects = self.srv_socks + self.connections
            to_read, _, to_error = select.select(objects, [], objects, 0.1)
            for sock in to_read:
                if sock in self.srv_socks:
                    if sock.proto == socket.IPPROTO_TCP:
                        conn, _ = sock.accept()
                        self.connections.append(conn)
                    else:
                        self.handle_query(sock)
                elif sock in self.connections:
                    if not self.handle_query(sock):
                        sock.close()
                        self.connections.remove(sock)
                else:
                    raise Exception(
                        "[query_io] Socket IO internal error {}, exit".format(sock.getsockname()))
            for sock in to_error:
                raise Exception("[query_io] Socket IO error {}, exit".format(sock.getsockname()))

    def start_srv(self, address, family, proto=socket.IPPROTO_UDP):
        """ Starts listening thread if necessary """
        assert address
        assert address[0]  # host
        assert address[1]  # port
        assert family
        assert proto
        if family == socket.AF_INET6:
            if not socket.has_ipv6:
                raise NotImplementedError("[start_srv] IPv6 is not supported by socket {0}"
                                          .format(socket))
        elif family != socket.AF_INET:
            raise NotImplementedError("[start_srv] unsupported protocol family {0}".format(family))

        if proto == socket.IPPROTO_TCP:
            socktype = socket.SOCK_STREAM
        elif proto == socket.IPPROTO_UDP:
            socktype = socket.SOCK_DGRAM
        else:
            raise NotImplementedError("[start_srv] unsupported protocol {0}".format(proto))

        if self.thread is None:
            self.thread = threading.Thread(target=self.query_io)
            self.thread.start()
            with self.condition:
                self.condition.wait()

        for srv_sock in self.srv_socks:
            if (srv_sock.family == family
                    and srv_sock.getsockname() == address
                    and srv_sock.proto == proto):
                return srv_sock.getsockname()

        sock = socket.socket(family, socktype, proto)
        sock.bind(address)
        sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        if proto == socket.IPPROTO_TCP:
            sock.listen(5)
        self.srv_socks.append(sock)
        sockname = sock.getsockname()
        return sockname, proto

    def _bind_sockets(self):
        """
        Bind test server to port 53 on all addresses referenced by test scenario.
        """
        # Bind to test servers
        for r in self.scenario.ranges:
            for addr in r.addresses:
                family = socket.AF_INET6 if ':' in addr else socket.AF_INET
                self.start_srv((addr, 53), family)

        # Bind addresses in ad-hoc REPLYs
        for s in self.scenario.steps:
            if s.type == 'REPLY':
                reply = s.data[0].message
                for rr in itertools.chain(reply.answer,
                                          reply.additional,
                                          reply.question,
                                          reply.authority):
                    for rd in rr:
                        if rd.rdtype == dns.rdatatype.A:
                            self.start_srv((rd.address, 53), socket.AF_INET)
                        elif rd.rdtype == dns.rdatatype.AAAA:
                            self.start_srv((rd.address, 53), socket.AF_INET6)

    def play(self, subject_addr):
        self.scenario.play({'': (subject_addr, 53)})


def empty_test_case():
    """
    Return (scenario, config) pair which answers to any query on 127.0.0.10.
    """
    # Mirror server
    empty_test_path = os.path.dirname(os.path.realpath(__file__)) + "/empty.rpl"
    test_config = {'ROOT_ADDR': '127.0.0.10',
                   '_SOCKET_FAMILY': socket.AF_INET}
    return scenario.parse_file(empty_test_path)[0], test_config


def standalone_self_test():
    """
    Self-test code

    Usage:
    LD_PRELOAD=libsocket_wrapper.so SOCKET_WRAPPER_DIR=/tmp $PYTHON -m pydnstest.testserver --help
    """
    logging.basicConfig(level=logging.DEBUG)
    argparser = argparse.ArgumentParser()
    argparser.add_argument('--scenario', help='absolute path to test scenario',
                           required=False)
    argparser.add_argument('--step', help='step # in the scenario (default: first)',
                           required=False, type=int)
    args = argparser.parse_args()
    if args.scenario:
        test_scenario, test_config_text = scenario.parse_file(args.scenario)
        test_config, _ = scenario.parse_config(test_config_text, True, os.getcwd())
    else:
        test_scenario, test_config = empty_test_case()

    if args.step:
        for step in test_scenario.steps:
            if step.id == args.step:
                test_scenario.current_step = step
        if not test_scenario.current_step:
            raise ValueError('step ID %s not found in scenario' % args.step)
    else:
        test_scenario.current_step = test_scenario.steps[0]

    server = TestServer(test_scenario, test_config['ROOT_ADDR'], test_config['_SOCKET_FAMILY'])
    server.start()

    logging.info("[==========] Mirror server running at %s", server.address())

    def exit(signum, frame):
        logging.info("[==========] Shutdown.")
        server.stop()
        sys.exit(128 + signum)

    signal.signal(signal.SIGINT, exit)
    signal.signal(signal.SIGTERM, exit)

    while True:
        time.sleep(0.5)


if __name__ == '__main__':
    # this is done to avoid creating global variables
    standalone_self_test()
