neptune/core/lib/rules.py

130 lines
4.6 KiB
Python

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, priority=str(rule.priority))
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