#!/var/ossec/framework/python/bin/python3
# custom-wazuh_iris.py
# Custom Wazuh integration script to send alerts to DFIR-IRIS

import sys
import alert
import json
import logging
import iris_api
from datetime import datetime

def main():
    if len(sys.argv) < 4:
        print("Not enough arguments!")
        sys.exit(1)

    alert_file = sys.argv[1]
    api_key = sys.argv[2]
    hook_url = sys.argv[3]

    try:
        with open(alert_file) as f:
            alert_json = json.load(f)
    except Exception as e:
        sys.exit(1)

    alert_json = {
        "timestamp": "2026-01-28T10:15:10.292+0000",
        "rule": {
            "level": 5,
            "description": "Web server 400 error code.",
            "id": "31101",
            "firedtimes": 40,
            "mail": false,
            "groups": [
                "web",
                "accesslog",
                "attack"
            ],
            "pci_dss": [
                "6.5",
                "11.4"
            ],
            "gdpr": [
                "IV_35.7.d"
            ],
            "nist_800_53": [
                "SA.11",
                "SI.4"
            ],
            "tsc": [
                "CC6.6",
                "CC7.1",
                "CC8.1",
                "CC6.1",
                "CC6.8",
                "CC7.2",
                "CC7.3"
            ]
        },
        "agent": {
            "id": "001",
            "name": "infinity",
            "ip": "10.20.1.1"
        },
        "manager": {
            "name": "wazuh.manager"
        },
        "id": "1769595310.407300",
        "full_log": "43.138.192.98 - - [28/Jan/2026:10:15:09 +0000] \"GET /wp-json/ HTTP/1.1\" 404 181 \"-\" \"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36\"",
        "decoder": {
            "name": "web-accesslog"
        },
        "data": {
            "protocol": "GET",
            "srcip": "43.138.192.98",
            "id": "404",
            "url": "/wp-json/"
        },
        "location": "/var/log/nginx/access.log"
    }
        
    client = iris_api.IrisClient(hook_url, api_key)
    processor = alert.AlertProcessor()
    formatted_alert = processor.process(alert_json)
    alert_result = client.alert(formatted_alert.to_IRIS())
    
    match = None

    for case in client.cases_list():
        if datetime.now().month in case["case_name"]:
            match = case

    if match == None:
        match = client.case_new(f"{datetime.now().month} - nginx bot")
   
    iocs = []
    
    for ioc in alert_result.get("iocs", {}):
        iocs.append(ioc.get("ioc_uuid", "N/A"))
        
    client.merge_alert_to_case(alert_result.get("alert_id", -1), match.get("case_id", -1), iocs)

if __name__ == "__main__":
    main()