__doc__ = """
tartiflette

Program to analyse real-time traceroute inormation for routing changes.

Usage:
    tartiflette --num_procs=<NUM>

Options:
    --num_procs=<NUM>   Number of worker processes to spin up to handle
                        load. Uses one asyncio event loop per process.
"""
import asyncio
import docopt
import ipaddress
import multiprocessing
from ripe.atlas.cousteau import AtlasStream

WORK_QUEUE = multiprocessing.Queue()

class IPMatcher(multiprocessing.Process):
    NETWORKS = {
        4: [
            ipaddress.ip_network(u'{}'.format(net.strip()), strict=False) for
            net in open('../v4.txt').readlines()
        ],
        6: [
            ipaddress.ip_network(u'{}'.format(net.strip()), strict=False) for
            net in open('../Comcast-v6-Space').readlines()
        ],
    }

    def __init__(self, queue):
        self.QUEUE = queue
        policy = asyncio.get_event_loop_policy()
        policy.set_event_loop(policy.new_event_loop())
        self.LOOP = asyncio.get_event_loop()
        super().__init__()

    @asyncio.coroutine
    def main(self):
        """Loop forever looking for work from the queue"""
        while True:
            if not self.QUEUE.empty():
                traceroute = self.QUEUE.get()
                yield from self.filter_hop_rtt(traceroute)

    def run(self):
        self.LOOP.run_until_complete(self.main())

    @asyncio.coroutine
    def filter_hop_rtt(self, traceroute):
        """Given a traceroute result, filter out the unnecessary data and
        hand off for analysis"""
        m_result = traceroute
        if 'result' in m_result.keys() and m_result['result']:
            for hop in m_result['result']:
                if not 'result' in hop.keys():
                    continue
                for address in hop['result']:
                    if 'from' in address.keys():
                        res = yield from self.is_comcast_ip(address['from'])
                        if res:
                            yield print(m_result)
                            return None

    @asyncio.coroutine
    def is_comcast_ip(self, ip_address):
        """Returns true if this is a comcast IP address"""
        address = ipaddress.ip_address(ip_address)
        for network in self.NETWORKS[address.version]:
            if address in network:
                return True
        return False


def on_result_recieved(*args):
    """Add the trqceroute result to a queue to be processed"""
    WORK_QUEUE.put(args[0])

def stream_results(filters={}):
    """Set up the atlas stream for all traceroute results"""
    atlas_stream = AtlasStream()
    atlas_stream.connect()
    atlas_stream.bind_channel('result', on_result_recieved)
    stream_parameters = {"type": "traceroute"}
    stream_parameters.update(filters)
    atlas_stream.start_stream(stream_type="result", **stream_parameters)
    atlas_stream.timeout()

if __name__ == '__main__':
    """Start up one worker process to deal with handling checking traceroute
    results, and just use the main thread to read from atlas."""
    args = docopt.docopt(__doc__)
    policy = asyncio.get_event_loop_policy()
    policy.set_event_loop(policy.new_event_loop())
    for i in range(int(args['--num_procs'])):
        proc = IPMatcher(WORK_QUEUE).start()
    stream_results()


@@ -133,16 +133,35 @@
 cs=%3CNUM%3E
+ %5B--time=%3CSECONDS%3E%5D
 %0A%0AOption
@@ -306,16 +306,133 @@
 rocess.%0A
+    --time=%3CSECONDS%3E    Number of seconds to run the analysis for. If%0A                        ommitted, run forever.%0A
 %22%22%22%0Aimpo
@@ -2665,16 +2665,30 @@
 results(
+seconds=None, 
 filters=
@@ -2692,16 +2692,16 @@
 rs=%7B%7D):%0A
-
     %22%22%22S
@@ -3056,16 +3056,61 @@
 timeout(
+seconds=seconds)%0A    atlas_stream.disconnect(
 )%0A%0Aif __
@@ -3399,24 +3399,39 @@
 ent_loop())%0A
+    procs = %5B%5D%0A
     for i in
@@ -3504,35 +3504,234 @@
 EUE)
-.start()%0A    stream_results()%0A
+%0A        procs.append(proc)%0A        proc.start()%0A    if args%5B'--time'%5D:%0A        seconds = int(args%5B'--time'%5D)%0A    else:%0A        seconds = None%0A    stream_results(seconds)%0A    for proc in procs:%0A        proc.terminate()%0A    exit()
 %0A

Add in variable timeout