neptune/core/lib/rules.py

from yaml import dump, load
from yaml.parser import ParserError
from yaml.scanner import ScannerError

from core.db.storage import db
from core.models import NotificationRule

try:
    from yaml import CDumper as Dumper
    from yaml import CLoader as Loader
except ImportError:
    from yaml import Loader, Dumper

from core.lib.notify import sendmsg
from core.util import logs

log = logs.get_logger("rules")


def rule_matched(rule, message, matched_fields):
    title = f"Rule {rule.name} matched"

    # Dump the message in YAML for readability
    message = dump(message, Dumper=Dumper, default_flow_style=False)
    matched_fields = ", ".join(matched_fields)

    notify_message = f"{rule.name} matched on {matched_fields}\n{message}"
    notify_message = notify_message.encode("utf-8", "replace")
    sendmsg(rule.user, notify_message, title=title)


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:
                parsed_rule = rule.parse()
                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

                rule_field_length = len(parsed_rule.keys())
                matched_field_number = 0
                matched_fields = []
                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_fields.append(field)
                                    # Break out of the token matching loop
                                    break
                        # Continue to next field
                        continue
                    if field in message and message[field] in value:
                        matched_field_number += 1
                        matched_fields.append(field)
                if matched_field_number == rule_field_length - 2:
                    rule_matched(rule, message, matched_fields)


class NotificationRuleData(object):
    def __init__(self, user, data):
        self.user = user
        self.data = data
        self.parsed = None

        self.parse_data()
        self.validate_permissions()

    def validate_permissions(self):
        """
        Validate permissions for the source and index variables.
        """
        if "index" in self.parsed:
            index = self.parsed["index"]
            if type(index) == list:
                for i in index:
                    db.parse_index(self.user, {"index": i}, raise_error=True)
            else:
                db.parse_index(self.user, {"index": index}, raise_error=True)
        else:
            # Get the default value for the user if not present
            index = db.parse_index(self.user, {}, raise_error=True)
            self.parsed["index"] = index

        if "source" in self.parsed:
            source = self.parsed["source"]
            if type(source) == list:
                for i in source:
                    db.parse_source(self.user, {"source": i}, raise_error=True)
            else:
                db.parse_source(self.user, {"source": source}, raise_error=True)
        else:
            # Get the default value for the user if not present
            source = db.parse_source(self.user, {}, raise_error=True)
            self.parsed["source"] = source

    def parse_data(self):
        """
        Parse the data in the text field to YAML.
        """
        try:
            self.parsed = load(self.data, Loader=Loader)
        except (ScannerError, ParserError) as e:
            raise ValueError(f"Invalid YAML: {e}")

    def __str__(self):
        return dump(self.parsed, Dumper=Dumper)

    def get_data(self):
        return self.parsed