# Source code for pathspider.observer

    import collections
    import logging
    import base64
    import queue
    import math

    from pathspider.base import SHUTDOWN_SENTINEL

    def _flow4_ids(ip):
        # FIXME keep map of fragment IDs to keys (#144)

        icmp_with_payload = {3, 4, 5, 11, 12}
        quotation_fid = False
        if ip.proto == 1 and ip.icmp.type in icmp_with_payload:
            if ip.icmp.payload is not None and len(ip.icmp.payload.data) >= 20:
                ip = ip.icmp.payload
                quotation_fid = True

        protos_with_ports = {6, 17, 132, 136}
        if ip.proto in protos_with_ports:
            # key includes ports
            fid = ip.src_prefix.addr + ip.dst_prefix.addr + ip.data[
                9:10] + ip.payload[0:4]
            rid = ip.dst_prefix.addr + ip.src_prefix.addr + ip.data[
                9:10] + ip.payload[2:4] + ip.payload[0:2]
        else:
            # no ports, just 3-tuple
            fid = ip.src_prefix.addr + ip.dst_prefix.addr + ip.data[9:10]
            rid = ip.dst_prefix.addr + ip.src_prefix.addr + ip.data[9:10]

        if quotation_fid:
            # If the fid is based on an ICMP quotation, need to be reversed
            return (base64.b64encode(rid), base64.b64encode(fid))
        else:
            return (base64.b64encode(fid), base64.b64encode(rid))

    def _flow6_ids(ip6):
        icmp_with_payload = {1, 2, 3, 4}
        quotation_fid = False

        if ip6.proto == 58 and ip6.icmp6.type in icmp_with_payload:
            if ip6.icmp6.payload is not None and len(ip6.icmp6.payload.data) >= 40:
                ip6 = ip6.icmp6.payload
                quotation_fid = True

        protos_with_ports = {6, 17, 132, 136}
        if ip6.proto in protos_with_ports:
            # key includes ports
            fid = ip6.src_prefix.addr + ip6.dst_prefix.addr + ip6.data[
                6:7] + ip6.payload[0:4]
            rid = ip6.dst_prefix.addr + ip6.src_prefix.addr + ip6.data[
                6:7] + ip6.payload[2:4] + ip6.payload[0:2]
        else:
            # no ports, just 3-tuple
            fid = ip6.src_prefix.addr + ip6.dst_prefix.addr + ip6.data[6:7]
            rid = ip6.dst_prefix.addr + ip6.src_prefix.addr + ip6.data[6:7]

        if quotation_fid:
            # If the fid is based on an ICMP quotation, need to be reversed
            return (base64.b64encode(rid), base64.b64encode(fid))
        else:
            return (base64.b64encode(fid), base64.b64encode(rid))

    PacketClockTimer = collections.namedtuple("PacketClockTimer", ("time", "fn"))

    class Observer:
        """
        Wraps a packet source identified by a libtrace URI,
        parses packets to divide them into flows, passing these
        packets and flows onto a function chain to allow
        data to be associated with each flow.
        """

        def __init__(self, lturi, chains=None, idle_timeout=30, expiry_timeout=5):
            """
            Create an Observer.

            :param chains: Array of Observer chain classes
            :see also: :ref:`Observer Documentation <observer>`
            """

            # Only import this when needed
            import plt as libtrace

            # Control
            self._irq = None
            self._irq_fired = False

            # Libtrace initialization
            self._trace = libtrace.trace(lturi)  # pylint: disable=no-member
            self._trace.start()
            self._pkt = libtrace.packet()  # pylint: disable=no-member

            # Chains of functions to evaluate
            chains = chains if chains is not None else []
            self._chains = [chain() for chain in chains]

            # Packet timer and bintables
            self._ptq = 0  # current packet timer, quantized
            self._idle_bins = {}  # map bin number to set of fids
            self._expiry_bins = {}  # map bin number to set of fids
            self._idle_timeout = idle_timeout
            self._expiry_timeout = expiry_timeout
            self._bin_quantum = 1

            #self._tq = []                  # packet timer queue (heap)

            # Flow tables
            self._active = {}
            self._expiring = {}
            self._ignored = set()

            # Emitter queue
            self._emitted = collections.deque()

            # Statistics and logging
            self._logger = logging.getLogger("observer")
            self._ct_pkt = 0
            self._ct_nonip = 0
            self._ct_shortkey = 0
            self._ct_ignored = 0
            self._ct_flow = 0

        def _interrupted(self):
            if not self._irq_fired and self._irq is not None:
                try:
                    self._irq.get_nowait()
                    self._irq_fired = True
                except queue.Empty:
                    pass

            return self._irq_fired

        def _get_chains(self, name):
            return [
                c.__getattribute__(name) for c in self._chains if hasattr(c, name)
            ]

        def _next_packet(self):
            # Import only when needed
            import plt as libtrace

            # see if someone told us to stop
            if self._interrupted():
                return False

            # see if we're done iterating
            if not self._trace.read_packet(self._pkt):
                return False

            # count the packet
            self._ct_pkt += 1

            # advance the packet clock
            self._tick(self._pkt.seconds)

            # get a flow ID and associated flow record for the packet
            (fid, rec, rev) = self._get_flow()

            # don't dispatch if we don't have a record
            # (this happens for non-IP packets and flows
            #  we know we want to ignore)
            if not rec:
                return True

            keep_flow = True

            # run IP header chains
            if self._pkt.ip:
                for fn in self._get_chains("ip4"):
                    keep_flow = keep_flow and fn(rec, self._pkt.ip, rev=rev)
                if self._pkt.icmp:
                    for fn in self._get_chains("icmp4"):
                        q = self._pkt.icmp.payload # pylint: disable=no-member
                        keep_flow = keep_flow and fn(rec, self._pkt.ip, q, rev=rev)

            elif self._pkt.ip6:
                for fn in self._get_chains("ip6"):
                    keep_flow = keep_flow and fn(rec, self._pkt.ip6, rev=rev)
                if self._pkt.icmp6:
                    for fn in self._get_chains("icmp6"):
                        q = self._pkt.icmp6.payload # pylint: disable=no-member
                        keep_flow = keep_flow and fn(
                            rec, self._pkt.ip6, q, rev=rev)

            # run transport header chains
            if self._pkt.tcp:
                for fn in self._get_chains("tcp"):
                    keep_flow = keep_flow and fn(rec, self._pkt.tcp, rev=rev)
            elif self._pkt.udp:
                for fn in self._get_chains("udp"):
                    keep_flow = keep_flow and fn(rec, self._pkt.udp, rev=rev)

            # complete the flow if any chain function asked us to
            if not keep_flow:
                self._flow_complete(fid)

            # we processed a packet, keep going
            return True

        # def _set_timer(self, delay, fid):
        #     # add to queue
        #     heapq.heappush(self._tq, PacketClockTimer(self._pt + delay,
        #                    self._finish_expiry_tfn(fid)))

        def _get_flow(self):
            """
            Get a flow record for the given packet.
            Create a new basic flow record
            """
            # get possible a flow IDs for the packet
            try:
                if self._pkt.ip:
                    (ffid, rfid) = _flow4_ids(self._pkt.ip)
                    ip = self._pkt.ip
                elif self._pkt.ip6:
                    (ffid, rfid) = _flow6_ids(self._pkt.ip6)
                    ip = self._pkt.ip6
                else:
                    # we don't care about non-IP packets
                    self._ct_nonip += 1
                    return (None, None, False)
            except ValueError:
                self._ct_shortkey += 1
                return (None, None, False)

            # now look for forward and reverse in ignored, active,
            # and expiring tables.
            if ffid in self._ignored:
                return (None, None, False)
            elif rfid in self._ignored:
                return (None, None, False)
            elif ffid in self._active:
                (fid, rec, active) = (ffid, self._active[ffid], True)
                #self._logger.debug("found forward flow for "+str(ffid))
            elif ffid in self._expiring:
                (fid, rec, active) = (ffid, self._expiring[ffid], False)
                #self._logger.debug("found expiring forward flow for "+str(ffid))
            elif rfid in self._active:
                (fid, rec, active) = (rfid, self._active[rfid], True)
                #self._logger.debug("found reverse flow for "+str(rfid))
            elif rfid in self._expiring:
                (fid, rec, active) = (rfid, self._expiring[rfid], False)
                #self._logger.debug("found expiring reverse flow for "+str(rfid))
            else:
                # nowhere to be found. new flow.
                rec = {'pkt_first': ip.seconds, '_idle_bin': 0}
                for fn in self._get_chains("new_flow"):
                    if not fn(rec, ip):
                        # self._logger.debug("ignoring "+str(ffid))
                        self._ignored.add(ffid)
                        self._ct_ignored += 1
                        return (None, None, False)

                # wasn't vetoed. add to active table.
                fid = ffid
                self._active[ffid] = rec
                active = True
                # self._logger.debug("new flow for "+str(ffid))
                self._ct_flow += 1

            # update time and idle bin and return record
            rec['pkt_last'] = ip.seconds

            # update idle bin if we're not expiring
            if active:
                new_idle_bin = math.ceil((rec['pkt_last'] + self._idle_timeout) /
                                         self._bin_quantum) * self._bin_quantum

                if new_idle_bin > rec["_idle_bin"]:

                    if rec['_idle_bin'] in self._idle_bins:
                        self._idle_bins[rec['_idle_bin']] -= set((fid, ))
                    if new_idle_bin in self._idle_bins:
                        self._idle_bins[new_idle_bin] |= set((fid, ))
                    else:
                        self._idle_bins[new_idle_bin] = set((fid, ))

                    rec['_idle_bin'] = new_idle_bin

            return (fid, rec, bool(fid == rfid))

        def _flow_complete(self, fid):
            """
            Mark a given flow ID as complete
            """
            # skip all of this unless the flow is still in the active table
            if fid not in self._active:
                return

            # remove flow ID from idle bin
            rec = self._active[fid]
            self._idle_bins[rec['_idle_bin']] -= set((fid, ))

            del rec['_idle_bin']

            # move record to expiring table
            self._expiring[fid] = rec
            del self._active[fid]

            # assign expiry bin
            expiry_bin = math.ceil((self._ptq + self._expiry_timeout) /
                                   self._bin_quantum) * self._bin_quantum
            #self._logger.debug("Completing flow "+str(fid)+" at "+str(self._ptq)+" to expire "+str(expiry_bin)+" (in "+str(expiry_bin-self._ptq)+"s)")

            if expiry_bin in self._expiry_bins:
                self._expiry_bins[expiry_bin] |= set((fid, ))
            else:
                self._expiry_bins[expiry_bin] = set((fid, ))

        def _emit_flow(self, rec):
            self._emitted.append(rec)

        def _next_flow(self):
            while len(self._emitted) == 0:
                if not self._next_packet():
                    return None

            return self._emitted.popleft()

        def _tick(self, pt):
            # quantize and skip if we're not advancing
            next_ptq = math.ceil(pt / self._bin_quantum) * self._bin_quantum
            if next_ptq <= self._ptq:
                return
            elif self._ptq == 0:
                # handle zero case
                self._ptq = next_ptq
                return

            # advance quantum
            for bint in range(self._ptq + self._bin_quantum,
                              next_ptq + self._bin_quantum, self._bin_quantum):
                self._logger.debug("tick: " + str(bint))

                # process idle
                if bint in self._idle_bins:
                    if len(self._idle_bins[bint]) > 0:
                        for fid in self._idle_bins[bint].copy():
                            self._flow_complete(fid)
                    del self._idle_bins[bint]

                # process expiry
                if bint in self._expiry_bins:
                    if len(self._expiry_bins[bint]) > 0:
                        for fid in self._expiry_bins[bint].copy():
                            self._emit_flow(self._expiring[fid])
                            del self._expiring[fid]
                    del self._expiry_bins[bint]

            self._ptq = next_ptq

        # def _tick(self, pt):
        #     # Advance packet clock
        #     self._pt = pt

        #     # fire all timers whose time has come
        #     while len(self._tq) > 0 and pt > min(self._tq, key=lambda x: x.time).time:
        #         try:
        #             heapq.heappop(self._tq).fn()
        #         except:
        #             type, value, tb = sys.exc_info()
        #             traceback.print_exc()
        #             pdb.post_mortem(tb)

        # def _finish_expiry_tfn(self, fid):
        #     """
        #     On expiry timer, emit the flow
        #     and delete it from the expiring queue
        #     """
        #     def tfn():
        #         if fid in self._expiring:
        #             self._emit_flow(self._expiring[fid])
        #             del self._expiring[fid]
        #             # self._logger.debug("emitted "+str(fid)+" on expiry")
        #     return tfn

        # def purge_idle(self, timeout=30):
        #     # TODO test this, it's probably pretty slow.
        #     for fid in self._active:
        #         if self._pt - self._active['fid']['last'] > timeout:
        #             self._flow_complete(fid)

        def flush(self):
            for fid in self._expiring:
                self._emit_flow(self._expiring[fid])
                # self._logger.debug("emitted "+str(fid)+" expiring during flush")
            self._expiring.clear()

            for fid in self._active:
                self._emit_flow(self._active[fid])
                # self._logger.debug("emitted "+str(fid)+" active during flush")
            self._active.clear()

            self._ignored.clear()

        def run_flow_enqueuer(self, flowqueue, irqueue=None):
            if irqueue:
                self._irq = irqueue
                self._irq_fired = None

            # Run main loop until last packet seen
            # then flush active flows and run again
            for _ in range(2):
                while True:
                    f = self._next_flow()
                    if f:
                        flowqueue.put(f)
                    else:
                        self.flush()
                        break

            # log observer info on shutdown
            self._logger.info(("processed %u packets "
                               "(%u dropped, %u short, %u non-ip) "
                               "into %u flows (%u ignored)"), self._ct_pkt,
                              self._trace.pkt_drops(), self._ct_shortkey,
                              self._ct_nonip, self._ct_flow, self._ct_ignored)

            flowqueue.put(SHUTDOWN_SENTINEL)

    class DummyObserver:  # pylint: disable=R0903
        """
        The dummy observer provides a class compatible with the API of the Observer
        class without actually performing any operations.  This is primarily used
        for PATHspider's test suite.
        """

        def run_flow_enqueuer(self, flowqueue, irqueue=None):  # pylint: disable=R0201
            """
            When running the flow enqueuer, no network operation is performed and
            the thread will block until given a shutdown signal. When the shutdown
            signal is received it will cascade the signal onto the flowqueue
            in the same way that a real Observer instance would.
            """

            irqueue.get()
            flowqueue.put(SHUTDOWN_SENTINEL)