#!/usr/bin/python3
# Wazuh 5.x custom active response: scan a changed file with YARA.
#
# Structure follows the official Wazuh 5.0 Python active response sample:
# https://documentation.wazuh.com/5.0-beta/user-manual/capabilities/active-response/python-active-response-script-sample.html
#
# The documented parts are: reading the JSON message from STDIN, the
# "command" field ("enable" / "disable"), the active-responses.log paths,
# and the exit codes. The YARA invocation and the JSON result line are NOT
# documented by Wazuh in the 5.0 beta (there is no YARA module or YARA page)
# and are this script's own addition. Treat them as a starting point.
#
# Type: stateless. A scan has nothing to revert, so only "enable" arrives.
#
# Install (Linux):
#   sudo cp yara-scan.py /var/ossec/active-response/bin/yara-scan
#   sudo chmod 750 /var/ossec/active-response/bin/yara-scan
#   sudo chown root:wazuh /var/ossec/active-response/bin/yara-scan
# macOS: /Library/Ossec/active-response/bin/
# Windows: C:\Program Files (x86)\ossec-agent\active-response\bin\
#          (convert with PyInstaller or call via a launcher.cmd wrapper)

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

# ---------------------------------------------------------------- TODO block
# Set these for your environment before deploying.
YARA_BINARY = "/usr/local/bin/yara"        # TODO: path to the yara executable
YARA_RULES = "/var/ossec/etc/yara/rules.compiled"  # TODO: compiled rule file
SCAN_TIMEOUT_SECONDS = 60                  # TODO: cap the scan, see README
MAX_FILE_BYTES = 64 * 1024 * 1024          # TODO: skip files larger than this
# ---------------------------------------------------------------------------

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"

OS_SUCCESS = 0
OS_INVALID = -1


def write_debug_file(ar_name, msg):
    """Append a line to the agent's active response log (documented path)."""
    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(
            str(datetime.datetime.now().strftime("%Y/%m/%d %H:%M:%S"))
            + " "
            + ar_name_posix
            + ": "
            + msg
            + "\n"
        )


def emit_result(ar_name, payload):
    """Emit one JSON object per scan so a decoder can parse it with parse_json.

    Wazuh ships no YARA decoder in 5.0, so this format is yours to define and
    yours to decode. Keep it flat and stable.
    """
    write_debug_file(ar_name, "yara-scan-result " + json.dumps(payload, sort_keys=True))


def run_yara(file_path):
    """Return (matched_rule_names, error_string)."""
    cmd = [YARA_BINARY, "--no-warnings", YARA_RULES, file_path]
    try:
        proc = subprocess.run(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            timeout=SCAN_TIMEOUT_SECONDS,
            check=False,
        )
    except FileNotFoundError:
        return [], "yara binary not found at " + YARA_BINARY
    except subprocess.TimeoutExpired:
        return [], "yara scan timed out after %ds" % SCAN_TIMEOUT_SECONDS

    if proc.returncode not in (0, 1):
        return [], (proc.stderr.decode("utf-8", "replace").strip() or
                    "yara exited with code %d" % proc.returncode)

    # yara prints "<rule_name> <file_path>" per match, one per line.
    matches = []
    for line in proc.stdout.decode("utf-8", "replace").splitlines():
        line = line.strip()
        if line:
            matches.append(line.split(" ", 1)[0])
    return matches, None


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

    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")
        sys.exit(OS_INVALID)

    command = data.get("command")
    if command != "enable":
        # Stateless responses only ever receive "enable". Guard anyway.
        write_debug_file(
            argv[0],
            "Command is not 'enable' (got: " + str(command) + "), nothing to do",
        )
        sys.exit(OS_SUCCESS)

    file_path = data.get("file", {}).get("path")
    if not file_path:
        write_debug_file(argv[0], "No file.path present in the active response payload")
        sys.exit(OS_INVALID)

    agent_id = data.get("wazuh", {}).get("agent", {}).get("id", "unknown")

    if not os.path.isfile(file_path):
        emit_result(argv[0], {"file_path": file_path, "agent_id": agent_id,
                              "status": "skipped", "reason": "not a regular file"})
        sys.exit(OS_SUCCESS)

    try:
        size = os.path.getsize(file_path)
    except OSError as exc:
        emit_result(argv[0], {"file_path": file_path, "agent_id": agent_id,
                              "status": "error", "reason": str(exc)})
        sys.exit(OS_INVALID)

    if size > MAX_FILE_BYTES:
        emit_result(argv[0], {"file_path": file_path, "agent_id": agent_id,
                              "status": "skipped", "reason": "file larger than limit",
                              "size": size})
        sys.exit(OS_SUCCESS)

    matches, error = run_yara(file_path)
    if error:
        emit_result(argv[0], {"file_path": file_path, "agent_id": agent_id,
                              "status": "error", "reason": error})
        write_debug_file(argv[0], "Ended")
        sys.exit(OS_INVALID)

    emit_result(argv[0], {
        "file_path": file_path,
        "agent_id": agent_id,
        "status": "matched" if matches else "clean",
        "matched_rules": matches,
        "size": size,
    })

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


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