stub

#!/usr/bin/env python

import broker
import broker.bro
import uuid
import sys
import time

from ipaddress import *

# -- constants ----------------------------------------------------------------

PORT = 43000
CONTROL_TOPIC = "/vast/control/"
DATA_TOPIC = "/vast/data/"

# -- utilities ----------------------------------------------------------------

def log(*args):
    """Prints logging output to standard error"""
    print('\033[36m>>>\033[0m', *args, file=sys.stderr)

# -- Bro ----------------------------------------------------------------------

class Bro:
    def __init__(self, port):
        self.endpoint = broker.Endpoint()
        # Create a subscriber and register control and data channel.
        self.subscriber = self.endpoint.make_subscriber([DATA_TOPIC])
        log("peering with VAST")
        self.endpoint.peer("127.0.0.1", port)
        log("established peering successfully")

    def query(self, expression):
        query_id = str(uuid.uuid4())
        event = broker.bro.Event("VAST::query", query_id, expression)
        self.endpoint.publish(CONTROL_TOPIC, event)
        log("performing lookup for", expression)
        while True:
            topic, data = self.subscriber.get()
            event = broker.bro.Event(data)
            log(topic, event.args())
            (qid, result) = event.args()
            if result is None:
                break;

    def run(self):
        self.query(":addr in 10.0.0.0/8")

# -- VAST ---------------------------------------------------------------------

class VAST:
    def __init__(self, port):
        self.endpoint = broker.Endpoint()
        self.subscriber = self.endpoint.make_subscriber([CONTROL_TOPIC])
        self.endpoint.listen("127.0.0.1", port)

    def lookup(self, query_id, expression):
        log("answering query '{}'".format(expression))
        name = "VAST::result"
        make_result_event = lambda *xs: broker.bro.Event(name, query_id, *xs)
        for x in self.answer(expression):
            self.endpoint.publish(DATA_TOPIC, make_result_event(x))
        self.endpoint.publish(DATA_TOPIC, make_result_event(None))

    def answer(self, expression):
        # We just generate dummy data here.
        f = lambda x: [x, time.ctime(), [IPv4Address("10.0.0.1")]]
        return map(f, range(10))

    def run(self):
        while True:
            log("waiting for commands")
            (topic, data) = self.subscriber.get()
            event = broker.bro.Event(data)
            (query_id, expression) = event.args()
            self.lookup(query_id, expression)

# -- main ---------------------------------------------------------------------

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("usage: {} <mode> [port]".format(sys.argv[0]))
        exit(1)
    mode = sys.argv[1]
    port = int(sys.argv[2]) if len(sys.argv) > 2 else PORT
    if mode == "bro":
        bro = Bro(port)
        bro.run()
    elif mode == "vast":
        vast = VAST(port)
        vast.run()
    else:
        print("invalid mode:", mode)
        exit(1)