#!/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()