#!/usr/bin/python3
# Copyright (C) 2024, Atomicorp, Inc.
# March 15, 2024
# Copyright (C) 2015-2020, Wazuh Inc.
# March 13, 2018.
#
# This program is free software; you can redistribute it
# and/or modify it under the terms of the GNU General Public
# License (version 2) as published by the FSF - Free Software
# Foundation.

import json
import sys
import time
import os

try:
    import requests
    from requests.auth import HTTPBasicAuth
except Exception as e:
    print("No module 'requests' found. Install: pip install requests")
    sys.exit(1)

# Global vars

version = "1.8.0"
debug_enabled = False
pwd = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
json_alert = {}
now = time.strftime("%a %b %d %H:%M:%S %Z %Y")

# Set paths
log_file = '{0}/logs/integrations.log'.format(pwd)
config_file_path = '{0}/etc/alert-config.json'.format(pwd)

def print_help():
    help_text = """
    Usage: atomicorp-alerter [options] <alert_file>
    
    Options:
    --help          Show this help message and exit
    --version       Show the version of the script
    --debug         Enable debug mode

    Arguments:
    alert_file      Path to the JSON file containing the alert data

    Example:
    ./atomicorp-alerter --debug alert.json
    """
    print(help_text)

def print_version():
    print(f"process_alert.py version {version}")

def load_config():
    global debug_enabled
    try:
        with open(config_file_path) as config_file:
            config = json.load(config_file)
            debug_enabled = config.get('debug', False)
    except Exception as e:
        print(f"Error loading config file: {e}")
        sys.exit(1)

def main(args):
    global debug_enabled
    
    debug("# Starting")

    # Read args
    if '--debug' in args:
        debug_enabled = True
        args.remove('--debug')

    alert_file_location = args[1]

    # Load alert. Parse JSON object.
    with open(alert_file_location) as alert_file:
        json_alert = json.load(alert_file)
    debug("# Processing alert")
    debug(json_alert)
    debug("# Processing alert: MARK")

    # Load filter rules from config file
    with open(config_file_path) as config_file:
        filter_config = json.load(config_file)
    debug("# Loaded filter config")
    debug(filter_config)

    # Check if alert matches any filter rules
    matched_rules = []
    for match in filter_config['filters']['matches']:
        match_name = match['name']
        match_rules = match['rules']
        match_passed = True

        for rule in match_rules:
            rule_type = rule['type']
            debug(f"Processing rule: {rule}")  # Add debug statement for the rule being processed

            # Split the rule_type to get the key parts
            key_parts = rule_type.split('.')

            # Check if the key exists in json_alert
            if check_key_exists(json_alert, key_parts):
                # Perform appropriate comparison based on rule type
                key_value = json_alert
                for key in key_parts:
                    key_value = key_value[key]
                if rule.get('min_value') is not None:
                    if key_value < rule['min_value']:
                        match_passed = False
                        break
                elif 'contains' in rule:
                    if not check_contains(key_value, rule['contains']):
                        match_passed = False
                        break        
                else:
                    if key_value != rule['value']:
                        match_passed = False
                        break
            else:
                # Key doesn't exist, rule doesn't match
                match_passed = False
                break

        if match_passed:
            matched_rules.append(match_name)

    # Log matched filters
    for match_name in matched_rules:
        debug(f"Filter {match_name} matched {alert_file_location}")

        for match in filter_config['filters']['matches']:
            if match['name'] == match_name:
                action_type = match['action']['type']
                action_url = match['action']['url']
                if action_type == 'webhook':
                    debug(f"# Sending message for matched rule: {match_name}")
                    send_msg(generate_msg(json_alert, matched_rules), action_url)
                    break  # Assuming we only execute action for the first matched rule

def check_contains(key_value, target_value):
    """
    Check if the key_value contains any part of the target_value.
    Args:
    - key_value (str): The value of the key to check.
    - target_value (str): The target value to search for.

    Returns:
    - bool: True if key_value contains any part of target_value, False otherwise.
    """
    return target_value in key_value

def check_key_exists(json_data, key_parts):
    """
    Check if the specified key exists in the JSON data.
    Args:
    - json_data (dict): The JSON data to search.
    - key_parts (list): A list containing the parts of the key.

    Returns:
    - bool: True if the key exists, False otherwise.
    """
    if len(key_parts) == 1:
        return key_parts[0] in json_data
    else:
        if key_parts[0] in json_data:
            return check_key_exists(json_data[key_parts[0]], key_parts[1:])
        else:
            return False

def debug(*args):
    if debug_enabled:
        msg = "{0}: {1}\n".format(now, ' '.join(map(str, args)))
        print(msg)
        with open(log_file, "a") as f:
            f.write(msg)

def generate_msg(alert, matched_filter_rules):
    level = alert['rule']['level']

    if level <= 4:
        color = "good"
    elif 5 <= level <= 7:
        color = "warning"
    else:
        color = "danger"

    msg = {
        'color': color,
        'pretext': "Atomic OSSEC Alert",
        'title': alert['rule'].get('description', "N/A"),
        'text': alert.get('full_log'),
        'fields': []
    }
    if 'agent' in alert:
        agent_info = f"({alert['agent'].get('id', 'N/A')}) - {alert['agent'].get('name', 'N/A')}"
    
    if 'labels' in alert['agent']:
        agent_group = alert['agent']['labels'].get('group', 'N/A')
        agent_info += f" Group: {agent_group}"
    msg['fields'].append({
        "title": "Agent",
        "value": agent_info
    })

    if 'agentless' in alert:
        msg['fields'].append({
            "title": "Agentless Host",
            "value": alert['agentless']['host']
        })
    msg['fields'].extend([
        {"title": "Location", "value": alert['location']},
        {"title": "Rule ID", "value": f"{alert['rule']['id']} _(Level {level})_"},
        {"title": "Rule Groups", "value": ", ".join(alert['rule']['groups'])},  # Joining the groups into a string
        {"title": "Matched Alert Filter Rules", "value": ", ".join(matched_filter_rules)}  # Add matched filter rules
    ])

    # Add syscheck details if available
    if 'syscheck' in alert and 'audit' in alert['syscheck']:
        audit_info = alert['syscheck']['audit']
        syscheck_fields = {
            "title": "FIM Extended Attributes",
            "value": "\n".join([
                f"- User Name: {audit_info['user']['name']}",
                f"- Process Name: {audit_info['process']['name']}",
                f"- Process CWD: {audit_info['process']['cwd']}",
                f"- Parent Name: {audit_info['process']['parent_name']}",
                f"- Parent CWD: {audit_info['process']['parent_cwd']}",
                f"- Group Name: {audit_info['group']['name']}",
                f"- Effective User Name: {audit_info['effective_user']['name']}"
            ])
        }
        msg['fields'].append(syscheck_fields)

    msg['ts'] = alert['id']
    attach = {'attachments': [msg]}
    return json.dumps(attach)

def send_msg(msg, url):
    headers = {'content-type': 'application/json', 'Accept-Charset': 'UTF-8'}
    try:
        res = requests.post(url, data=msg, headers=headers)
        res.raise_for_status()  # Raises an HTTPError if the HTTP request returned an unsuccessful status code
        debug(f"Message sent successfully: {res.status_code} {res.text}")
    except requests.exceptions.RequestException as e:
        error_msg = f"ERROR: Error sending message: {e}"
        debug(error_msg)
        with open(log_file, "a") as f:
            f.write(f"{now}: {error_msg}\n")


if __name__ == "__main__":
    try:
        # Load configuration
        load_config()

        # Read arguments
        bad_arguments = False
        if len(sys.argv) >= 2:
            if '--help' in sys.argv or '-h' in sys.argv:
                print_help()
                sys.exit(0)
            if '--version' in sys.argv:
                print_version()
                sys.exit(0)
            if '--debug' in sys.argv:
                debug_enabled = True
                sys.argv.remove('--debug')
            script_name = os.path.basename(sys.argv[0])
            msg = f"{now} {script_name} {' '.join(sys.argv[1:])}"
        else:
            script_name = os.path.basename(sys.argv[0])
            msg = f"{now} {script_name} Wrong arguments"
            bad_arguments = True

        # Logging the call
        with open(log_file, 'a') as f:
            f.write(msg + '\n')

        if bad_arguments:
            debug("# Exiting: Bad arguments.")
            print_help()
            sys.exit(1)

        # Main function
        main(sys.argv)

    except Exception as e:
        debug(str(e))
        raise
