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")

SECONDS_PER_UNIT = {"s": 1, "m": 60, "h": 3600, "d": 86400, "w": 604800}

MAX_WINDOW = 2592000


class RuleParseError(Exception):
    def __init__(self, message, field):
        super().__init__(message)
        self.field = field


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

    # Dump the message in YAML for readability
    message = dump(message, Dumper=Dumper, default_flow_style=False)
    matched = ", ".join([f"{k}: {v}" for k, v in matched.items()])

    notify_message = f"{rule.name} match: {matched}\n{message}"
    notify_message = notify_message.encode("utf-8", "replace")

    cast = {
        "title": title,
        "priority": str(rule.priority),
    }
    if rule.topic is not None:
        cast["topic"] = rule.topic
    sendmsg(rule.user, notify_message, **cast)


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()
                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 NotificationRuleData(object):
    def __init__(self, user, cleaned_data):
        self.user = user
        self.cleaned_data = cleaned_data
        self.data = self.cleaned_data.get("data")
        self.parsed = None

        self.validate_user_permissions()

        self.parse_data()
        self.validate_permissions()
        self.validate_time_fields()

    def validate_user_permissions(self):
        """
        Ensure the user can use notification rules.
        """
        if not self.user.has_perm("core.use_rules"):
            raise RuleParseError("User does not have permission to use rules", "data")

    def validate_time_fields(self):
        """
        Validate the interval and window fields.
        Prohibit window being specified with an ondemand interval.
        """
        interval = self.cleaned_data.get("interval")
        window = self.cleaned_data.get("window")
        if interval == 0 and window is not None:
            raise RuleParseError(
                "Window cannot be specified with on-demand interval", "window"
            )

        if interval is not None and window is None:
            raise RuleParseError(
                "Window must be specified with non-on-demand interval", "window"
            )

        if window is not None:
            window_number = window[:-1]
            if not window_number.isdigit():
                raise RuleParseError("Window prefix must be a number", "window")
            window_number = int(window_number)
            window_unit = window[-1]
            if window_unit not in SECONDS_PER_UNIT:
                raise RuleParseError(
                    f"Window unit must be one of {', '.join(SECONDS_PER_UNIT.keys())}, not '{window_unit}'",
                    "window",
                )
            window_seconds = window_number * SECONDS_PER_UNIT[window_unit]
            print("Window seconds", window_seconds)
            if window_seconds > MAX_WINDOW:
                raise RuleParseError(
                    f"Window cannot be larger than {MAX_WINDOW} seconds (30 days)",
                    "window",
                )

    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 RuleParseError("data", f"Invalid YAML: {e}")

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

    def get_data(self):
        return self.parsed