Exemplo n.º 1
0
    def send_next_message(self, conn, msg):
        """
        Publish next custom BGP update via the autoconf RPC.
        """

        with nested(
                conn.Consumer(
                    on_message=self.handle_config_notify,
                    queues=[self.config_queue],
                    accept=["ujson"],
                )):
            with conn.Producer() as producer:
                print("[+] Sending message '{}'".format(msg))
                producer.publish(
                    msg,
                    exchange=self.config_exchange,
                    routing_key="autoconf-update",
                    retry=True,
                    priority=4,
                    serializer="ujson",
                )
                print("[+] Sent message '{}'".format(msg))
                try:
                    conn.drain_events(timeout=10)
                except socket.timeout:
                    # avoid infinite loop by timeout
                    assert (self.config_notify_received
                            ), "[-] Config notify consumer timeout"
                assert self.config_notify_received, "[-] Config notify consumer timeout"
            print("[+] Async received config notify")
    def receive_event(self):
        def handle_notification(body, message):
            print("Receive Event!")

            action_content = self.mapping(trigger_id)

            action_id = ''

            self.call_to_action(action_id, action_content)

        try:
            self.consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(self.consumer_connection,
                             queues=self.queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    self.consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            print('Connection lost')
        except self.consumer_connection.connection_errors:
            print('Connection error')
        except:
            print("some error")
Exemplo n.º 3
0
    def run(self):
        queue_notification = Queue(name='collector.request.notification', exchange=self.exchange,
                                   routing_key='collector.request.notification', message_ttl=20)
        queue_list_platforms = Queue(name='registry.response.collector.api_get_list_platforms', exchange=self.exchange,
                                     routing_key='registry.response.collector.api_get_list_platforms', message_ttl=20)
        queue_get_states = Queue(name='driver.response.collector.api_get_states', exchange=self.exchange,
                                 routing_key='driver.response.collector.api_get_states', message_ttl=20)

        # if self.mode == 'PULL':
        #     print("Collector use Mode: PULL Data")
        self.get_list_platforms()
        #     self.collect()

        while 1:
            try:
                self.consumer_connection.ensure_connection(max_retries=1)
                with nested(Consumer(self.consumer_connection, queues=queue_notification, callbacks=[self.handle_notification],
                                     no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_list_platforms, callbacks=[self.handle_get_list],
                                     no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_get_states,
                                     callbacks=[self.handle_collect_by_platform_id], no_ack=True)):
                    while True:
                        self.consumer_connection.drain_events()
            except (ConnectionRefusedError, exceptions.OperationalError):
                print('Connection lost')
            except self.consumer_connection.connection_errors:
                print('Connection error')
Exemplo n.º 4
0
def run():
    queue_get_states = Queue(
        name='monitor.request.collector',
        exchange=exchange,
        routing_key='monitor.request.collector')  #, message_ttl=20)
    queue_get_past_states = Queue(
        name='dbreader.respond.monitor',
        exchange=exchange,
        routing_key='dbreader.respond.monitor')  #, message_ttl=20)
    queue_get_rules = Queue(
        name='monitor.request.rule_engine',
        exchange=exchange,
        routing_key='monitor.request.rule_engine')  #, message_ttl=20)
    while 1:
        try:
            consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(consumer_connection,
                             queues=queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True),
                    Consumer(consumer_connection,
                             queues=queue_get_past_states,
                             callbacks=[handle_dbreader_notification],
                             no_ack=True),
                    Consumer(consumer_connection,
                             queues=queue_get_rules,
                             callbacks=[handle_rule],
                             no_ack=True)):
                while True:
                    consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            print('Connection lost')
        except consumer_connection.connection_errors:
            print('Connection error')
Exemplo n.º 5
0
    def run(self):

        queue_get_metric = Queue(name='dbreader.request.api_get_metric',
                                 exchange=self.exchange,
                                 routing_key='dbreader.request.api_get_metric',
                                 message_ttl=20)
        queue_get_metric_history = Queue(
            name='dbreader.request.api_get_metric_history',
            exchange=self.exchange,
            routing_key='dbreader.request.api_get_metric_history',
            message_ttl=20)
        while 1:
            try:
                self.consumer_connection.ensure_connection(max_retries=1)
                with nested(
                        Consumer(self.consumer_connection,
                                 queues=queue_get_metric,
                                 callbacks=[self.api_get_metric],
                                 no_ack=True),
                        Consumer(self.consumer_connection,
                                 queues=queue_get_metric_history,
                                 callbacks=[self.api_get_metric_history],
                                 no_ack=True)):
                    while True:
                        self.consumer_connection.drain_events()
            except (ConnectionRefusedError, exceptions.OperationalError):
                print('Connection lost')
            except self.consumer_connection.connection_errors:
                print('Connection error')
Exemplo n.º 6
0
    def receiveCallToAction(self):
        def handle_notification(body, message):
            print("Receive Event!")
            action_id = json.loads(body)["action_id"]
            action_name = json.loads(body)["action_name"]
            action_type = json.loads(body)["action_type"]

            self.execute(action_type, action_id)
            print("Executed action: ", action_id)

        # End handle_notification

        try:
            self.consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(self.consumer_connection,
                             queues=self.queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    self.consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            print('Connection lost')
        except self.consumer_connection.connection_errors:
            print('Connection error')
Exemplo n.º 7
0
    def receive_states(self):
        def handle_notification(body, message):
            print("Receive state!")
            item_id = json.loads(body)["item_id"]
            item_name = json.loads(body)["item_name"]
            item_type = json.loads(body)["item_type"]
            item_state = json.loads(body)["item_state"]
            # time = json.loads(body)[""]
            time = datetime.now()
            time = str(time)

            # Write new item to the database
            write_result = self.write_item_to_database(item_id, item_name,
                                                       item_type, item_state,
                                                       time)

            #if write new item success to the database, consider the trigger with that item
            if (write_result == 1):
                list_event_condition = self.read_event_condition()
                for event_condition in list_event_condition:
                    trigger_id = event_condition[0]
                    trigger_type = event_condition[1]
                    trigger_content = event_condition[2]

                    # Check the item with each trigger condition
                    result = self.check_trigger_condition(
                        trigger_id, trigger_type, trigger_content, item_id)

                    if (result == None):
                        return None

                    # If checkTriggerCondition success, create an event
                    if (result == True):
                        # event_source = item_id
                        # event_name = str(item_name) + "_" + str(datetime.now())
                        # event_id   = randint(1, 1000000)
                        # my_event = Event(event_name, event_id, event_source)
                        self.create_event(trigger_id)

            # End handle_notification

        try:
            self.consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(self.consumer_connection,
                             queues=self.queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    self.consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            print('Connection lost')
        except self.consumer_connection.connection_errors:
            print('Connection error')
        except:
            print("some error")
Exemplo n.º 8
0
    def test_fanout(self, name='test_mongodb_fanout'):
        if not self.verify_alive():
            return
        c = self.connection
        self.e = Exchange(name, type='fanout')
        self.q = Queue(name, exchange=self.e, routing_key=name)
        self.q2 = Queue(name + '2', exchange=self.e, routing_key=name + '2')

        channel = c.default_channel
        producer = Producer(channel, self.e)
        consumer1 = Consumer(channel, self.q)
        consumer2 = Consumer(channel, self.q2)
        self.q2(channel).declare()

        for i in range(10):
            producer.publish({'foo': i}, routing_key=name)
        for i in range(10):
            producer.publish({'foo': i}, routing_key=name + '2')

        _received1 = []
        _received2 = []

        def callback1(message_data, message):
            _received1.append(message)
            message.ack()

        def callback2(message_data, message):
            _received2.append(message)
            message.ack()

        consumer1.register_callback(callback1)
        consumer2.register_callback(callback2)

        with nested(consumer1, consumer2):

            while 1:
                if len(_received1) + len(_received2) == 20:
                    break
                c.drain_events(timeout=60)
        self.assertEqual(len(_received1) + len(_received2), 20)

        # queue.delete
        for i in range(10):
            producer.publish({'foo': i}, routing_key=name)
        self.assertTrue(self.q(channel).get())
        self.q(channel).delete()
        self.q(channel).declare()
        self.assertIsNone(self.q(channel).get())

        # queue.purge
        for i in range(10):
            producer.publish({'foo': i}, routing_key=name + '2')
        self.assertTrue(self.q2(channel).get())
        self.q2(channel).purge()
        self.assertIsNone(self.q2(channel).get())
Exemplo n.º 9
0
    def test_fanout(self, name='test_mongodb_fanout'):
        if not self.verify_alive():
            return
        c = self.connection
        self.e = Exchange(name, type='fanout')
        self.q = Queue(name, exchange=self.e, routing_key=name)
        self.q2 = Queue(name + '2', exchange=self.e, routing_key=name + '2')

        channel = c.default_channel
        producer = Producer(channel, self.e)
        consumer1 = Consumer(channel, self.q)
        consumer2 = Consumer(channel, self.q2)
        self.q2(channel).declare()

        for i in range(10):
            producer.publish({'foo': i}, routing_key=name)
        for i in range(10):
            producer.publish({'foo': i}, routing_key=name + '2')

        _received1 = []
        _received2 = []

        def callback1(message_data, message):
            _received1.append(message)
            message.ack()

        def callback2(message_data, message):
            _received2.append(message)
            message.ack()

        consumer1.register_callback(callback1)
        consumer2.register_callback(callback2)

        with nested(consumer1, consumer2):

            while 1:
                if len(_received1) + len(_received2) == 20:
                    break
                c.drain_events(timeout=60)
        self.assertEqual(len(_received1) + len(_received2), 20)

        # queue.delete
        for i in range(10):
            producer.publish({'foo': i}, routing_key=name)
        self.assertTrue(self.q(channel).get())
        self.q(channel).delete()
        self.q(channel).declare()
        self.assertIsNone(self.q(channel).get())

        # queue.purge
        for i in range(10):
            producer.publish({'foo': i}, routing_key=name + '2')
        self.assertTrue(self.q2(channel).get())
        self.q2(channel).purge()
        self.assertIsNone(self.q2(channel).get())
    def receive_event(self):
        def process_handle_notification(event_id, event_source, trigger_id):
            # mapping trigger_id to condition_id and action_id
            condition_id, condition_type, condition_content, action_id, action_content = self.mapping(
                trigger_id)

            # Check if condition has condition_id is true
            is_condition_satisfice = self.check_condition(
                condition_id, condition_type, condition_content)

            print("check condition result: ", is_condition_satisfice)

            if (is_condition_satisfice == True):
                # Execute an action
                self.call_to_action(action_id, action_content)

        def handle_notification(body, message):
            print("Receive Event!")

            event_id = json.loads(body)["event_id"]
            event_source = json.loads(body)["event_source"]
            trigger_id = json.loads(body)["trigger_id"]

            try:
                p = Process(target=process_handle_notification,
                            args=(event_id, event_source, trigger_id))
                p.start()
                p.join()
                time.sleep(1)

            except:
                print("Error: unable to start the new process")

            # End handle_notification

        try:
            self.consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(self.consumer_connection,
                             queues=self.queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    self.consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            print('Connection lost')
        except self.consumer_connection.connection_errors:
            print('Connection error')
Exemplo n.º 11
0
    def receive_event(self):
        def handle_notification(body, message):
            self.logger.info("Receive Event!")
            # event_name = json.loads(body)["event_name"]
            event_id = json.loads(body)["event_id"]
            event_source = json.loads(body)["event_source"]
            trigger_id = json.loads(body)["trigger_id"]
            # event_generator_id = json.loads(body)["event_generator_id"]
            time = json.loads(body)["time"]

            # mapping trigger_id to condition_id and action_id
            condition_id, condition_type, condition_content, action_id, action_content = self.mapping(
                trigger_id)

            # Check if condition has condition_id is true
            is_condition_satisfice = self.check_condition(
                condition_id, condition_type, condition_content)

            self.logger.info("check condition result: " +
                             str(is_condition_satisfice))

            if (is_condition_satisfice == True):
                # Execute an action
                self.call_to_action(action_id, action_content)

            self.logger.info("Finish receive one event!\n\n\n")
            # End handle_notification

        try:
            self.consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(self.consumer_connection,
                             queues=self.queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    self.consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            self.logger.error("Connection lost", exc_info=True)
        except self.consumer_connection.connection_errors:
            self.logger.error("Connection error", exc_info=True)
        except Exception as e:
            self.logger.error("Error receive event", exc_info=True)
Exemplo n.º 12
0
def run():
    queue_notification = Queue(name='monitor.request.alert',
                               exchange=exchange,
                               routing_key='monitor.request.alert')

    while 1:
        try:
            consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(consumer_connection,
                             queues=queue_notification,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            print('Connection lost')
        except consumer_connection.connection_errors:
            print('Connection error')
Exemplo n.º 13
0
    def run(self):
        self.client_fog.on_disconnect = self.on_disconnect
        self.client_fog.on_connect = self.on_connect

        # declare rabbitmq resources
        queue_add_platform = Queue(
            name='registry.response.driver.api_add_platform',
            exchange=self.exchange,
            routing_key='registry.response.driver.api_add_platform')
        queue_check_config = Queue(
            name='driver.request.api_check_configuration_changes',
            exchange=self.exchange,
            routing_key='driver.request.api_check_configuration_changes')
        queue_collect = Queue(name='driver.request.api_get_states',
                              exchange=self.exchange,
                              routing_key='driver.request.api_get_states')
        queue_set_state = Queue(name='driver.request.api_set_state',
                                exchange=self.exchange,
                                routing_key='driver.request.api_set_state')

        while 1:
            try:
                self.rabbitmq_connection.ensure_connection(max_retries=3)
                with nested(
                        Consumer(self.rabbitmq_connection,
                                 queues=queue_add_platform,
                                 callbacks=[self.on_message_add_platform]),
                        Consumer(self.rabbitmq_connection,
                                 queues=queue_check_config,
                                 callbacks=[self.on_message_check_config]),
                        Consumer(self.rabbitmq_connection,
                                 queues=queue_collect,
                                 callbacks=[self.on_message_collect]),
                        Consumer(self.rabbitmq_connection,
                                 queues=queue_set_state,
                                 callbacks=[self.on_message_set_state])):
                    while True:
                        self.rabbitmq_connection.drain_events()
            except (ConnectionRefusedError, exceptions.OperationalError):
                print('Connection lost')
            except self.rabbitmq_connection.connection_errors:
                print('Connection error')
Exemplo n.º 14
0
    def run(self):
        queue_get_things = Queue(name='registry.request.api_get_things', exchange=self.exchange,
                                 routing_key='registry.request.api_get_things', ttl=20)
        queue_get_list_platforms = Queue(name='registry.request.api_get_list_platforms', exchange=self.exchange,
                                         routing_key='registry.request.api_get_list_platforms')
        queue_add_platform = Queue(name='registry.request.api_add_platform', exchange=self.exchange,
                                   routing_key='registry.request.api_add_platform')
        queue_check_config = Queue(name='driver.response.registry.api_check_configuration_changes', exchange=self.exchange,
                                   routing_key='driver.response.registry.api_check_configuration_changes')
        queue_get_thing_by_global_id = Queue(name='registry.request.api_get_thing_by_global_id', exchange=self.exchange,
                                             routing_key='registry.request.api_get_thing_by_global_id')
        queue_get_things_by_platform_id = Queue(name='registry.request.api_get_things_by_platform_id',
                                                exchange=self.exchange,
                                                routing_key='registry.request.api_get_things_by_platform_id')

        if self.mode == 'PULL':
            self.update_all_config_changes()

        self.check_platform_active()

        while 1:
            try:
                self.consumer_connection.ensure_connection(max_retries=1)
                with nested(Consumer(self.consumer_connection, queues=queue_get_things_by_platform_id,
                                     callbacks=[self.api_get_things_by_platform_id], no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_get_thing_by_global_id,
                                     callbacks=[self.api_get_thing_by_global_id], no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_add_platform, callbacks=[self.api_add_platform],
                                     no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_get_things, callbacks=[self.api_get_things],
                                     no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_get_list_platforms,
                                     callbacks=[self.api_get_list_platforms], no_ack=True),
                            Consumer(self.consumer_connection, queues=queue_check_config,
                                     callbacks=[self.handle_configuration_changes], no_ack=True)):
                    while True:
                        self.consumer_connection.drain_events()
            except (ConnectionRefusedError, exceptions.OperationalError):
                print('Connection lost')
            except self.consumer_connection.connection_errors:
                print('Connection error')
Exemplo n.º 15
0
 def run(self):
     queues = []
     for queue, routing_key, function in self.app.queues:
         exchange = Exchange(self.app.conf.EXCHANGE_NAME,
                             'topic',
                             durable=True)
         queue_name = '{}_{}'.format(self.app.SERVICE_NAME, queue)
         queue = Queue(queue_name,
                       exchange=exchange,
                       routing_key=routing_key)
         queues.append((queue, function))
     with Connection(self.app.conf.BROKER_URL) as connection:
         with connection.channel() as channel:
             consumers = []
             for queue, function in queues:
                 consumer = Consumer(channel,
                                     queue,
                                     callback=function,
                                     accept=['json'])
                 consumers.append(consumer)
             with nested(*consumers):
                 connection.drain_events(timeout=1)
Exemplo n.º 16
0
consumer_connection = Connection(BROKER_CLOUD)
MODE = "PUSH"  # PUSH or PULL

exchange = Exchange("IOT", type="direct")


def handle_notification(body, message):
    # receive list items

    print('Receive message from Alert')
    print(body)
    time.sleep(3)


test = Queue(name='alert_in', exchange=exchange, routing_key='alert_in')

while 1:
    try:
        consumer_connection.ensure_connection(max_retries=1)
        with nested(
                Consumer(consumer_connection,
                         queues=test,
                         callbacks=[handle_notification],
                         no_ack=True)):
            while True:
                consumer_connection.drain_events()
    except (ConnectionRefusedError, exceptions.OperationalError):
        print('Connection lost')
    except consumer_connection.connection_errors:
        print('Connection error')
Exemplo n.º 17
0
    def test(self):
        """
        Loads a test file that includes crafted bgp updates as
        input and expected messages as output.
        """
        RABBITMQ_USER = os.getenv("RABBITMQ_USER", "guest")
        RABBITMQ_PASS = os.getenv("RABBITMQ_PASS", "guest")
        RABBITMQ_HOST = os.getenv("RABBITMQ_HOST", "rabbitmq")
        RABBITMQ_PORT = os.getenv("RABBITMQ_PORT", 5672)
        RABBITMQ_URI = "amqp://{}:{}@{}:{}//".format(RABBITMQ_USER,
                                                     RABBITMQ_PASS,
                                                     RABBITMQ_HOST,
                                                     RABBITMQ_PORT)
        RPKI_VALIDATOR_HOST = os.getenv("RPKI_VALIDATOR_HOST", "routinator")
        RPKI_VALIDATOR_PORT = os.getenv("RPKI_VALIDATOR_PORT", 3323)

        # check RPKI RTR manager connectivity
        while True:
            try:
                rtrmanager = RTRManager(RPKI_VALIDATOR_HOST,
                                        RPKI_VALIDATOR_PORT)
                rtrmanager.start()
                print("Connected to RPKI VALIDATOR '{}:{}'".format(
                    RPKI_VALIDATOR_HOST, RPKI_VALIDATOR_PORT))
                rtrmanager.stop()
                break
            except Exception:
                print("Could not connect to RPKI VALIDATOR '{}:{}'".format(
                    RPKI_VALIDATOR_HOST, RPKI_VALIDATOR_PORT))
                print("Retrying in 30 seconds...")
                time.sleep(30)

        # exchanges
        self.update_exchange = Exchange("bgp-update",
                                        type="direct",
                                        durable=False,
                                        delivery_mode=1)

        self.hijack_exchange = Exchange("hijack-update",
                                        type="direct",
                                        durable=False,
                                        delivery_mode=1)

        self.pg_amq_bridge = Exchange("amq.direct",
                                      type="direct",
                                      durable=True,
                                      delivery_mode=1)

        # queues
        self.update_queue = Queue(
            "detection-testing",
            exchange=self.pg_amq_bridge,
            routing_key="update-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_queue = Queue(
            "hijack-testing",
            exchange=self.hijack_exchange,
            routing_key="update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_db_queue = Queue(
            "hijack-db-testing",
            exchange=self.pg_amq_bridge,
            routing_key="hijack-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        with Connection(RABBITMQ_URI) as connection:
            print("Waiting for pg_amq exchange..")
            Tester.waitExchange(self.pg_amq_bridge, connection.default_channel)
            print("Waiting for hijack exchange..")
            Tester.waitExchange(self.hijack_exchange,
                                connection.default_channel)
            print("Waiting for update exchange..")
            Tester.waitExchange(self.update_exchange,
                                connection.default_channel)

            # query database for the states of the processes
            db_con = self.getDbConnection()
            db_cur = db_con.cursor()
            query = "SELECT name FROM process_states WHERE running=True"
            running_modules = set()
            # wait until all 5 modules are running
            while len(running_modules) < 5:
                db_cur.execute(query)
                entries = db_cur.fetchall()
                for entry in entries:
                    running_modules.add(entry[0])
                db_con.commit()
                print("Running modules: {}".format(running_modules))
                print("{}/5 modules are running.".format(len(running_modules)))
                time.sleep(1)

            Tester.config_request_rpc(connection)

            time.sleep(10)

            for testfile in os.listdir("testfiles/"):
                self.clear()

                self.curr_test = testfile
                self.messages = {}
                # load test
                with open("testfiles/{}".format(testfile), "r") as f:
                    self.messages = json.load(f)

                send_len = len(self.messages)

                with nested(
                        connection.Consumer(
                            self.hijack_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson"],
                        ),
                        connection.Consumer(
                            self.update_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson", "txtjson"],
                        ),
                        connection.Consumer(
                            self.hijack_db_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson", "txtjson"],
                        ),
                ):
                    send_cnt = 0
                    # send and validate all messages in the messages.json file
                    while send_cnt < send_len:
                        self.curr_idx = send_cnt
                        self.send_next_message(connection)
                        send_cnt += 1
                        # sleep until we receive all expected messages
                        while self.curr_idx != send_cnt:
                            time.sleep(0.1)
                            try:
                                connection.drain_events(timeout=10)
                            except socket.timeout:
                                # avoid infinite loop by timeout
                                assert False, "Consumer timeout"

            connection.close()

        time.sleep(5)
        self.supervisor.supervisor.stopAllProcesses()

        self.waitProcess("listener", 0)  # 0 STOPPED
        self.waitProcess("clock", 0)  # 0 STOPPED
        self.waitProcess("detection", 0)  # 0 STOPPED
        self.waitProcess("configuration", 0)  # 0 STOPPED
        self.waitProcess("database", 0)  # 0 STOPPED
        self.waitProcess("observer", 0)  # 0 STOPPED
Exemplo n.º 18
0
    def test(self):
        with Connection(RABBITMQ_URI) as connection:
            # exchanges
            self.update_exchange = Exchange("bgp-update",
                                            type="direct",
                                            durable=False,
                                            delivery_mode=1)

            self.pg_amq_bridge = Exchange("amq.direct",
                                          type="direct",
                                          durable=True,
                                          delivery_mode=1)

            # queues
            self.update_queue = Queue(
                "detection-testing",
                exchange=self.pg_amq_bridge,
                routing_key="update-update",
                durable=False,
                auto_delete=True,
                max_priority=1,
                consumer_arguments={"x-priority": 1},
            )

            self.hijack_db_queue = Queue(
                "hijack-db-testing",
                exchange=self.pg_amq_bridge,
                routing_key="hijack-update",
                durable=False,
                auto_delete=True,
                max_priority=1,
                consumer_arguments={"x-priority": 1},
            )

            print("Waiting for pg_amq exchange..")
            AutoignoreTester.waitExchange(self.pg_amq_bridge,
                                          connection.default_channel)
            print("Waiting for update exchange..")
            AutoignoreTester.waitExchange(self.update_exchange,
                                          connection.default_channel)

            # query database for the states of the processes
            db_con = self.getDbConnection()
            db_cur = db_con.cursor()
            query = "SELECT name FROM process_states WHERE running=True"
            running_modules = set()
            # wait until all 6 modules are running
            while len(running_modules) < 6:
                db_cur.execute(query)
                entries = db_cur.fetchall()
                for entry in entries:
                    running_modules.add(entry[0])
                db_con.commit()
                print("[+] Running modules: {}".format(running_modules))
                print("[+] {}/6 modules are running. Re-executing query...".
                      format(len(running_modules)))
                time.sleep(1)

            AutoignoreTester.config_request_rpc(connection)

            time.sleep(5)

            db_cur.close()
            db_con.close()

            for testfile in os.listdir("testfiles/"):
                self.clear()

                self.curr_test = testfile
                self.messages = {}

                # load test
                with open("testfiles/{}".format(testfile), "r") as f:
                    self.messages = json.load(f)

                send_len = len(self.messages)

                with nested(
                        connection.Consumer(
                            self.hijack_db_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson", "txtjson"],
                        )):
                    send_cnt = 0
                    # send and validate all messages in the messages.json file
                    while send_cnt < send_len:
                        self.curr_idx = send_cnt
                        self.send_next_message(connection)
                        send_cnt += 1
                        # sleep until we receive all expected messages
                        while self.curr_idx != send_cnt:
                            time.sleep(0.1)
                            try:
                                connection.drain_events(timeout=60)
                            except socket.timeout:
                                # avoid infinite loop by timeout
                                assert False, "Consumer timeout"
                        # sleep for at least 20 seconds between messages so that we check that the ignore mechanism
                        # is not triggered by mistake
                        if send_cnt < send_len:
                            print(
                                "[+] Sleeping for 20 seconds to ensure auto-ignore works correctly"
                            )
                            time.sleep(20)

            connection.close()

        print("[+] Sleeping for 5 seconds...")
        time.sleep(5)
        print("[+] Instructing all processes to stop...")
        self.supervisor.supervisor.stopAllProcesses()

        self.waitProcess("autoignore", 0)  # 0 STOPPED
        self.waitProcess("listener", 0)  # 0 STOPPED
        self.waitProcess("clock", 0)  # 0 STOPPED
        self.waitProcess("configuration", 0)  # 0 STOPPED
        self.waitProcess("database", 0)  # 0 STOPPED
        self.waitProcess("observer", 0)  # 0 STOPPED
        self.waitProcess("detection", 0)  # 0 STOPPED
        print(
            "[+] All processes (listener, clock, conf, db, detection, autoignore and observer) are stopped."
        )
Exemplo n.º 19
0
    def send_next_message(self, conn, msg):
        """
        Publish next custom BGP update via the autoconf RPC.
        """
        correlation_id = uuid()
        callback_queue = Queue(
            uuid(),
            channel=conn.default_channel,
            durable=False,
            auto_delete=True,
            max_priority=4,
            consumer_arguments={"x-priority": 4},
        )

        with nested(
                conn.Consumer(
                    on_message=self.handle_config_notify,
                    queues=[self.config_queue],
                    accept=["ujson"],
                ),
                conn.Consumer(
                    on_message=self.handle_autoconf_update_goahead_reply,
                    queues=[callback_queue],
                    accept=["ujson"],
                ),
        ):
            self.autoconf_goahead = False
            with conn.Producer() as producer:
                print("[+] Sending message '{}'".format(msg))
                producer.publish(
                    msg,
                    exchange="",
                    routing_key="conf-autoconf-update-queue",
                    reply_to=callback_queue.name,
                    correlation_id=correlation_id,
                    retry=True,
                    declare=[
                        Queue(
                            "conf-autoconf-update-queue",
                            durable=False,
                            max_priority=4,
                            consumer_arguments={"x-priority": 4},
                        ),
                        callback_queue,
                    ],
                    priority=4,
                    serializer="ujson",
                )
                print("[+] Sent message '{}'".format(msg))
                conn.drain_events()
                try:
                    conn.drain_events(timeout=10)
                except socket.timeout:
                    # avoid infinite loop by timeout
                    assert self.autoconf_goahead, "[-] Autoconf consumer timeout"
                print("[+] Concluded autoconf RPC")
                try:
                    conn.drain_events(timeout=10)
                except socket.timeout:
                    # avoid infinite loop by timeout
                    assert (self.config_notify_received
                            ), "[-] Config notify consumer timeout"
                print("[+] Async received config notify")
Exemplo n.º 20
0
    def test(self):
        """
        Loads a test file that includes crafted bgp updates as
        input and expected messages as output.
        """
        RABBITMQ_USER = os.getenv("RABBITMQ_USER", "guest")
        RABBITMQ_PASS = os.getenv("RABBITMQ_PASS", "guest")
        RABBITMQ_HOST = os.getenv("RABBITMQ_HOST", "rabbitmq")
        RABBITMQ_PORT = os.getenv("RABBITMQ_PORT", 5672)
        RABBITMQ_URI = "amqp://{}:{}@{}:{}//".format(RABBITMQ_USER,
                                                     RABBITMQ_PASS,
                                                     RABBITMQ_HOST,
                                                     RABBITMQ_PORT)

        # exchanges
        self.update_exchange = Exchange("bgp-update",
                                        type="direct",
                                        durable=False,
                                        delivery_mode=1)

        self.hijack_exchange = Exchange("hijack-update",
                                        type="direct",
                                        durable=False,
                                        delivery_mode=1)

        self.pg_amq_bridge = Exchange("amq.direct",
                                      type="direct",
                                      durable=True,
                                      delivery_mode=1)

        # queues
        self.update_queue = Queue(
            "detection-testing",
            exchange=self.pg_amq_bridge,
            routing_key="update-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_queue = Queue(
            "hijack-testing",
            exchange=self.hijack_exchange,
            routing_key="update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_db_queue = Queue(
            "hijack-db-testing",
            exchange=self.pg_amq_bridge,
            routing_key="hijack-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        with Connection(RABBITMQ_URI) as connection:
            print("Waiting for pg_amq exchange..")
            Tester.waitExchange(self.pg_amq_bridge, connection.default_channel)
            print("Waiting for hijack exchange..")
            Tester.waitExchange(self.hijack_exchange,
                                connection.default_channel)
            print("Waiting for update exchange..")
            Tester.waitExchange(self.update_exchange,
                                connection.default_channel)

            # query database for the states of the processes
            db_con = self.getDbConnection()
            db_cur = db_con.cursor()
            query = "SELECT name FROM process_states WHERE running=True"
            running_modules = set()
            # wait until all 6 modules are running
            while len(running_modules) < 6:
                db_cur.execute(query)
                entries = db_cur.fetchall()
                for entry in entries:
                    running_modules.add(entry[0])
                db_con.commit()
                print("Running modules: {}".format(running_modules))
                print("{}/6 modules are running.".format(len(running_modules)))
                time.sleep(1)

            Tester.config_request_rpc(connection)

            time.sleep(10)

            # call all helper functions
            Helper.hijack_resolve(db_con, connection, "a", "139.5.46.0/24",
                                  "S|0|-|-", 133720)
            Helper.hijack_mitigate(db_con, connection, "b", "10.91.236.0/24")
            Helper.hijack_ignore(db_con, connection, "c", "139.5.237.0/24",
                                 "S|0|-|-", 136334)
            Helper.hijack_comment(db_con, connection, "d", "test")
            Helper.hijack_ack(db_con, connection, "e", "true")
            Helper.hijack_multiple_action(db_con, connection, ["f", "g"],
                                          "hijack_action_acknowledge")
            Helper.hijack_multiple_action(db_con, connection, ["f", "g"],
                                          "hijack_action_acknowledge_not")
            Helper.hijack_multiple_action(db_con, connection, ["f"],
                                          "hijack_action_resolve")
            Helper.hijack_multiple_action(db_con, connection, ["g"],
                                          "hijack_action_ignore")
            # multi-action delete a hijack purged from cache
            Helper.hijack_multiple_action(db_con, connection, ["f"],
                                          "hijack_action_delete")
            # delete a hijack purged from cache
            Helper.hijack_delete(db_con, connection, "g", "139.5.16.0/22",
                                 "S|0|-|-", 133676)
            # multi-action delete a hijack using cache
            Helper.hijack_multiple_action(db_con, connection, ["h"],
                                          "hijack_action_delete")
            # delete a hijack using cache
            Helper.hijack_delete(db_con, connection, "i", "139.5.24.0/24",
                                 "S|0|-|-", 133720)
            Helper.hijack_mitigate(db_con, connection, "j",
                                   "2001:db8:abcd:12::0/80")
            Helper.load_as_sets(connection)

            time.sleep(10)

            db_cur.close()
            db_con.close()

            for testfile in os.listdir("testfiles/"):
                self.clear()

                self.curr_test = testfile
                self.messages = {}
                # load test
                with open("testfiles/{}".format(testfile), "r") as f:
                    self.messages = json.load(f)

                send_len = len(self.messages)

                with nested(
                        connection.Consumer(
                            self.hijack_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson"],
                        ),
                        connection.Consumer(
                            self.update_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson", "txtjson"],
                        ),
                        connection.Consumer(
                            self.hijack_db_queue,
                            callbacks=[self.validate_message],
                            accept=["ujson", "txtjson"],
                        ),
                ):
                    send_cnt = 0
                    # send and validate all messages in the messages.json file
                    while send_cnt < send_len:
                        self.curr_idx = send_cnt
                        self.send_next_message(connection)
                        send_cnt += 1
                        # sleep until we receive all expected messages
                        while self.curr_idx != send_cnt:
                            time.sleep(0.1)
                            try:
                                connection.drain_events(timeout=10)
                            except socket.timeout:
                                # avoid infinite loop by timeout
                                assert False, "Consumer timeout"

            connection.close()

        with open("configs/config.yaml") as f1, open(
                "configs/config2.yaml") as f2:
            new_data = f2.read()
            old_data = f1.read()

        Helper.change_conf(connection, new_data, old_data, "test")

        time.sleep(5)
        self.supervisor.supervisor.stopAllProcesses()

        self.waitProcess("listener", 0)  # 0 STOPPED
        self.waitProcess("clock", 0)  # 0 STOPPED
        self.waitProcess("detection", 0)  # 0 STOPPED
        self.waitProcess("mitigation", 0)  # 0 STOPPED
        self.waitProcess("configuration", 0)  # 0 STOPPED
        self.waitProcess("database", 0)  # 0 STOPPED
        self.waitProcess("observer", 0)  # 0 STOPPED
Exemplo n.º 21
0
                      exchange=exchange,
                      routing_key='driver.request.api_get_states')
queue_set_state = Queue(name='driver.request.api_set_state',
                        exchange=exchange,
                        routing_key='driver.request.api_set_state')
# queue_monitor = Queue(name='driver.request.api_get_states', exchange=exchange, routing_key='driver.request.api_get_states')

while 1:
    try:
        rabbitmq_connection.ensure_connection(max_retries=3)
        with nested(
                Consumer(rabbitmq_connection,
                         queues=queue_add_platform,
                         callbacks=[on_message_add_platform]),
                Consumer(rabbitmq_connection,
                         queues=queue_check_config,
                         callbacks=[on_message_check_config]),
                Consumer(rabbitmq_connection,
                         queues=queue_collect,
                         callbacks=[on_message_collect]),
                # Consumer(rabbitmq_connection, queonitorues=queue_monitor, callbacks=[on_message_monitor]),
                Consumer(rabbitmq_connection,
                         queues=queue_set_state,
                         callbacks=[on_message_set_state])):
            while True:
                rabbitmq_connection.drain_events()
    except (ConnectionRefusedError, exceptions.OperationalError) as hihi:
        print('Connection lost')
    except rabbitmq_connection.connection_errors:
        print('Connection error')
Exemplo n.º 22
0
    def test(self):
        """
        Loads a test file that includes crafted bgp updates as
        input and expected messages as output.
        """
        RABBITMQ_USER = os.getenv("RABBITMQ_USER", "guest")
        RABBITMQ_PASS = os.getenv("RABBITMQ_PASS", "guest")
        RABBITMQ_HOST = os.getenv("RABBITMQ_HOST", "rabbitmq")
        RABBITMQ_PORT = os.getenv("RABBITMQ_PORT", 5672)
        RABBITMQ_URI = "amqp://{}:{}@{}:{}//".format(
            RABBITMQ_USER, RABBITMQ_PASS, RABBITMQ_HOST, RABBITMQ_PORT
        )
        RPKI_VALIDATOR_HOST = os.getenv("RPKI_VALIDATOR_HOST", "routinator")
        RPKI_VALIDATOR_PORT = os.getenv("RPKI_VALIDATOR_PORT", 3323)

        # check RPKI RTR manager connectivity
        while True:
            try:
                rtrmanager = RTRManager(RPKI_VALIDATOR_HOST, RPKI_VALIDATOR_PORT)
                rtrmanager.start()
                print(
                    "Connected to RPKI VALIDATOR '{}:{}'".format(
                        RPKI_VALIDATOR_HOST, RPKI_VALIDATOR_PORT
                    )
                )
                rtrmanager.stop()
                break
            except Exception:
                print(
                    "Could not connect to RPKI VALIDATOR '{}:{}'".format(
                        RPKI_VALIDATOR_HOST, RPKI_VALIDATOR_PORT
                    )
                )
                print("Retrying in 30 seconds...")
                time.sleep(30)

        # exchanges
        self.update_exchange = Exchange(
            "bgp-update", type="direct", durable=False, delivery_mode=1
        )

        self.hijack_exchange = Exchange(
            "hijack-update", type="direct", durable=False, delivery_mode=1
        )

        self.pg_amq_bridge = Exchange(
            "amq.direct", type="direct", durable=True, delivery_mode=1
        )

        # queues
        self.update_queue = Queue(
            "detection-testing",
            exchange=self.pg_amq_bridge,
            routing_key="update-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_queue = Queue(
            "hijack-testing",
            exchange=self.hijack_exchange,
            routing_key="update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_db_queue = Queue(
            "hijack-db-testing",
            exchange=self.pg_amq_bridge,
            routing_key="hijack-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        with Connection(RABBITMQ_URI) as connection:
            # print("Waiting for pg_amq exchange..")
            # Tester.waitExchange(self.pg_amq_bridge, connection.default_channel)
            # print("Waiting for hijack exchange..")
            # Tester.waitExchange(self.hijack_exchange, connection.default_channel)
            # print("Waiting for update exchange..")
            # Tester.waitExchange(self.update_exchange, connection.default_channel)

            # wait for dependencies data workers to start
            wait_data_worker_dependencies(DATA_WORKER_DEPENDENCIES)

            while True:
                try:
                    r = requests.get(
                        "http://{}:{}/config".format(CONFIGURATION_HOST, REST_PORT)
                    )
                    result = r.json()
                    assert len(result) > 0
                    break
                except Exception:
                    print("exception")
                time.sleep(1)

            time.sleep(1)

            for testfile in os.listdir("testfiles/"):
                self.clear()

                self.curr_test = testfile
                self.messages = {}
                # load test
                with open("testfiles/{}".format(testfile), "r") as f:
                    self.messages = json.load(f)

                send_len = len(self.messages)

                with nested(
                    connection.Consumer(
                        self.hijack_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson"],
                    ),
                    connection.Consumer(
                        self.update_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson", "txtjson"],
                    ),
                    connection.Consumer(
                        self.hijack_db_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson", "txtjson"],
                    ),
                ):
                    send_cnt = 0
                    # send and validate all messages in the messages.json file
                    while send_cnt < send_len:
                        self.curr_idx = send_cnt
                        self.send_next_message(connection)
                        send_cnt += 1
                        # sleep until we receive all expected messages
                        while self.curr_idx != send_cnt:
                            time.sleep(0.1)
                            try:
                                connection.drain_events(timeout=10)
                            except socket.timeout:
                                # avoid infinite loop by timeout
                                assert False, "Consumer timeout"
Exemplo n.º 23
0
    def test(self):
        # exchanges
        self.update_exchange = Exchange(
            "bgp-update", type="direct", durable=False, delivery_mode=1
        )

        self.pg_amq_bridge = Exchange(
            "amq.direct", type="direct", durable=True, delivery_mode=1
        )

        # queues
        self.update_queue = Queue(
            "detection-testing",
            exchange=self.pg_amq_bridge,
            routing_key="update-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_db_queue = Queue(
            "hijack-db-testing",
            exchange=self.pg_amq_bridge,
            routing_key="hijack-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        with Connection(RABBITMQ_URI) as connection:
            # print("Waiting for pg_amq exchange..")
            # AutoignoreTester.waitExchange(
            #     self.pg_amq_bridge, connection.default_channel
            # )
            # print("Waiting for update exchange..")
            # AutoignoreTester.waitExchange(
            #     self.update_exchange, connection.default_channel
            # )

            # wait for dependencies data workers to start
            wait_data_worker_dependencies(DATA_WORKER_DEPENDENCIES)

            while True:
                try:
                    r = requests.get(
                        "http://{}:{}/config".format(CONFIGURATION_HOST, REST_PORT)
                    )
                    result = r.json()
                    assert len(result) > 0
                    break
                except Exception:
                    print("exception")
                time.sleep(1)

            time.sleep(1)

            for testfile in os.listdir("testfiles/"):
                self.clear()

                self.curr_test = testfile
                self.messages = {}

                # load test
                with open("testfiles/{}".format(testfile), "r") as f:
                    self.messages = json.load(f)

                send_len = len(self.messages)

                with nested(
                    connection.Consumer(
                        self.hijack_db_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson", "txtjson"],
                    )
                ):
                    send_cnt = 0
                    # send and validate all messages in the messages.json file
                    while send_cnt < send_len:
                        self.curr_idx = send_cnt
                        self.send_next_message(connection)
                        send_cnt += 1
                        # sleep until we receive all expected messages
                        while self.curr_idx != send_cnt:
                            time.sleep(0.1)
                            try:
                                connection.drain_events(timeout=60)
                            except socket.timeout:
                                # avoid infinite loop by timeout
                                assert False, "Consumer timeout"
                        # sleep for at least 20 = 2 x interval seconds between messages so that we check that the ignore mechanism
                        # is not triggered by mistake
                        if send_cnt < send_len:
                            print(
                                "[+] Sleeping for 20 seconds to ensure auto-ignore works correctly"
                            )
                            time.sleep(20)

            connection.close()

        print("[+] Sleeping for 5 seconds...")
        time.sleep(5)
Exemplo n.º 24
0
    def run(self):

        queue_get_sources = Queue(
            name='registry.request.api_get_sources',
            exchange=self.exchange,
            routing_key='registry.request.api_get_sources',
            message_ttl=20)
        queue_get_list_platforms = Queue(
            name='registry.request.api_get_list_platforms',
            exchange=self.exchange,
            routing_key='registry.request.api_get_list_platforms',
            message_ttl=20)
        queue_add_platform = Queue(
            name='registry.request.api_add_platform',
            exchange=self.exchange,
            routing_key='registry.request.api_add_platform',
            message_ttl=20)
        queue_check_config = Queue(
            name='driver.response.registry.api_check_configuration_changes',
            exchange=self.exchange,
            routing_key=
            'driver.response.registry.api_check_configuration_changes',
            message_ttl=20)
        queue_check_platform_active = Queue(
            name='driver.response.registry.api_check_platform_active',
            exchange=self.exchange,
            routing_key='driver.response.registry.api_check_platform_active',
            message_ttl=20)

        thread_check_active = threading.Thread(
            target=self.check_platform_active)
        thread_check_active.setDaemon(True)
        thread_check_active.start()

        while 1:
            try:
                self.consumer_connection.ensure_connection(max_retries=1)
                with nested(
                        Consumer(self.consumer_connection,
                                 queues=queue_add_platform,
                                 callbacks=[self.api_add_platform],
                                 no_ack=True),
                        Consumer(self.consumer_connection,
                                 queues=queue_get_sources,
                                 callbacks=[self.api_get_sources],
                                 no_ack=True),
                        Consumer(self.consumer_connection,
                                 queues=queue_get_list_platforms,
                                 callbacks=[self.api_get_list_platforms],
                                 no_ack=True),
                        Consumer(self.consumer_connection,
                                 queues=queue_check_config,
                                 callbacks=[self.handle_configuration_changes],
                                 no_ack=True),
                        Consumer(self.consumer_connection,
                                 queues=queue_check_platform_active,
                                 callbacks=[self.handle_check_platform_active],
                                 no_ack=True)):
                    while True:
                        self.consumer_connection.drain_events()
            except (ConnectionRefusedError, exceptions.OperationalError):
                self.logger.error('Connection to Broker Cloud is lost')
            except self.consumer_connection.connection_errors:
                self.logger.error('Connection to Broker Cloud is error')
Exemplo n.º 25
0
    def test(self):
        """
        Loads a test file that includes crafted bgp updates as
        input and expected messages as output.
        """
        RABBITMQ_USER = os.getenv("RABBITMQ_USER", "guest")
        RABBITMQ_PASS = os.getenv("RABBITMQ_PASS", "guest")
        RABBITMQ_HOST = os.getenv("RABBITMQ_HOST", "rabbitmq")
        RABBITMQ_PORT = os.getenv("RABBITMQ_PORT", 5672)
        RABBITMQ_URI = "amqp://{}:{}@{}:{}//".format(
            RABBITMQ_USER, RABBITMQ_PASS, RABBITMQ_HOST, RABBITMQ_PORT
        )

        # exchanges
        self.update_exchange = Exchange(
            "bgp-update", type="direct", durable=False, delivery_mode=1
        )

        self.hijack_exchange = Exchange(
            "hijack-update", type="direct", durable=False, delivery_mode=1
        )

        self.pg_amq_bridge = Exchange(
            "amq.direct", type="direct", durable=True, delivery_mode=1
        )

        # queues
        self.update_queue = Queue(
            "detection-testing",
            exchange=self.pg_amq_bridge,
            routing_key="update-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_queue = Queue(
            "hijack-testing",
            exchange=self.hijack_exchange,
            routing_key="update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        self.hijack_db_queue = Queue(
            "hijack-db-testing",
            exchange=self.pg_amq_bridge,
            routing_key="hijack-update",
            durable=False,
            auto_delete=True,
            max_priority=1,
            consumer_arguments={"x-priority": 1},
        )

        with Connection(RABBITMQ_URI) as connection:
            # print("Waiting for pg_amq exchange..")
            # Tester.waitExchange(self.pg_amq_bridge, connection.default_channel)
            # print("Waiting for hijack exchange..")
            # Tester.waitExchange(self.hijack_exchange, connection.default_channel)
            # print("Waiting for update exchange..")
            # Tester.waitExchange(self.update_exchange, connection.default_channel)

            # establish connection to DB
            db_con = self.getDbConnection()
            db_cur = db_con.cursor()

            # wait for dependencies data workers to start
            wait_data_worker_dependencies(DATA_WORKER_DEPENDENCIES)

            while True:
                try:
                    r = requests.get(
                        "http://{}:{}/config".format(CONFIGURATION_HOST, REST_PORT)
                    )
                    result = r.json()
                    assert len(result) > 0
                    break
                except Exception:
                    print("exception")
                time.sleep(1)

            time.sleep(1)

            # call all helper functions
            Helper.hijack_resolve(
                db_con, connection, "a", "139.5.46.0/24", "S|0|-|-", 133720
            )
            Helper.hijack_mitigate(db_con, connection, "b", "10.91.236.0/24")
            Helper.hijack_ignore(
                db_con, connection, "c", "139.5.237.0/24", "S|0|-|-", 136334
            )
            Helper.hijack_comment(db_con, connection, "d", "test")
            Helper.hijack_ack(db_con, connection, "e", "true")
            Helper.hijack_multiple_action(
                db_con, connection, ["f", "g"], "hijack_action_acknowledge"
            )
            Helper.hijack_multiple_action(
                db_con, connection, ["f", "g"], "hijack_action_acknowledge_not"
            )
            Helper.hijack_multiple_action(
                db_con, connection, ["f"], "hijack_action_resolve"
            )
            Helper.hijack_multiple_action(
                db_con, connection, ["g"], "hijack_action_ignore"
            )
            # multi-action delete a hijack purged from cache
            Helper.hijack_multiple_action(
                db_con, connection, ["f"], "hijack_action_delete"
            )
            # delete a hijack purged from cache
            Helper.hijack_delete(
                db_con, connection, "g", "139.5.16.0/22", "S|0|-|-", 133676
            )
            # multi-action delete a hijack using cache
            Helper.hijack_multiple_action(
                db_con, connection, ["h"], "hijack_action_delete"
            )
            # delete a hijack using cache
            Helper.hijack_delete(
                db_con, connection, "i", "139.5.24.0/24", "S|0|-|-", 133720
            )
            Helper.hijack_mitigate(db_con, connection, "j", "2001:db8:abcd:12::0/80")
            Helper.load_as_sets(connection)

            time.sleep(1)

            db_cur.close()
            db_con.close()

            for testfile in sorted(os.listdir("testfiles/")):
                self.clear()

                self.curr_test = testfile
                self.messages = {}
                # load test
                with open("testfiles/{}".format(testfile), "r") as f:
                    self.messages = json.load(f)

                send_len = len(self.messages)

                with nested(
                    connection.Consumer(
                        self.hijack_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson"],
                    ),
                    connection.Consumer(
                        self.update_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson", "txtjson"],
                    ),
                    connection.Consumer(
                        self.hijack_db_queue,
                        callbacks=[self.validate_message],
                        accept=["ujson", "txtjson"],
                    ),
                ):
                    send_cnt = 0
                    # send and validate all messages in the messages.json file
                    while send_cnt < send_len:
                        self.curr_idx = send_cnt
                        self.send_next_message(connection)
                        send_cnt += 1
                        # sleep until we receive all expected messages
                        while self.curr_idx != send_cnt:
                            time.sleep(0.1)
                            try:
                                connection.drain_events(timeout=10)
                            except socket.timeout:
                                # avoid infinite loop by timeout
                                assert False, "Consumer timeout"

        with open("configs/config.yaml") as f1, open("configs/config2.yaml") as f2:
            new_data = f2.read()
            old_data = f1.read()

        Helper.change_conf(connection, new_data, old_data, "test")

        time.sleep(5)
Exemplo n.º 26
0
    def receive_states(self):
        """Wait message at a specific topic and callback when receive a message
        """
        def handle_notification(body, message):
            #time = datetime.now()
            #print ("time: " + str(datetime.now()))
            #time = str(time)

            self.logger.info("Receive state!\n")
            message = json.loads(body)["body"]["states"]
            # self.logger.debug(("message: ", message))

            for i in range(len(message)):
                item_id = json.loads(body)["body"]["states"][i]["MetricId"]
                item_name = json.loads(
                    body)["body"]["states"][i]["MetricLocalId"]
                item_type = json.loads(
                    body)["body"]["states"][i]["DataPoint"]["DataType"]
                item_state = json.loads(
                    body)["body"]["states"][i]["DataPoint"]["Value"]
                self.logger.info("item_name: " + str(item_name))
                self.logger.info("item_state: " + str(item_state))

                # time = json.loads(body)[""]

                # Write new item to the database
                write_result = self.write_item_to_database(
                    item_id, item_name, item_type, item_state,
                    str(datetime.now()))

                #if write new item success to the database, --> consider the trigger with that item
                if (write_result == 1):
                    list_event_trigger = self.read_event_trigger()
                    for event_trigger in list_event_trigger:
                        trigger_id = event_trigger[0]
                        trigger_type = event_trigger[1]
                        trigger_content = event_trigger[2]

                        # Check the item with each trigger condition
                        result = self.check_trigger_condition(
                            trigger_id, trigger_type, trigger_content, item_id,
                            item_state)
                        self.logger.info("check trigger : " + str(result))

                        if (result == None):
                            return None

                        # If checkTriggerCondition success, create an event
                        if (result == True):
                            # event_source = item_id
                            # event_name = str(item_name) + "_" + str(datetime.now())
                            # event_id   = randint(1, 1000000)
                            # my_event = Event(event_name, event_id, event_source)
                            self.create_event(trigger_id)

            self.logger.info("Finish handle one notification! \n\n\n")

            # End handle_notification

        try:
            self.consumer_connection.ensure_connection(max_retries=1)
            with nested(
                    Consumer(self.consumer_connection,
                             queues=self.queue_get_states,
                             callbacks=[handle_notification],
                             no_ack=True)):
                while True:
                    self.consumer_connection.drain_events()
        except (ConnectionRefusedError, exceptions.OperationalError):
            self.logger.error('Connection lost', exc_info=True)
        except self.consumer_connection.connection_errors:
            self.logger.error('Connection error', exc_info=True)