# Source code for pathspider.base

    """
    Basic framework for Pathspider: coordinate active measurements on large target
    lists with both system-level network stack state (sysctls, iptables rules, etc)
    as well as information derived from flow-level passive observation of traffic at
    the sender.

    .. moduleauthor:: Brian Trammell <brian@trammell.ch>

    """

    import sys
    import time
    import logging
    import collections
    import threading
    import multiprocessing as mp
    import queue
    from datetime import datetime

    from pathspider.network import ipv4_address
    from pathspider.network import ipv6_address
    from pathspider.network import ipv4_address_public
    from pathspider.network import ipv6_address_public
    from pathspider.network import ipv4_asn
    from pathspider.network import ipv6_asn

    __version__ = "2.0.1"

    CONN_OK = 0
    CONN_FAILED = 1
    CONN_TIMEOUT = 2
    CONN_SKIPPED = 3
    CONN_DISCARD = 4

    PORT_FAILED = 0
    PORT_FAILED_AGAIN = -1

    QUEUE_SIZE = 1000
    QUEUE_SLEEP = 0.5

    SHUTDOWN_SENTINEL = "SHUTDOWN_SENTINEL"
    NO_FLOW = None

    class Spider:
        """
        A spider consists of a configurator (which alternates between two system
        configurations), a large number of workers (for performing some network
        action for each configuration), an Observer which derives information from
        passively observed traffic, and a thread that merges results from the
        workers with flow records from the collector.

        """

        chains = [] # Disable the observer by default

        def __init__(self, worker_count, libtrace_uri, args, server_mode):
            """
            The initialisation of a pathspider plugin.

            :param worker_count: The number of workers to use.
            :type worker_count: int
            :param libtrace_uri: The URI to pass to the Observer to describe the
                                 interface on which packets should be captured.
            :type libtrace_uri: str

            It is expected that this function will be overloaded by plugins, though
            the plugin should always make a call to the __init__() function of the
            abstract Spider class as this initialises all of the base functionality:

            .. code-block:: python

             super().__init__(worker_count=worker_count,
                              libtrace_uri=libtrace_uri,
                              args=args)

            This can be used to initialise any variables which may be required in
            the object.
            """

            self.args = args

            self.running = False
            self.stopping = False
            self.terminating = False
            self.server_mode = server_mode

            self.worker_count = worker_count
            self.active_worker_count = 0
            self.active_worker_lock = threading.Lock()

            self.libtrace_uri = libtrace_uri

            self.jobqueue = queue.Queue(QUEUE_SIZE)
            self.resqueue = queue.Queue(QUEUE_SIZE)
            self.flowqueue = mp.Queue(QUEUE_SIZE)
            self.observer_shutdown_queue = mp.Queue(QUEUE_SIZE)

            self.jobtab = {}
            self.comparetab = {}
            self.restab = {}
            self.flowtab = {}
            self.flowreap = collections.deque()
            self.flowreap_size = min(self.worker_count * 100, 10000)

            self.outqueue = queue.Queue(QUEUE_SIZE)

            self.observer = None

            self.worker_threads = []
            self.configurator_thread = None
            self.merger_thread = None

            self.observer_process = None

            self.lock = threading.Lock()
            self.exception = None

            if libtrace_uri.startswith('int'):
                # TODO: Refactor this
                self.source = (ipv4_address(self.libtrace_uri[4:]),
                               ipv6_address(self.libtrace_uri[4:]))
                self.source_public = (ipv4_address_public(self.libtrace_uri[4:]),
                                      ipv6_address_public(self.libtrace_uri[4:]))
                self.source_asn = (ipv4_asn(self.libtrace_uri[4:]),
                                   ipv6_asn(self.libtrace_uri[4:]))
            else:
                self.source = ("127.0.0.1", "::1")

            self.__logger = logging.getLogger('pathspider')

            if hasattr(self.args, 'connect') and self.args.connect.startswith('tor'):
                logging.getLogger("stem").setLevel(logging.ERROR)
                import stem.control
                self.controller = stem.control.Controller.from_port()
                self.controller.authenticate()

        def configurator(self):
            raise NotImplementedError("Cannot instantiate an abstract Spider")

        def worker(self, worker_number):
            raise NotImplementedError("Cannot instantiate an abstract Spider")

        def pre_connect(self, job):
            """
            Performs pre-connection operations.

            :param job: The job record
            :type job: dict

            The pre_connect function can be used to perform any operations that
            must be performed before each connection. It will be run only once
            per job, with the same result passed to both the A and B connect
            calls. This function is not synchronized with the configurator.

            Plugins to PATHspider can optionally implement this function. If this
            function is not overloaded, it will be a noop.
            """

            pass

        def _connect_wrapper(self, job, config, connect=None):
            start = str(datetime.utcnow())
            if connect is None:
                conn = self.connect(job, config) # pylint: disable=no-member
            else:
                if not hasattr(connect, '__self__'):
                    connect = connect.__get__(self)
                conn = connect(job, config)
            conn['spdr_start'] = start
            return conn

        def post_connect(self, job, rec, config):
            """
            Performs post-connection operations.

            :param job: The job record.
            :type job: dict
            :param rec: The result of the connection operation(s).
            :type rec: dict
            :param config: The state of the configurator during
                           :func:`pathspider.base.Spider.connect`.
            :type config: int

            The post_connect function can be used to perform any operations that
            must be performed after each connection. It will be run for both the
            A and the B configuration, and is not synchronized with the
            configurator.

            Plugins to PATHspider can optionally implement this function. If this
            function is not overloaded, it will be a noop.

            Any sockets or other file handles that were opened during
            :func:`pathspider.base.Spider.connect` should be closed in this
            function if they have not been already.
            """

            pass

        def create_observer(self):
            """
            Create a flow observer.

            This function is called by the base Spider logic to get an instance
            of :class:`pathspider.observer.Observer` configured with the function
            chains that are requried by the plugin.
            """

            self.__logger.info("Creating observer")
            if len(self.chains) > 0:
                from pathspider.observer import Observer
                return Observer(self.libtrace_uri,
                                chains=self.chains) # pylint: disable=no-member
            else:
                from pathspider.observer import DummyObserver
                return DummyObserver()

        def _key(self, obj):
            if self.server_mode:
                objkey = (obj['sip'], obj['sp'])
            else:
                objkey = (obj['dip'], obj['sp'])
            return objkey

        def _merge_flows(self):
            try:
                flow = self.flowqueue.get_nowait()
            except queue.Empty:
                time.sleep(QUEUE_SLEEP)
                return True
            else:
                if flow == SHUTDOWN_SENTINEL:
                    self.__logger.debug("stopping flow merging on sentinel")
                    return False

                flowkey = self._key(flow)
                self.__logger.debug("got a flow (" + repr(flowkey) + ")")

                if flowkey in self.restab:
                    self.__logger.debug("merging flow")
                    self.merge(flow, self.restab[flowkey])
                    del self.restab[flowkey]
                elif flowkey in self.flowtab:
                    self.__logger.debug("won't merge duplicate flow")
                else:
                    # Create a new flow
                    self.flowtab[flowkey] = flow

                    # And reap the oldest, if the reap queue is full
                    self.flowreap.append(flowkey)
                    if len(self.flowreap) > self.flowreap_size:
                        try:
                            del self.flowtab[self.flowreap.popleft()]
                        except KeyError:
                            pass
                return True

        def _merge_results(self):
            try:
                res = self.resqueue.get_nowait()
            except queue.Empty:
                time.sleep(QUEUE_SLEEP)
                self.__logger.debug("result queue is empty")
                return True
            else:
                if res == SHUTDOWN_SENTINEL:
                    self.__logger.debug("stopping result merging on sentinel")
                    return False
                if 'spdr_state' in res.keys() and res['spdr_state'] == CONN_SKIPPED:
                    # handle skipped results
                    return True

                reskey = self._key(res)
                self.__logger.debug("got a result (" + repr(reskey) + ")")

                if reskey in self.restab and res['sp'] == PORT_FAILED:
                    # both connections failed, but need to be distinguished
                    reskey = (reskey[0], PORT_FAILED_AGAIN)

                if reskey in self.flowtab:
                    self.__logger.debug("merging result")
                    self.merge(self.flowtab[reskey], res)
                    del self.flowtab[reskey]
                elif reskey in self.restab:
                    self.__logger.debug("won't merge duplicate result")
                else:
                    self.restab[reskey] = res

                self.resqueue.task_done()
                return True

        def merger(self):
            """
            Thread to merge results from the workers and the observer.
            """

            if len(self.chains) == 0:
                self.__logger.warning("Merger is not expecting flows from the Observer")
                # Immediately merge with NO_FLOW when there's no chains, as there's
                # going to also be no observer and so no flows.
                while self.running:
                    try:
                        res = self.resqueue.get_nowait()
                    except queue.Empty:
                        time.sleep(QUEUE_SLEEP)
                        self.__logger.debug("result queue is empty")
                        continue
                    if res == SHUTDOWN_SENTINEL:
                        break
                    if not ('spdr_state' in res.keys() and res['spdr_state'] == CONN_SKIPPED):
                        self.merge(NO_FLOW, res)
            else:
                merging_flows = True
                merging_results = True

                while self.running and (merging_results or merging_flows):

                    if merging_flows and self.flowqueue.qsize() >= self.resqueue.qsize():
                        merging_flows = self._merge_flows()

                    elif merging_results:
                        merging_results = self._merge_results()

            # One more pass during shutdown, to clean up any leftovers
            for res_item in self.restab.items():
                res = res_item[1]
                self.merge(NO_FLOW, res)

        def merge(self, flow, res):
            """
            Merge a job record with a flow record.

            :param flow: The flow record.
            :type flow: dict
            :param res: The job record.
            :type res: dict
            :return: tuple -- Final record for job.

            In order to create a final record for reporting on a job, the final job
            record must be merged with the flow record. This function should
            be implemented by any plugin to provide the logic for this merge as
            the keys used in these records cannot be known by PATHspider in advance.

            This method is not implemented in the abstract
            :class:`pathspider.base.Spider` class and must be implemented by any
            plugin.
            """

            if flow == NO_FLOW:
                flow = {'observed': False}
            else:
                flow['observed'] = True

            for key in res.keys():
                if key in flow.keys():
                    if res[key] == flow[key]:
                        continue
                    else:
                        self.__logger.warning("Dropping flow due to mismatch with "
                                              "observations on key %s", key)
                        return
                flow[key] = res[key]

            # Remove private keys - we need to make a copy of the keys as we
            #                       modify the dict during the iteration
            for key in [x for x in flow.keys()]:
                if key.startswith("_"):
                    flow.pop(key)

            self.__logger.debug("Result: " + str(flow))

            if flow['jobId'] not in self.comparetab:
                self.comparetab[flow['jobId']] = []
            self.comparetab[flow['jobId']].append(flow)

            if len(self.comparetab[flow['jobId']]) == self._config_count: # pylint: disable=no-member
                flows = self.comparetab.pop(flow['jobId'])
                flows.sort(key=lambda x: x['config'])
                start = min([flow['spdr_start'] for flow in flows])
                stop = max([flow['spdr_stop'] for flow in flows])
                job = self.jobtab.pop(flow['jobId'])
                job['flow_results'] = flows
                job['time'] = {'from': start, 'to': stop}
                job['missed_flows'] = 0
                for flow in flows:
                    if not flow['observed']:
                        job['missed_flows'] = job['missed_flows'] + 1
                job['conditions'] = self.combine_flows(flows)
                if job['conditions'] is not None:
                    if "pathspider.not_observed" in job['conditions']:
                        self.__logger.debug("At least one flow was not observed and so conditions could not be fully generated (if at all)")
                    if job['missed_flows'] > 0:
                        job['conditions'].append("pathspider.missed_flows:" + str(job['missed_flows']))
                else:
                    job.pop('conditions')
                self.outqueue.put(job)

        def combine_flows(self, flows):
            pass

        def exception_wrapper(self, target, *args, **kwargs):
            try:
                target(*args, **kwargs)
            except: # pylint: disable=W0702
                self.__logger = logging.getLogger('pathspider')
                self.__logger.exception("exception occurred. terminating.")
                if self.exception is None:
                    self.exception = sys.exc_info()[1]

                self.terminate()

        def _finalise_conns(self, job, jobId, conns):
            # Pass results on for merge
            config = 0
            for conn in conns:
                self.post_connect(job, conn, config)
                conn['spdr_stop'] = str(datetime.utcnow())
                conn['config'] = config
                if self.server_mode:
                    conn['sip'] = job['sip']
                else:
                    conn['dip'] = job['dip']
                conn['jobId'] = jobId
                self.resqueue.put(conn)
                config += 1

        def start(self):
            """
            This function starts a PATHspider plugin by:

             * Setting the running flag
             * Create and start an observer
             * Start the merger thread
             * Start the configurator thread
             * Start the worker threads

            The number of worker threads to start was given when activating the
            plugin.
            """

            self.__logger.info("starting pathspider")

            with self.lock:
                # set the running flag
                self.running = True

                # create an observer and start its process
                self.observer = self.create_observer()
                self.observer_process = mp.Process(
                    args=(self.observer.run_flow_enqueuer,
                          self.flowqueue,
                          self.observer_shutdown_queue),
                    target=self.exception_wrapper,
                    name='observer',
                    daemon=True)
                self.observer_process.start()
                self.__logger.debug("observer forked")

                # now start up ecnspider, backwards
                self.merger_thread = threading.Thread(
                    args=(self.merger,),
                    target=self.exception_wrapper,
                    name="merger",
                    daemon=True)
                self.merger_thread.start()
                self.__logger.debug("merger up")

                self.configurator_thread = threading.Thread(
                    args=(self.configurator,),
                    target=self.exception_wrapper,
                    name="configurator",
                    daemon=True)
                self.configurator_thread.start()
                self.__logger.debug("configurator up")
                self.worker_threads = []
                with self.active_worker_lock:
                    self.active_worker_count = self.worker_count
                for i in range(self.worker_count):
                    worker_thread = threading.Thread(
                        args=(self.worker, i),
                        target=self.exception_wrapper,
                        name='worker_{}'.format(i),
                        daemon=True)
                    self.worker_threads.append(worker_thread)
                    worker_thread.start()
                self.__logger.debug("workers up")

        def shutdown(self):
            """
            Shut down PathSpider in an orderly fashion,
            ensuring that all queued jobs complete,
            and all available results are merged.

            """

            self.__logger.info("beginning shutdown")

            with self.lock:
                # Set stopping flag
                self.stopping = True

                # Put a bunch of shutdown signals in the job queue
                for _ in range(self.worker_count * 2):
                    self.jobqueue.put(SHUTDOWN_SENTINEL)

                # Wait for worker threads to shut down
                for worker in self.worker_threads:
                    if threading.current_thread() != worker:
                        self.__logger.debug("joining worker: " + repr(worker))
                        worker.join()
                self.__logger.debug("all workers joined")

                # Tell observer to shut down
                self.observer_shutdown_queue.put(True)
                self.observer_process.join()
                self.__logger.debug("observer shutdown")

                # Tell merger to shut down
                self.resqueue.put(SHUTDOWN_SENTINEL)
                self.merger_thread.join()
                self.__logger.debug("merger shutdown")

                # Wait for merged results to be written
                self.outqueue.join()
                self.__logger.debug("all results retrieved")

                # Propagate shutdown sentinel and tell threads to stop
                self.outqueue.put(SHUTDOWN_SENTINEL)

                # Tell threads we've stopped
                self.running = False

                # Join configurator
                # if threading.current_thread() != self.configurator_thread:
                #     self.configurator_thread.join()

                self.stopping = False

            self.__logger.info("shutdown complete")

        def terminate(self):
            """
            Shut down PathSpider as quickly as possible,
            without any regard to completeness of results.

            """
            self.__logger.info("terminating pathspider")

            # tell threads to stop
            self.stopping = True
            self.running = False

            # terminate observer
            self.observer_shutdown_queue.put(True)

            # drain queues
            try:
                while True:
                    self.jobqueue.task_done()
            except ValueError:
                pass

            try:
                while True:
                    self.resqueue.task_done()
            except ValueError:
                pass

            try:
                while True:
                    self.flowqueue.get_nowait()
            except queue.Empty:
                pass

            # Join remaining threads
            for worker in self.worker_threads:
                if threading.current_thread() != worker:
                    self.__logger.debug("joining worker: " + repr(worker))
                    worker.join()
            self.__logger.debug("all workers joined")

            if self.configurator_thread and \
                    (threading.current_thread() != self.configurator_thread):
                self.configurator_thread.join()
                self.__logger.debug("configurator joined")

            if threading.current_thread() != self.merger_thread:
                self.merger_thread.join()
                self.__logger.debug("merger joined")

            self.observer_process.join()
            self.__logger.debug("observer joined")

            self.outqueue.put(SHUTDOWN_SENTINEL)
            self.__logger.info("termination complete")

        def add_job(self, job):
            """
            Adds a job to the job queue.

            If PATHspider is currently stopping, the job will not be added to the
            queue.
            """

            if self.stopping:
                return

            if not self.server_mode:
                if 'dip' in job.keys():
                    sourceindex = 1 if ':' in job['dip'] else 0
                    job['sip'] = self.source[sourceindex]
                    job['path'] = [job['sip']]
                    job['sip_public'] = self.source_public[sourceindex]
                    if not ( job['sip'] == job['sip_public'] ):
                        job['path'].append(job['sip_public'])
                    if self.source_asn[sourceindex] is not None:
                        job['sip_asn'] = self.source_asn[sourceindex]
                        job['path'].append("AS" + str(job['sip_asn']))
                    if 'dip_asn' in job.keys(): # This may be generated by other tools
                        job['path'].append("AS" + job['dip_asn'])
                    elif 'info' in job.keys(): # Hellfire does it this way
                        if 'ASN' in job['info'].keys():
                            job['path'].append("AS" + str(job['info']['ASN']))
                    job['path'].append(job['dip'])

            self.jobqueue.put(job)

    class PluggableSpider:
        @staticmethod
        def register_args(subparsers):
            raise NotImplementedError("Cannot register an abstract plugin")