#!/usr/bin/env python3
import sys
import json
import requests  # if not available, you can shell out to curl instead

def main():
    # Args from Wazuh integrator:
    #   1: alert file path
    #   2: api_key (if defined, often like "api_key:VALUE")
    #   3: hook_url (from <hook_url>)
    alert_file_path = sys.argv[1]
    api_key_arg = sys.argv[2] if len(sys.argv) > 2 else ""
    hook_url = sys.argv[3] if len(sys.argv) > 3 else ""

    with open(alert_file_path, "r") as f:
        alert = json.load(f)

    # Extract data from Wazuh alert
    rule       = alert.get("rule", {})
    agent      = alert.get("agent", {})
    rule_id    = rule.get("id")
    rule_desc  = rule.get("description", "No description")
    level      = rule.get("level", 0)
    full_log   = alert.get("full_log", "")  # <- default to "" so you don't get "None"
    agent_name = agent.get("name", "unknown-agent")

    # Map Wazuh level -> your notification level
    if level >= 13:
        notif_level = "critical"
    elif level >= 10:
        notif_level = "error"
    elif level >= 7:
        notif_level = "warning"
    else:
        notif_level = "info"

    # Build the payload your webhook expects
    payload = {
        "title": f"Wazuh alert on {agent_name}",
        "level": notif_level,
        "ttl": 15000,
        "message": f"Rule {rule_id} - {rule_desc}\n{full_log}"
    }

    headers = {"Content-Type": "application/json"}

    # If your API key should go in a header, you could do something like:
    # if api_key_arg.startswith("api_key:"):
    #     api_key = api_key_arg.split(":", 1)[1]
    #     headers["Authorization"] = f"Bearer {api_key}"
    api_key_arg = (api_key_arg or "").strip()
    if api_key_arg:
        # Wazuh sometimes passes "api_key:VALUE"
        if api_key_arg.startswith("api_key:"):
            api_key_arg = api_key_arg.split(":", 1)[1].strip()

        headers["Authorization"] = f"Bearer {api_key_arg}"

    try:
        r = requests.post(hook_url, headers=headers, json=payload, timeout=5)
        r.raise_for_status()
    except Exception as e:
        # Name here is just a tag for the logs, you can match it to your integration name if you want
        sys.stderr.write(f"[custom-notifiq] Error sending webhook: {e}\n")
        sys.exit(1)

if __name__ == "__main__":
    main()
