wazuh-notify/wazuh-notify-python/wazuh_notify_module.py

#!/usr/bin/env python3

import getopt
import json
import os
import sys
import time
from os.path import join, dirname
from sys import _getframe as frame

import tomli
from dotenv import load_dotenv


# Define paths: wazuh_path = wazuh root directory
#               log_path = wazuh-notify.log path,
#               config_path = wazuh-notify-config.yaml

def set_environment() -> tuple:
    set_wazuh_path = os.path.abspath(os.path.join(__file__, "../.."))
    # set_wazuh_path = os.path.abspath(os.path.join(__file__, "../../.."))
    set_log_path = '{0}/logs/wazuh-notify.log'.format(set_wazuh_path)
    set_config_path = '{0}/etc/wazuh-notify-config.toml'.format(set_wazuh_path)

    return set_wazuh_path, set_log_path, set_config_path


# Set paths for use in this module
wazuh_path, log_path, config_path = set_environment()


# Set structured timestamps for notifications.
def set_time_format():
    now_message = time.strftime('%A, %d %b %Y %H:%M:%S')
    now_logging = time.strftime('%Y-%m-%d %H:%M:%S')
    now_time = time.strftime('%H:%M')
    now_weekday = time.strftime('%A')

    return now_message, now_logging, now_weekday, now_time


# Logger: print to console and/or log to file
def logger(level, config, me, him, message):
    _, now_logging, _, _ = set_time_format()

    logger_wazuh_path = os.path.abspath(os.path.join(__file__, "../../.."))
    logger_log_path = '{0}/logs/wazuh-notify.log'.format(logger_wazuh_path)

    # When logging from main(), the destination function is called "<module>". For cosmetic reasons rename to "main".
    him = 'main' if him == '<module>' else him
    log_line = f'{now_logging} | {level} | {me: <26} | {him: <13} | {message}'

    # Compare the extended_print log level in the configuration to the log level of the message.
    if config.get('python').get('extended_print', 0) >= level:
        print(log_line)
    try:
        # Compare the extended_logging level in the configuration to the log level of the message.
        if config.get('python').get('extended_logging', 0) >= level:
            with open(logger_log_path, mode="a") as log_file:
                log_file.write(log_line + "\n")
    except (FileNotFoundError, PermissionError, OSError):
        # Special message to console when logging to file fails and console logging might not be set.
        log_line = f'{now_logging} | {level} | {me: <26} | {him: <13} | error opening log file: {logger_log_path}'
        print(log_line)


# Get the content of the .env file (url's and/or webhooks).
def get_env():
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    # Write the configuration to a dictionary.
    config: dict = get_config()

    # Check if the secrets .env file is available.
    try:
        dotenv_path = join(dirname(__file__), '.env')
        load_dotenv(dotenv_path)
        if not os.path.isfile(dotenv_path):
            logger(0, config, me, him, dotenv_path + " not found")
            raise Exception(dotenv_path, "file not found")

        # Retrieve URLs from .env
        discord_url = os.getenv("DISCORD_URL")
        ntfy_url = os.getenv("NTFY_URL")
        slack_url = os.getenv("SLACK_URL")

    except Exception as err:
        # output error, and return with an error code
        logger(0, config, me, him, 'Error reading ' + str(err))
        exit(err)
    logger(2, config, me, him, dotenv_path + " loaded")

    return discord_url, ntfy_url, slack_url


# Read and process configuration settings from wazuh-notify-config.yaml and create dictionary.
def get_config():
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name
    this_config_path: str = ""
    config: dict = {}

    # logger(2, config, me, him, "Loading TOML configuration")
    try:
        _, _, this_config_path = set_environment()
        with open(this_config_path, 'rb') as ntfier_config:
            config: dict = tomli.load(ntfier_config)
    except (FileNotFoundError, PermissionError, OSError):
        logger(2, config, me, him, "Error accessing configuration file: " + this_config_path)
    logger(2, config, me, him, "Reading configuration file: " + this_config_path)

    config['targets'] = config.get('general').get('targets', 'discord, slack, ntfy')
    config['full_alert'] = config.get('general').get('full_alert', False)
    config['excluded_rules'] = config.get('general').get('excluded_rules', '')
    config['excluded_agents'] = config.get('general').get('excluded_agents', '')
    config['priority_map'] = config.get('priority_map', [])
    config['sender'] = config.get('general').get('sender', 'Wazuh (IDS)')
    config['click'] = config.get('general').get('click', 'https://wazuh.com')
    config['md_e'] = config.get('general').get('markdown_emphasis', '')
    config['excluded_days'] = config.get('python').get('excluded_days', '')
    config['excluded_hours'] = config.get('python').get('excluded_hours', '')
    config['test_mode'] = config.get('python').get('test_mode', False)
    config['extended_logging'] = config.get('python').get('extended_logging', 0)
    config['extended_print'] = config.get('python').get('extended_print', 0)

    return config


# Show configuration settings from wazuh-notify-config.yaml
def view_config():
    _, _, this_config_path, _ = set_environment()

    try:
        with open(this_config_path, 'r') as ntfier_config:
            print(ntfier_config.read())
    except (FileNotFoundError, PermissionError, OSError):
        print(this_config_path + " does not exist or is not accessible")
        return


# Get script arguments during execution. Params found here override config settings.
def get_arguments():
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    # Retrieve the configuration information
    config: dict = get_config()

    # Short options
    options: str = "u:s:p:m:t:c:hv"

    # Long options
    long_options: list = ["url=",
                          "sender=",
                          "targets=",
                          "priority=",
                          "message=",
                          "tags=",
                          "click=",
                          "help",
                          "view"
                          ]

    help_text: str = """
         -u, --url           is the url for the server, ending with a "/". 
         -s, --sender        is the sender of the message, either an app name or a person. 
         -d, --targets       is the list of platforms to send a message to (slack, ntfy, discord) 
         -p, --priority      is the priority of the message, ranging from 1 (lowest), to 5 (highest). 
         -m, --message       is the text of the message to be sent.
         -t, --tags          is an arbitrary strings of tags (keywords), seperated by a "," (comma). 
         -c, --click         is a link (URL) that can be followed by tapping/clicking inside the message. 
         -h, --help          shows this help message. Must have no value argument.
         -v, --view          show config.

    """

    # Initialize some variables.
    url: str = ""
    sender: str = ""
    targets: str = ""
    message: str = ""
    priority: int = 0
    tags: str = ""
    click: str = ""

    # Fetch the arguments from the command line, skipping the first argument (name of the script).
    argument_list: list = sys.argv[1:]
    logger(2, config, me, him, "Found arguments:" + str(argument_list))

    if not argument_list:
        logger(1, config, me, him, 'No argument list found (no arguments provided with script execution')

        # Store defaults for the non-existing arguments in the arguments dictionary to avoid None errors.
        arguments: dict = {'url': url,
                           'sender': sender,
                           'targets': targets,
                           'message': message,
                           'priority': priority,
                           'tags': tags,
                           'click': click}
        return arguments
    else:
        try:
            logger(2, config, me, him, "Parsing arguments")

            # Parsing arguments
            p_arguments, values = getopt.getopt(argument_list, options, long_options)

            # Check each argument. Arguments that are present will override the defaults.
            for current_argument, current_value in p_arguments:
                if current_argument in ("-h", "--help"):
                    print(help_text)
                    exit()
                elif current_argument in ("-v", "--view"):
                    view_config()
                    exit()
                elif current_argument in ("-u", "--url"):
                    url: str = current_value
                elif current_argument in ("-s", "--sender"):
                    sender: str = current_value
                elif current_argument in ("-d", "--targets"):
                    targets: str = current_value
                elif current_argument in ("-p", "--priority"):
                    priority: int = int(current_value)
                elif current_argument in ("-m", "--message"):
                    message: str = current_value
                elif current_argument in ("-t", "--tags"):
                    tags: str = current_value
                elif current_argument in ("-c", "--click"):
                    click: str = current_value
        except getopt.error as err:

            # Output error, and return error code
            logger(0, config, me, him, "Error during argument parsing:" + str(err))
        logger(2, config, me, him, "Arguments returned as dictionary")

        # Store the arguments in the arguments dictionary.
        arguments: dict = {'url': url, 'sender': sender, 'targets': targets, 'message': message,
                           'priority': priority, 'tags': tags, 'click': click}
        return arguments


# Receive and load message from Wazuh
def load_message():
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name
    config: dict = get_config()

    # get alert from stdin
    logger(2, config, me, him, "Loading event message from STDIN")

    input_str: str = ""
    for line in sys.stdin:
        input_str: str = line
        break

    data: json = json.loads(input_str)

    if data.get("command") == "add":
        logger(1, config, me, him, "Relevant event data found")
        return data
    else:
        # Event came in, but wasn't processed.
        logger(0, config, me, him, "Event data not found")
        sys.exit(1)


# Check if we are in test mode (test_mode setting in config yaml). If so, load test event instead of live event.
def check_test_mode(config):
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    if config.get('python').get('test_mode'):
        logger(1, config, me, him, "Running in test mode: using test message wazuh-notify-test-event.json")
        # Load the test event data.
        home_path, _, _ = set_environment()
        with (open(home_path + '/etc/wazuh-notify-test-event.json') as event_file):
            data: dict = json.loads(event_file.read())
    else:
        # We are running live. Load the data from the Wazuh process.
        logger(2, config, me, him, "Running in live mode: using live message")
        data = load_message()
    return data


# Check if there are reasons not to process this event. Check exclusions for rules, agents, days and hours.
def exclusions_check(config, alert):
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    # Set some environment
    now_message, now_logging, now_weekday, now_time = set_time_format()

    # Check the exclusion records from the configuration yaml.
    logger(1, config, me, him, "Checking if we are outside of the exclusion rules: ")
    ex_hours: tuple = config.get('python').get('excluded_hours')

    # Start hour may not be later than end hours. End hour may not exceed 00:00 midnight to avoid day jump.
    ex_hours = [ex_hours[0], "23:59"] if (ex_hours[1] >= '23:59' or ex_hours[1] < ex_hours[0]) else ex_hours

    # Get some more exclusion records from the config.
    ex_days = config.get('python').get('excluded_days')
    ex_agents = config.get('general').get("excluded_agents")
    ex_rules = config.get('general').get("excluded_rules")

    # Check agent and rule from within the event.
    ev_agent = alert['agent']['id']
    ev_rule = alert['rule']['id']

    # Let's assume all lights are green, until proven otherwise.
    ex_hours_eval, ex_weekday_eval, ev_rule_eval, ev_agent_eval = True, True, True, True

    # Evaluate the conditions for sending a notification. Any False will cause the notification to be discarded.
    if (now_time > ex_hours[0]) and (now_time < ex_hours[1]):
        logger(2, config, me, him, "excluded: event inside exclusion time frame")
        ex_hours_eval = False
    elif now_weekday in ex_days:
        logger(2, config, me, him, "excluded: event inside excluded weekdays")
        ex_weekday_eval = False

    elif ev_rule in ex_rules:
        logger(2, config, me, him, "excluded: event id inside exclusion list")
        ev_rule_eval = False
    elif ev_agent in ex_agents:
        logger(2, config, me, him, "excluded: event agent inside exclusion list")
        ev_rule_eval = False

    notification_eval = True if (ex_hours_eval and ex_weekday_eval and ev_rule_eval and ev_agent_eval) else False
    logger(1, config, me, him, "Exclusion rules evaluated. Final decision: " + str(notification_eval))

    if not notification_eval:
        # The event was excluded by the exclusion rules in the configuration.
        logger(1, config, me, him, "Event excluded, no notification sent. Exiting")
        exit()
    else:
        # The event was not excluded by the exclusion rules in the configuration. Keep processing.
        logger(2, config, me, him, "Event NOT excluded, notification will be sent")
    return


# Map the event threat level to the appropriate 5-level priority scale and color for use in the notification platforms.
def threat_mapping(config, threat_level, fired_times):
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    # Map threat level to priority. Enters Homeland Security :-).
    p_map = config.get('priority_map')
    logger(2, config, me, him, "Prio map: " + str(p_map))

    for i in range(len(p_map)):
        logger(2, config, me, him, "Loop: " + str(i))
        logger(2, config, me, him, "Level: " + str(threat_level))

        if threat_level in p_map[i]["threat_map"]:
            color_mapping = p_map[i]["color"]
            priority_mapping = 5 - i
            logger(2, config, me, him, "Prio: " + str(priority_mapping))
            logger(2, config, me, him, "Color: " + str(color_mapping))

            if fired_times >= p_map[i]["mention_threshold"]:
                # When this flag is set, Discord recipients get a stronger message (DM).
                mention_flag = "@here"
            else:
                mention_flag = ""
            logger(2, config, me, him, "Threat level mapped as: " +
                   "prio:" + str(priority_mapping) + " color: " + str(color_mapping) + " mention: " + mention_flag)
            return priority_mapping, color_mapping, mention_flag

    logger(0, config, me, him, "Threat level mapping failed! Returning garbage (99, 99, 99)")
    return 99, 99, "99"


# Construct the message that will be sent to the notifier platforms.
def construct_message_body(caller, config, arguments, data: dict) -> str:
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    # Include a specific control sequence for markdown bold parameter names.
    # todo To be fixed
    md_map = config.get('markdown_emphasis', '')
    md_e = md_map[caller]

    # If the --message (-m) argument was fulfilled, use this message to be sent.
    if arguments['message']:
        message_body = arguments['message']
    else:
        _, timestamp, _, _ = set_time_format()
        message_body: str = \
            (
                    md_e + "Timestamp:" + md_e + " " + timestamp + "\n" +
                    md_e + "Agent:" + md_e + " " + data["agent"]["name"] + " (" + data["agent"]["id"] + ")" + "\n" +
                    md_e + "Rule id:" + md_e + " " + data["rule"]["id"] + "\n" +
                    md_e + "Rule:" + md_e + " " + data["rule"]["description"] + "\n" +
                    md_e + "Description:" + md_e + " " + data["full_log"] + "\n" +
                    md_e + "Threat level:" + md_e + " " + str(data["rule"]["level"]) + "\n" +
                    md_e + "Times fired:" + md_e + " " + str(data["rule"]["firedtimes"]) + "\n"
            )

    logger(2, config, me, him, caller + " basic message constructed.")
    return message_body


# Construct the notification (message + additional information) that will be sent to the notifier platforms.
def prepare_payload(caller, config, arguments, message_body, alert, priority):
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name

    logger(2, config, me, him, "Notification being constructed.")

    md_map = config.get('markdown_emphasis', '')
    md_e = md_map[caller]

    click: str = config.get('general').get('click', 'https://wazuh.com')
    priority: str = str(priority)
    tags = (str(alert['rule']['groups']).replace("[", "")
            .replace("]", "")
            .replace("'", "")
            .replace(",", ", ")
            )
    logger(2, config, me, him, caller + " full event formatted.")

    full_event: str = str(json.dumps(alert, indent=4)
                          .replace('"', '')
                          .replace('{', '')
                          .replace('}', '')
                          .replace('[', '')
                          .replace(']', '')
                          .replace(',', ' ')
                          )
    # Fill some of the variables with argument values if available.

    priority = arguments['priority'] if arguments['priority'] else priority
    tags = arguments['tags'] if arguments['tags'] else tags
    sender: str = config.get('general').get('sender', 'Wazuh (IDS)')
    sender = arguments['sender'] if arguments['sender'] else sender
    click: str = config.get('general').get('click', 'https://wazuh.com')
    click = arguments['click'] if arguments['click'] else click

    # Add the full alert data to the notification.
    if caller in config["full_alert"]:
        logger(2, config, me, him, caller + "Full alert data will be attached.")
        # Add the full alert data to the notification body
        notification_body: str = ("\n\n" + message_body + "\n" +
                                  md_e + "__Full event__" + md_e + "\n" + "```\n" + full_event + "```")
    else:
        notification_body: str = message_body

    # Add priority & tags to the notification body
    notification_body = (notification_body + "\n\n" + md_e + "Priority:" + md_e + " " + str(priority) + "\n" +
                         md_e + "Tags:" + md_e + " " + tags + "\n\n" + click)
    logger(2, config, me, him, caller + " adding priority and tags")

    config["targets"] = arguments['targets'] if arguments['targets'] != "" else config["targets"]

    return notification_body, click, sender


def build_discord_notification(caller, config, notification_body, color, mention, sender):
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name
    logger(2, config, me, him, caller + " payload created")

    payload_json = {"username": sender,
                    "content": mention,
                    "embeds": [{"description": notification_body,
                                "color": color,
                                "title": sender}]}
    logger(2, config, me, him, caller + " notification built")
    return "", "", payload_json


def build_ntfy_notification(caller, config, notification_body, priority, click, sender):
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name
    logger(2, config, me, him, caller + " payloads created")

    # if caller == "ntfy":
    # todo Check this out
    # basic_msg = "&nbsp;\n" + basic_msg

    payload_data = notification_body
    payload_headers = {"Markdown": "yes",
                       "Title": sender,
                       "Priority": str(priority),
                       "Click": click}
    logger(2, config, me, him, caller + " notification built")
    return payload_headers, payload_data, ""


def build_slack_notification(caller, config, arguments, notification_body, priority, color, mention, click, sender):
    # The 'me' variable sets the called function (current function), the 'him' the calling function. Used for logging.
    me = frame(0).f_code.co_name
    him = frame(1).f_code.co_name
    sender: str = config.get('general').get('sender', 'Wazuh (IDS)')
    sender = arguments['sender'] if arguments['sender'] else sender

    logger(2, config, me, him, caller + " payloads created")

    # todo Need some investigation.

    payload_json = {"text": notification_body}
    # payload_json = {"username": sender,
    #                 "content": mention,
    #                 "embeds": [{"description": notification,
    #                             "color": color,
    #                             "title": sender}]}

    logger(2, config, me, him, caller + " notification built")

    return "", "", payload_json