#!/usr/bin/python3
"""
Custom active response skeleton for Wazuh 5.0.

The message contract, the log file locations, the enable/disable command values and the
exit-code convention come from the Wazuh 5.0-beta documentation:

  https://documentation.wazuh.com/5.0-beta/user-manual/capabilities/active-response/custom-active-response-scripts.html
  https://documentation.wazuh.com/5.0-beta/user-manual/capabilities/active-response/python-active-response-script-sample.html

Everything marked TODO is yours to fill in. Everything marked GUARD is this skeleton's own
addition: the documented sample has no safety rails, and a response script without them is the
fastest way to take an estate offline.

Install (Linux):
    sudo cp custom-ar-skeleton.py /var/ossec/active-response/bin/custom-ar
    sudo chmod 750 /var/ossec/active-response/bin/custom-ar
    sudo chown root:wazuh /var/ossec/active-response/bin/custom-ar

Windows agents execute binaries, so convert this to an .exe with PyInstaller or call it from a
batch launcher placed in C:\\Program Files (x86)\\ossec-agent\\active-response\\bin\\.
"""

import datetime
import ipaddress
import json
import os
import platform
import sys
from pathlib import PurePosixPath, PureWindowsPath

if os.name == "nt":
    LOG_FILE = "C:\\Program Files (x86)\\ossec-agent\\active-response\\active-responses.log"
elif platform.system() == "Darwin":
    LOG_FILE = "/Library/Ossec/logs/active-responses.log"
else:
    LOG_FILE = "/var/ossec/logs/active-responses.log"

ENABLE_COMMAND = 0
DISABLE_COMMAND = 1

OS_SUCCESS = 0
OS_INVALID = -1

# GUARD: set to False only after the dry run has produced the log lines you expect.
DRY_RUN = True

# GUARD: never act on these. Add your domain controllers, jump hosts, load balancers,
# scanner sources, VPN concentrators and egress NAT addresses before you enable anything.
NEVER_TOUCH = {
    # "10.0.0.10",
    # "203.0.113.7",
}
NEVER_TOUCH_NETWORKS = [
    # ipaddress.ip_network("10.10.0.0/24"),
]


def write_debug_file(ar_name, msg):
    """Append a line to the agent's active response log."""
    with open(LOG_FILE, mode="a") as log_file:
        ar_name_posix = str(
            PurePosixPath(PureWindowsPath(ar_name[ar_name.find("active-response"):]))
        )
        log_file.write(
            datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S")
            + " "
            + ar_name_posix
            + ": "
            + msg
            + "\n"
        )


def read_message(argv):
    """Read one JSON message from STDIN and return (data, command_code)."""
    input_str = ""
    for line in sys.stdin:
        input_str = line
        break

    write_debug_file(argv[0], input_str)

    try:
        data = json.loads(input_str)
    except ValueError:
        write_debug_file(argv[0], "Decoding JSON has failed, invalid input format")
        return None, OS_INVALID

    command = data.get("command")
    if command == "enable":
        return data, ENABLE_COMMAND
    if command == "disable":
        return data, DISABLE_COMMAND

    write_debug_file(argv[0], "Not valid command: " + str(command))
    return data, OS_INVALID


def is_protected(argv, address):
    """GUARD: refuse to act on infrastructure you cannot afford to lose."""
    if address in NEVER_TOUCH:
        write_debug_file(argv[0], "Refusing: " + address + " is on the never-touch list")
        return True
    try:
        parsed = ipaddress.ip_address(address)
    except ValueError:
        write_debug_file(argv[0], "Refusing: " + str(address) + " is not a valid IP address")
        return True
    if parsed.is_loopback or parsed.is_link_local:
        write_debug_file(argv[0], "Refusing: " + address + " is loopback or link-local")
        return True
    for network in NEVER_TOUCH_NETWORKS:
        if parsed in network:
            write_debug_file(argv[0], "Refusing: " + address + " is in " + str(network))
            return True
    return False


def main(argv):
    write_debug_file(argv[0], "Started")

    data, command = read_message(argv)
    if command < 0:
        sys.exit(OS_INVALID)

    metadata = data.get("wazuh", {}).get("active_response", {})
    response_type = metadata.get("type", "stateless")
    stateful_timeout = metadata.get("stateful_timeout")

    # The key identifies what this response acted on. A stateful response needs it again on
    # the disable pass. Pick the WCS field your response actually operates on.
    key = data.get("source", {}).get("ip")
    if key is None:
        write_debug_file(argv[0], "No source.ip in the payload, nothing to act on")
        sys.exit(OS_INVALID)

    if is_protected(argv, key):
        # Exit successfully: the refusal is the correct outcome, not an error.
        write_debug_file(argv[0], "Ended without acting")
        sys.exit(OS_SUCCESS)

    if command == ENABLE_COMMAND:
        if DRY_RUN:
            write_debug_file(
                argv[0],
                "DRY RUN: would apply " + response_type + " response to " + key
                + " (timeout=" + str(stateful_timeout) + ")",
            )
        else:
            # TODO: the action. Keep it idempotent: applying it twice must be harmless.
            write_debug_file(argv[0], "Applied response to " + key)

    elif command == DISABLE_COMMAND:
        if DRY_RUN:
            write_debug_file(argv[0], "DRY RUN: would revert response for " + key)
        else:
            # TODO: the reversal. A stateful response is only as good as this branch.
            # It must succeed even if the enable pass partially failed.
            write_debug_file(argv[0], "Reverted response for " + key)

    write_debug_file(argv[0], "Ended")
    sys.exit(OS_SUCCESS)


if __name__ == "__main__":
    main(sys.argv)