guard-service/src/app.py

#!/usr/bin/env python3
import sys
import json
import logging

import sentry_sdk
import pika
import requests
from sentry_sdk.integrations.logging import LoggingIntegration

import jaeger_client
import opentracing
from opentracing.ext import tags
from opentracing.propagation import Format
from requests_opentracing import SessionTracing

import config
import uuid
from mqtt_helper import MQTT

"""
Main entry point
"""

__author__ = "@tormakris"
__copyright__ = "Copyright 2020, Birbnetes Team"
__module_name__ = "app"
__version__text__ = "1"


def setup_rabbit(mqtt_: MQTT) -> None:
    logging.info("Connecting to RabbitMQ...")
    credentials = pika.PlainCredentials(config.RABBIT_USERNAME, config.RABBIT_PASSWORD)
    while True:
        connection = pika.BlockingConnection(pika.ConnectionParameters(host=config.RABBIT_HOSTNAME,
                                                                       credentials=credentials,
                                                                       heartbeat=30,
                                                                       socket_timeout=45))
        channel = connection.channel()
        channel.exchange_declare(exchange=config.RABBIT_EXCHANGE, exchange_type='fanout')
        queue = channel.queue_declare(durable=True, auto_delete=True, queue=uuid.uuid4().urn.split(':')[2],
                                      exclusive=True).method.queue
        channel.queue_bind(exchange=config.RABBIT_EXCHANGE, queue=queue)
        channel.basic_consume(queue=queue, on_message_callback=on_message_creator(mqtt_), auto_ack=False)
        logging.debug("Starting consumption...")
        try:
            channel.start_consuming()  # this automagically responds to heartbeats
        except pika.exceptions.AMQPConnectionError as e:
            logging.warning(f"AMQP Error happened: {e}; Reconnecting...")


def on_message_creator(mqtt_: MQTT):
    """
    This generator is used, so that the mqtt object can be injected just when the callback is registered
    """

    requests_session = SessionTracing(propagate=True)

    def on_message(
            channel: pika.channel.Channel,
            method: pika.spec.Basic.Deliver,
            properties: pika.spec.BasicProperties,
            body: bytes
    ):
        try:
            msg_json = json.loads(body)
        except (json.JSONDecodeError, UnicodeDecodeError) as e:
            logging.error(f"Malformed message from classifier: {e}")
            channel.basic_ack(delivery_tag=method.delivery_tag)
            return

        span_ctx = opentracing.tracer.extract(Format.TEXT_MAP, msg_json)
        span_tags = {tags.SPAN_KIND: tags.SPAN_KIND_CONSUMER}

        with opentracing.tracer.start_active_span(
                'handleMessage', finish_on_close=True, child_of=span_ctx, tags=span_tags
        ) as scope:

            if ('probability' not in msg_json) or ('class' not in msg_json):
                logging.error("Malformed message from classifier: Missing fields")
                channel.basic_ack(delivery_tag=method.delivery_tag)
                return

            # TODO: strurnus should not be hardcoded here
            if (msg_json['class'] == 'sturnus') and (msg_json['probability'] > config.TRIGGER_LEVEL):
                scope.span.log_kv({'event': 'decisionMade', 'alerting': True})
                try:
                    r = requests_session.get(
                        f"http://{config.INPUT_HOSTNAME}/sample/{msg_json['tag']}",
                        timeout=config.INPUT_TIMEOUT
                    )
                except requests.exceptions.Timeout:
                    logging.error(f"Input-service timed out! (Timeout: {config.INPUT_TIMEOUT} sec)")
                    channel.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
                    return

                if r.status_code != 200:
                    logging.error(f"Input-service status code is not 200: {r.status_code}")
                    channel.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
                    return

                if 'device_id' not in r.json():
                    logging.error("Input-service response invalid")
                    channel.basic_nack(delivery_tag=method.delivery_tag, requeue=True)
                    return

                logging.info(f"Sending alert command to device {r.json()['device_id']}...")
                with opentracing.tracer.start_active_span(
                        'publishAlert',
                        tags={
                            tags.SPAN_KIND: tags.SPAN_KIND_PRODUCER,
                            "device_id": r.json()['device_id']
                        }
                ):
                    mqtt_.publish(
                        subtopic=r.json()['device_id'],
                        message=json.dumps({"command": "doAlert"})
                    )

            else:
                scope.span.log_kv({'event': 'decisionMade', 'alerting': False})
                logging.debug(f"Probability is either bellow trigger level, or not the target class. Nothing to do.")

        # This concludes the job
        channel.basic_ack(delivery_tag=method.delivery_tag)

    return on_message


def main():
    logging.basicConfig(
        stream=sys.stdout,
        format="%(asctime)s - %(name)s [%(levelname)s]: %(message)s",
        level=config.LOG_LEVEL
    )
    if config.SENTRY_DSN:
        sentry_logging = LoggingIntegration(
            level=logging.DEBUG,  # Capture info and above as breadcrumbs
            event_level=logging.ERROR  # Send errors as events
        )
        sentry_sdk.init(
            dsn=config.SENTRY_DSN,
            send_default_pii=True,
            integrations=[sentry_logging],
            traces_sample_rate=0.0,
            release=config.RELEASE_ID,
            environment=config.RELEASEMODE,
            _experiments={"auto_enabling_integrations": True}
        )
    jaeger_client.Config(config={}, service_name='guard-service', validate=True).initialize_tracer()
    logging.info("Guard service starting...")
    mqtt = MQTT()
    mqtt.topic = config.MQTT_TOPIC
    mqtt.connect()
    mqtt.client.loop_start()  # Start MQTT event loop on a different thread
    setup_rabbit(mqtt)


if __name__ == "__main__":
    main()