neptune/core/management/commands/processing.py

import msgpack
from django.core.management.base import BaseCommand
from redis import StrictRedis

from core.lib.rules import rule_matched
from core.models import NotificationRule
from core.util import logs

log = logs.get_logger("processing")


def process_rules(data):
    all_rules = NotificationRule.objects.filter(enabled=True)

    for index, index_messages in data.items():
        for message in index_messages:
            for rule in all_rules:
                # Quicker helper to get the data without spinning
                # up a NotificationRuleData object
                parsed_rule = rule.parse()
                matched = {}
                if "index" not in parsed_rule:
                    continue
                if "source" not in parsed_rule:
                    continue
                rule_index = parsed_rule["index"]
                rule_source = parsed_rule["source"]
                # if not type(rule_index) == list:
                #     rule_index = [rule_index]
                # if not type(rule_source) == list:
                #     rule_source = [rule_source]
                if index not in rule_index:
                    continue
                if message["src"] not in rule_source:
                    continue

                matched["index"] = index
                matched["source"] = message["src"]

                rule_field_length = len(parsed_rule.keys())
                matched_field_number = 0
                for field, value in parsed_rule.items():
                    # if not type(value) == list:
                    #     value = [value]
                    if field == "src":
                        continue
                    if field == "tokens":
                        for token in value:
                            if "tokens" in message:
                                if token in message["tokens"]:
                                    matched_field_number += 1
                                    matched[field] = token
                                    # Break out of the token matching loop
                                    break
                        # Continue to next field
                        continue

                    # Allow partial matches for msg
                    if field == "msg":
                        for msg in value:
                            if "msg" in message:
                                if msg.lower() in message["msg"].lower():
                                    matched_field_number += 1
                                    matched[field] = msg
                                    # Break out of the msg matching loop
                                    break
                        # Continue to next field
                        continue
                    if field in message and message[field] in value:
                        matched_field_number += 1
                        matched[field] = message[field]
                if matched_field_number == rule_field_length - 2:
                    rule_matched(rule, message, matched)


class Command(BaseCommand):
    def handle(self, *args, **options):
        r = StrictRedis(unix_socket_path="/var/run/socks/redis.sock", db=0)
        p = r.pubsub()
        p.psubscribe("messages")
        for message in p.listen():
            if message:
                if message["channel"] == b"messages":
                    data = message["data"]
                    try:
                        unpacked = msgpack.unpackb(data, raw=False)
                    except TypeError:
                        continue
                    process_rules(unpacked)